2 files changed, 26 insertions, 29 deletions

diff --git a/lib/Frontend/TextDiagnostic.cpp b/lib/Frontend/TextDiagnostic.cpp
index d1936e3a02..259e05b473 100644
--- a/lib/Frontend/TextDiagnostic.cpp
+++ b/lib/Frontend/TextDiagnostic.cpp
@@ -392,17 +392,9 @@ static bool printWordWrapped(raw_ostream &OS, StringRef Str,
 TextDiagnostic::TextDiagnostic(raw_ostream &OS,
                                const SourceManager &SM,
                                const LangOptions &LangOpts,
-                               const DiagnosticOptions &DiagOpts,
-                               FullSourceLoc LastLoc,
-                               FullSourceLoc LastIncludeLoc,
-                               DiagnosticsEngine::Level LastLevel)
-  : OS(OS), SM(SM), LangOpts(LangOpts), DiagOpts(DiagOpts),
-    LastLoc(LastLoc), LastIncludeLoc(LastIncludeLoc), LastLevel(LastLevel) {
-  if (LastLoc.isValid() && &SM != &LastLoc.getManager())
-    this->LastLoc = SourceLocation();
-  if (LastIncludeLoc.isValid() && &SM != &LastIncludeLoc.getManager())
-    this->LastIncludeLoc = SourceLocation();
-    }
+                               const DiagnosticOptions &DiagOpts)
+  : OS(OS), SM(SM), LangOpts(LangOpts), DiagOpts(DiagOpts), LastLevel() {
+}
 
 void TextDiagnostic::emitDiagnostic(SourceLocation Loc,
                                     DiagnosticsEngine::Level Level,
diff --git a/lib/Frontend/TextDiagnosticPrinter.cpp b/lib/Frontend/TextDiagnosticPrinter.cpp
index cf2826484f..b5d0bcb0cb 100644
--- a/lib/Frontend/TextDiagnosticPrinter.cpp
+++ b/lib/Frontend/TextDiagnosticPrinter.cpp
@@ -27,7 +27,7 @@ using namespace clang;
 TextDiagnosticPrinter::TextDiagnosticPrinter(raw_ostream &os,
                                              const DiagnosticOptions &diags,
                                              bool _OwnsOutputStream)
-  : OS(os), LangOpts(0), DiagOpts(&diags), LastLevel(),
+  : OS(os), LangOpts(0), DiagOpts(&diags), SM(0),
     OwnsOutputStream(_OwnsOutputStream) {
 }
 
@@ -36,6 +36,16 @@ TextDiagnosticPrinter::~TextDiagnosticPrinter() {
     delete &OS;
 }
 
+void TextDiagnosticPrinter::BeginSourceFile(const LangOptions &LO,
+                                            const Preprocessor *PP) {
+  LangOpts = &LO;
+}
+
+void TextDiagnosticPrinter::EndSourceFile() {
+  LangOpts = 0;
+  TextDiag.reset(0);
+}
+
 /// \brief Print the diagnostic name to a raw_ostream.
 ///
 /// This prints the diagnostic name to a raw_ostream if it has one. It formats
@@ -158,23 +168,18 @@ void TextDiagnosticPrinter::HandleDiagnostic(DiagnosticsEngine::Level Level,
   assert(DiagOpts && "Unexpected diagnostic without options set");
   assert(Info.hasSourceManager() &&
          "Unexpected diagnostic with no source manager");
-  const SourceManager &SM = Info.getSourceManager();
-  TextDiagnostic TextDiag(OS, SM, *LangOpts, *DiagOpts,
-                          LastLoc, LastIncludeLoc, LastLevel);
-
-  TextDiag.emitDiagnostic(Info.getLocation(), Level, DiagMessageStream.str(),
-                          Info.getRanges(),
-                          llvm::makeArrayRef(Info.getFixItHints(),
-                                             Info.getNumFixItHints()));
-
-  // Cache the LastLoc from the TextDiagnostic printing.
-  // FIXME: Rather than this, we should persist a TextDiagnostic object across
-  // diagnostics until the SourceManager changes. That will allow the
-  // TextDiagnostic object to form a 'session' of output where we can
-  // reasonably collapse redundant information.
-  LastLoc = FullSourceLoc(TextDiag.getLastLoc(), SM);
-  LastIncludeLoc = FullSourceLoc(TextDiag.getLastIncludeLoc(), SM);
-  LastLevel = TextDiag.getLastLevel();
+
+  // Rebuild the TextDiagnostic utility if missing or the source manager has
+  // changed.
+  if (!TextDiag || SM != &Info.getSourceManager()) {
+    SM = &Info.getSourceManager();
+    TextDiag.reset(new TextDiagnostic(OS, *SM, *LangOpts, *DiagOpts));
+  }
+
+  TextDiag->emitDiagnostic(Info.getLocation(), Level, DiagMessageStream.str(),
+                           Info.getRanges(),
+                           llvm::makeArrayRef(Info.getFixItHints(),
+                                              Info.getNumFixItHints()));
 
   OS.flush();
 }