1 files changed, 28 insertions, 19 deletions

diff --git a/clang/lib/ExtractAPI/ExtractAPIConsumer.cpp b/clang/lib/ExtractAPI/ExtractAPIConsumer.cpp
index bcc01db..10b2887 100644
--- a/clang/lib/ExtractAPI/ExtractAPIConsumer.cpp
+++ b/clang/lib/ExtractAPI/ExtractAPIConsumer.cpp
@@ -41,8 +41,8 @@ namespace {
 /// information.
 class ExtractAPIVisitor : public RecursiveASTVisitor<ExtractAPIVisitor> {
 public:
-  ExtractAPIVisitor(ASTContext &Context, Language Lang)
-      : Context(Context), API(Context.getTargetInfo().getTriple(), Lang) {}
+  ExtractAPIVisitor(ASTContext &Context, APISet &API)
+      : Context(Context), API(API) {}
 
   const APISet &getAPI() const { return API; }
 
@@ -489,43 +489,40 @@ private:
   }
 
   ASTContext &Context;
-  APISet API;
+  APISet &API;
 };
 
 class ExtractAPIConsumer : public ASTConsumer {
 public:
-  ExtractAPIConsumer(ASTContext &Context, StringRef ProductName, Language Lang,
-                     std::unique_ptr<raw_pwrite_stream> OS)
-      : Visitor(Context, Lang), ProductName(ProductName), OS(std::move(OS)) {}
+  ExtractAPIConsumer(ASTContext &Context, APISet &API)
+      : Visitor(Context, API) {}
 
   void HandleTranslationUnit(ASTContext &Context) override {
     // Use ExtractAPIVisitor to traverse symbol declarations in the context.
     Visitor.TraverseDecl(Context.getTranslationUnitDecl());
-
-    // Setup a SymbolGraphSerializer to write out collected API information in
-    // the Symbol Graph format.
-    // FIXME: Make the kind of APISerializer configurable.
-    SymbolGraphSerializer SGSerializer(Visitor.getAPI(), ProductName);
-    SGSerializer.serialize(*OS);
   }
 
 private:
   ExtractAPIVisitor Visitor;
-  std::string ProductName;
-  std::unique_ptr<raw_pwrite_stream> OS;
 };
 
 } // namespace
 
 std::unique_ptr<ASTConsumer>
 ExtractAPIAction::CreateASTConsumer(CompilerInstance &CI, StringRef InFile) {
-  std::unique_ptr<raw_pwrite_stream> OS = CreateOutputFile(CI, InFile);
+  OS = CreateOutputFile(CI, InFile);
   if (!OS)
     return nullptr;
-  return std::make_unique<ExtractAPIConsumer>(
-      CI.getASTContext(), CI.getInvocation().getFrontendOpts().ProductName,
-      CI.getFrontendOpts().Inputs.back().getKind().getLanguage(),
-      std::move(OS));
+
+  ProductName = CI.getFrontendOpts().ProductName;
+
+  // Now that we have enough information about the language options and the
+  // target triple, let's create the APISet before anyone uses it.
+  API = std::make_unique<APISet>(
+      CI.getTarget().getTriple(),
+      CI.getFrontendOpts().Inputs.back().getKind().getLanguage());
+
+  return std::make_unique<ExtractAPIConsumer>(CI.getASTContext(), *API);
 }
 
 bool ExtractAPIAction::PrepareToExecuteAction(CompilerInstance &CI) {
@@ -557,6 +554,18 @@ bool ExtractAPIAction::PrepareToExecuteAction(CompilerInstance &CI) {
   return true;
 }
 
+void ExtractAPIAction::EndSourceFileAction() {
+  if (!OS)
+    return;
+
+  // Setup a SymbolGraphSerializer to write out collected API information in
+  // the Symbol Graph format.
+  // FIXME: Make the kind of APISerializer configurable.
+  SymbolGraphSerializer SGSerializer(*API, ProductName);
+  SGSerializer.serialize(*OS);
+  OS->flush();
+}
+
 std::unique_ptr<raw_pwrite_stream>
 ExtractAPIAction::CreateOutputFile(CompilerInstance &CI, StringRef InFile) {
   std::unique_ptr<raw_pwrite_stream> OS =