Cache parsed source files in the PrintingDiagnosticConsumer

We were re-parsing each source file for each top-level diagnostic emitted,
which is... rather inefficient. Cache the parsed source files until
the PrintingDiagnosticConsumer goes away.

Author: DougGregor

include/swift/Frontend/PrintingDiagnosticConsumer.h

@@ -51,6 +51,9 @@ class PrintingDiagnosticConsumer : public DiagnosticConsumer {
   void *queuedDiagnostics = nullptr;
   llvm::DenseMap<unsigned, QueuedBuffer> queuedBuffers;
 
+  /// Source file syntax nodes cached by { source manager, buffer ID }.
+  llvm::DenseMap<std::pair<SourceManager *, unsigned>, void *> sourceFileSyntax;
+
 public:
   PrintingDiagnosticConsumer(llvm::raw_ostream &stream = llvm::errs());
   ~PrintingDiagnosticConsumer();
@@ -90,6 +93,9 @@ class PrintingDiagnosticConsumer : public DiagnosticConsumer {
   }
 
 private:
+  /// Retrieve the SourceFileSyntax for the given buffer.
+  void *getSourceFileSyntax(SourceManager &SM, unsigned bufferID);
+
   void queueBuffer(SourceManager &sourceMgr, unsigned bufferID);
   void printDiagnostic(SourceManager &SM, const DiagnosticInfo &Info);
 };

lib/Frontend/PrintingDiagnosticConsumer.cpp

@@ -304,17 +304,27 @@ static SmallVector<unsigned, 1> getSourceBufferStack(
   }
 }
 
-void PrintingDiagnosticConsumer::queueBuffer(
+void *PrintingDiagnosticConsumer::getSourceFileSyntax(
     SourceManager &sourceMgr, unsigned bufferID) {
-  QueuedBuffer knownSourceFile = queuedBuffers[bufferID];
-  if (knownSourceFile)
-    return;
+  auto known = sourceFileSyntax.find({&sourceMgr, bufferID});
+  if (known != sourceFileSyntax.end())
+    return known->second;
 
   auto bufferContents = sourceMgr.getEntireTextForBuffer(bufferID);
   auto sourceFile = swift_ASTGen_parseSourceFile(
       bufferContents.data(), bufferContents.size(),
       "module", "file.swift", /*ctx*/ nullptr);
 
+  sourceFileSyntax[{&sourceMgr, bufferID}] = sourceFile;
+  return sourceFile;
+}
+
+void PrintingDiagnosticConsumer::queueBuffer(
+    SourceManager &sourceMgr, unsigned bufferID) {
+  QueuedBuffer knownSourceFile = queuedBuffers[bufferID];
+  if (knownSourceFile)
+    return;
+
   // Find the parent and position in parent, if there is one.
   int parentID = -1;
   int positionInParent = 0;
@@ -345,6 +355,7 @@ void PrintingDiagnosticConsumer::queueBuffer(
         sourceMgr.getLocForBufferStart(bufferID)).str();
   }
 
+  auto sourceFile = getSourceFileSyntax(sourceMgr, bufferID);
   swift_ASTGen_addQueuedSourceFile(
       queuedDiagnostics, bufferID, sourceFile,
       (const uint8_t*)displayName.data(), displayName.size(),
@@ -423,9 +434,6 @@ void PrintingDiagnosticConsumer::flush(bool includeTrailingBreak) {
     }
     swift_ASTGen_destroyQueuedDiagnostics(queuedDiagnostics);
     queuedDiagnostics = nullptr;
-    for (const auto &buffer : queuedBuffers) {
-      swift_ASTGen_destroySourceFile(buffer.second);
-    }
     queuedBuffers.clear();
 
     if (includeTrailingBreak)
@@ -566,4 +574,9 @@ SourceManager::GetMessage(SourceLoc Loc, llvm::SourceMgr::DiagKind Kind,
 PrintingDiagnosticConsumer::PrintingDiagnosticConsumer(
     llvm::raw_ostream &stream)
     : Stream(stream) {}
-PrintingDiagnosticConsumer::~PrintingDiagnosticConsumer() = default;
+
+PrintingDiagnosticConsumer::~PrintingDiagnosticConsumer() {
+  for (const auto &sourceFileSyntax : sourceFileSyntax) {
+    swift_ASTGen_destroySourceFile(sourceFileSyntax.second);
+  }
+}