2 files changed, 159 insertions, 86 deletions

diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index 51147b6f949..29cdcd7df1d 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -1087,6 +1087,9 @@ bool clang::ParseDiagnosticArgs(DiagnosticOptions &Opts, ArgList &Args,
   Opts.SpellCheckingLimit = getLastArgIntValue(
       Args, OPT_fspell_checking_limit,
       DiagnosticOptions::DefaultSpellCheckingLimit, Diags);
+  Opts.SnippetLineLimit = getLastArgIntValue(
+      Args, OPT_fcaret_diagnostics_max_lines,
+      DiagnosticOptions::DefaultSnippetLineLimit, Diags);
   Opts.TabStop = getLastArgIntValue(Args, OPT_ftabstop,
                                     DiagnosticOptions::DefaultTabStop, Diags);
   if (Opts.TabStop == 0 || Opts.TabStop > DiagnosticOptions::MaxTabStop) {
diff --git a/clang/lib/Frontend/TextDiagnostic.cpp b/clang/lib/Frontend/TextDiagnostic.cpp
index a4937386b93..a18b8c888c2 100644
--- a/clang/lib/Frontend/TextDiagnostic.cpp
+++ b/clang/lib/Frontend/TextDiagnostic.cpp
@@ -928,6 +928,56 @@ void TextDiagnostic::emitBuildingModuleLocation(SourceLocation Loc,
     OS << "While building module '" << ModuleName << "':\n";
 }
 
+/// \brief Find the suitable set of lines to show to include a set of ranges.
+static llvm::Optional<std::pair<unsigned, unsigned>>
+findLinesForRange(const CharSourceRange &R, FileID FID,
+                  const SourceManager &SM) {
+  if (!R.isValid()) return None;
+
+  SourceLocation Begin = R.getBegin();
+  SourceLocation End = R.getEnd();
+  if (SM.getFileID(Begin) != FID || SM.getFileID(End) != FID)
+    return None;
+
+  return std::make_pair(SM.getExpansionLineNumber(Begin),
+                        SM.getExpansionLineNumber(End));
+}
+
+/// Add as much of range B into range A as possible without exceeding a maximum
+/// size of MaxRange. Ranges are inclusive.
+std::pair<unsigned, unsigned> maybeAddRange(std::pair<unsigned, unsigned> A,
+                                            std::pair<unsigned, unsigned> B,
+                                            unsigned MaxRange) {
+  // If A is already the maximum size, we're done.
+  unsigned Slack = MaxRange - (A.second - A.first + 1);
+  if (Slack == 0)
+    return A;
+
+  // Easy case: merge succeeds within MaxRange.
+  unsigned Min = std::min(A.first, B.first);
+  unsigned Max = std::max(A.second, B.second);
+  if (Max - Min + 1 <= MaxRange)
+    return {Min, Max};
+
+  // If we can't reach B from A within MaxRange, there's nothing to do.
+  // Don't add lines to the range that contain nothing interesting.
+  if ((B.first > A.first && B.first - A.first + 1 > MaxRange) ||
+      (B.second < A.second && A.second - B.second + 1 > MaxRange))
+    return A;
+
+  // Otherwise, expand A towards B to produce a range of size MaxRange. We
+  // attempt to expand by the same amount in both directions if B strictly
+  // contains A.
+
+  // Expand downwards by up to half the available amount, then upwards as
+  // much as possible, then downwards as much as possible.
+  A.second = std::min(A.second + (Slack + 1) / 2, Max);
+  Slack = MaxRange - (A.second - A.first + 1);
+  A.first = std::max(Min + Slack, A.first) - Slack;
+  A.second = std::min(A.first + MaxRange - 1, Max);
+  return A;
+}
+
 /// \brief Highlight a SourceRange (with ~'s) for any characters on LineNo.
 static void highlightRange(const CharSourceRange &R,
                            unsigned LineNo, FileID FID,
@@ -990,9 +1040,12 @@ static void highlightRange(const CharSourceRange &R,
       EndColNo = map.startOfPreviousColumn(EndColNo);
 
     // If the start/end passed each other, then we are trying to highlight a
-    // range that just exists in whitespace, which must be some sort of other
-    // bug.
-    assert(StartColNo <= EndColNo && "Trying to highlight whitespace??");
+    // range that just exists in whitespace. That most likely means we have
+    // a multi-line highlighting range that covers a blank line.
+    if (StartColNo > EndColNo) {
+      assert(StartLineNo != EndLineNo && "trying to highlight whitespace");
+      StartColNo = EndColNo;
+    }
   }
 
   assert(StartColNo <= map.getSourceLine().size() && "Invalid range!");
@@ -1103,7 +1156,7 @@ void TextDiagnostic::emitSnippetAndCaret(
   // Decompose the location into a FID/Offset pair.
   std::pair<FileID, unsigned> LocInfo = SM.getDecomposedLoc(Loc);
   FileID FID = LocInfo.first;
-  unsigned FileOffset = LocInfo.second;
+  unsigned CaretFileOffset = LocInfo.second;
 
   // Get information about the buffer it points into.
   bool Invalid = false;
@@ -1111,101 +1164,118 @@ void TextDiagnostic::emitSnippetAndCaret(
   if (Invalid)
     return;
 
-  const char *BufStart = BufData.data();
-  const char *BufEnd = BufStart + BufData.size();
+  unsigned CaretLineNo = SM.getLineNumber(FID, CaretFileOffset);
+  unsigned CaretColNo = SM.getColumnNumber(FID, CaretFileOffset);
 
-  unsigned LineNo = SM.getLineNumber(FID, FileOffset);
-  unsigned ColNo = SM.getColumnNumber(FID, FileOffset);
-  
   // Arbitrarily stop showing snippets when the line is too long.
   static const size_t MaxLineLengthToPrint = 4096;
-  if (ColNo > MaxLineLengthToPrint)
+  if (CaretColNo > MaxLineLengthToPrint)
     return;
 
-  // Rewind from the current position to the start of the line.
-  const char *TokPtr = BufStart+FileOffset;
-  const char *LineStart = TokPtr-ColNo+1; // Column # is 1-based.
-
-  // Compute the line end.  Scan forward from the error position to the end of
-  // the line.
-  const char *LineEnd = TokPtr;
-  while (*LineEnd != '\n' && *LineEnd != '\r' && LineEnd != BufEnd)
-    ++LineEnd;
-
-  // Arbitrarily stop showing snippets when the line is too long.
-  if (size_t(LineEnd - LineStart) > MaxLineLengthToPrint)
-    return;
-
-  // Trim trailing null-bytes.
-  StringRef Line(LineStart, LineEnd - LineStart);
-  while (Line.size() > ColNo && Line.back() == '\0')
-    Line = Line.drop_back();
+  // Find the set of lines to include.
+  const unsigned MaxLines = DiagOpts->SnippetLineLimit;
+  std::pair<unsigned, unsigned> Lines = {CaretLineNo, CaretLineNo};
+  for (SmallVectorImpl<CharSourceRange>::iterator I = Ranges.begin(),
+                                                  E = Ranges.end();
+       I != E; ++I)
+    if (auto OptionalRange = findLinesForRange(*I, FID, SM))
+      Lines = maybeAddRange(Lines, *OptionalRange, MaxLines);
+
+  for (unsigned LineNo = Lines.first; LineNo != Lines.second + 1; ++LineNo) {
+    const char *BufStart = BufData.data();
+    const char *BufEnd = BufStart + BufData.size();
+
+    // Rewind from the current position to the start of the line.
+    const char *LineStart =
+        BufStart +
+        SM.getDecomposedLoc(SM.translateLineCol(FID, LineNo, 1)).second;
+    if (LineStart == BufEnd)
+      break;
 
-  // Copy the line of code into an std::string for ease of manipulation.
-  std::string SourceLine(Line.begin(), Line.end());
+    // Compute the line end.
+    const char *LineEnd = LineStart;
+    while (*LineEnd != '\n' && *LineEnd != '\r' && LineEnd != BufEnd)
+      ++LineEnd;
 
-  // Build the byte to column map.
-  const SourceColumnMap sourceColMap(SourceLine, DiagOpts->TabStop);
+    // Arbitrarily stop showing snippets when the line is too long.
+    // FIXME: Don't print any lines in this case.
+    if (size_t(LineEnd - LineStart) > MaxLineLengthToPrint)
+      return;
 
-  // Create a line for the caret that is filled with spaces that is the same
-  // number of columns as the line of source code.
-  std::string CaretLine(sourceColMap.columns(), ' ');
+    // Trim trailing null-bytes.
+    StringRef Line(LineStart, LineEnd - LineStart);
+    while (!Line.empty() && Line.back() == '\0' &&
+           (LineNo != CaretLineNo || Line.size() > CaretColNo))
+      Line = Line.drop_back();
+
+    // Copy the line of code into an std::string for ease of manipulation.
+    std::string SourceLine(Line.begin(), Line.end());
+
+    // Build the byte to column map.
+    const SourceColumnMap sourceColMap(SourceLine, DiagOpts->TabStop);
+
+    // Create a line for the caret that is filled with spaces that is the same
+    // number of columns as the line of source code.
+    std::string CaretLine(sourceColMap.columns(), ' ');
+
+    // Highlight all of the characters covered by Ranges with ~ characters.
+    for (SmallVectorImpl<CharSourceRange>::iterator I = Ranges.begin(),
+                                                    E = Ranges.end();
+         I != E; ++I)
+      highlightRange(*I, LineNo, FID, sourceColMap, CaretLine, SM, LangOpts);
+
+    // Next, insert the caret itself.
+    if (CaretLineNo == LineNo) {
+      CaretColNo = sourceColMap.byteToContainingColumn(CaretColNo - 1);
+      if (CaretLine.size() < CaretColNo + 1)
+        CaretLine.resize(CaretColNo + 1, ' ');
+      CaretLine[CaretColNo] = '^';
+    }
 
-  // Highlight all of the characters covered by Ranges with ~ characters.
-  for (SmallVectorImpl<CharSourceRange>::iterator I = Ranges.begin(),
-                                                  E = Ranges.end();
-       I != E; ++I)
-    highlightRange(*I, LineNo, FID, sourceColMap, CaretLine, SM, LangOpts);
-
-  // Next, insert the caret itself.
-  ColNo = sourceColMap.byteToContainingColumn(ColNo-1);
-  if (CaretLine.size()<ColNo+1)
-    CaretLine.resize(ColNo+1, ' ');
-  CaretLine[ColNo] = '^';
-
-  std::string FixItInsertionLine = buildFixItInsertionLine(LineNo,
-                                                           sourceColMap,
-                                                           Hints, SM,
-                                                           DiagOpts.get());
-
-  // If the source line is too long for our terminal, select only the
-  // "interesting" source region within that line.
-  unsigned Columns = DiagOpts->MessageLength;
-  if (Columns)
-    selectInterestingSourceRegion(SourceLine, CaretLine, FixItInsertionLine,
-                                  Columns, sourceColMap);
-
-  // If we are in -fdiagnostics-print-source-range-info mode, we are trying
-  // to produce easily machine parsable output.  Add a space before the
-  // source line and the caret to make it trivial to tell the main diagnostic
-  // line from what the user is intended to see.
-  if (DiagOpts->ShowSourceRanges) {
-    SourceLine = ' ' + SourceLine;
-    CaretLine = ' ' + CaretLine;
-  }
+    std::string FixItInsertionLine = buildFixItInsertionLine(
+        LineNo, sourceColMap, Hints, SM, DiagOpts.get());
+
+    // If the source line is too long for our terminal, select only the
+    // "interesting" source region within that line.
+    unsigned Columns = DiagOpts->MessageLength;
+    if (Columns)
+      selectInterestingSourceRegion(SourceLine, CaretLine, FixItInsertionLine,
+                                    Columns, sourceColMap);
+
+    // If we are in -fdiagnostics-print-source-range-info mode, we are trying
+    // to produce easily machine parsable output.  Add a space before the
+    // source line and the caret to make it trivial to tell the main diagnostic
+    // line from what the user is intended to see.
+    if (DiagOpts->ShowSourceRanges) {
+      SourceLine = ' ' + SourceLine;
+      CaretLine = ' ' + CaretLine;
+    }
 
-  // Finally, remove any blank spaces from the end of CaretLine.
-  while (CaretLine[CaretLine.size()-1] == ' ')
-    CaretLine.erase(CaretLine.end()-1);
+    // Finally, remove any blank spaces from the end of CaretLine.
+    while (CaretLine[CaretLine.size() - 1] == ' ')
+      CaretLine.erase(CaretLine.end() - 1);
 
-  // Emit what we have computed.
-  emitSnippet(SourceLine);
+    // Emit what we have computed.
+    emitSnippet(SourceLine);
 
-  if (DiagOpts->ShowColors)
-    OS.changeColor(caretColor, true);
-  OS << CaretLine << '\n';
-  if (DiagOpts->ShowColors)
-    OS.resetColor();
+    if (!CaretLine.empty()) {
+      if (DiagOpts->ShowColors)
+        OS.changeColor(caretColor, true);
+      OS << CaretLine << '\n';
+      if (DiagOpts->ShowColors)
+        OS.resetColor();
+    }
 
-  if (!FixItInsertionLine.empty()) {
-    if (DiagOpts->ShowColors)
-      // Print fixit line in color
-      OS.changeColor(fixitColor, false);
-    if (DiagOpts->ShowSourceRanges)
-      OS << ' ';
-    OS << FixItInsertionLine << '\n';
-    if (DiagOpts->ShowColors)
-      OS.resetColor();
+    if (!FixItInsertionLine.empty()) {
+      if (DiagOpts->ShowColors)
+        // Print fixit line in color
+        OS.changeColor(fixitColor, false);
+      if (DiagOpts->ShowSourceRanges)
+        OS << ' ';
+      OS << FixItInsertionLine << '\n';
+      if (DiagOpts->ShowColors)
+        OS.resetColor();
+    }
   }
 
   // Print out any parseable fixit information requested by the options.