[libclang] Improve AST serialization done by ASTUnit::Save().

The ASTUnit needs to initialize an ASTWriter at the beginning of
parsing to fully handle serialization of a translation unit that
imports modules. Do this by introducing an option to enable it, which
corresponds to CXTranslationUnit_ForSerialization on the C API side.

llvm-svn: 165717

clang/lib/Frontend/ASTUnit.cpp

@@ -180,6 +180,14 @@ void OnDiskData::Cleanup() {
   CleanPreambleFile();
 }
 
+struct ASTUnit::ASTWriterData {
+  SmallString<128> Buffer;
+  llvm::BitstreamWriter Stream;
+  ASTWriter Writer;
+
+  ASTWriterData() : Stream(Buffer), Writer(Stream) { }
+};
+
 void ASTUnit::clearFileLevelDecls() {
   for (FileDeclsTy::iterator
          I = FileDecls.begin(), E = FileDecls.end(); I != E; ++I)
@@ -514,29 +522,26 @@ public:
 
   virtual bool ReadLanguageOptions(const serialization::ModuleFile &M,
                                    const LangOptions &LangOpts) {
-    if (InitializedLanguage || M.Kind != serialization::MK_MainFile)
+    if (InitializedLanguage)
       return false;
     
-    LangOpt = LangOpts;
-    
-    // Initialize the preprocessor.
-    PP.Initialize(*Target);
-    
-    // Initialize the ASTContext
-    Context.InitBuiltinTypes(*Target);
-    
-    InitializedLanguage = true;
+    assert(M.Kind == serialization::MK_MainFile);
 
-    applyLangOptsToTarget();
+    LangOpt = LangOpts;
+    InitializedLanguage = true;
+    
+    updated();
     return false;
   }
 
   virtual bool ReadTargetTriple(const serialization::ModuleFile &M,
                                 StringRef Triple) {
     // If we've already initialized the target, don't do it again.
-    if (Target || M.Kind != serialization::MK_MainFile)
+    if (Target)
       return false;
     
+    assert(M.Kind == serialization::MK_MainFile);
+
     // FIXME: This is broken, we should store the TargetOptions in the AST file.
     TargetOptions TargetOpts;
     TargetOpts.ABI = "";
@@ -546,7 +551,7 @@ public:
     TargetOpts.Triple = Triple;
     Target = TargetInfo::CreateTargetInfo(PP.getDiagnostics(), TargetOpts);
 
-    applyLangOptsToTarget();
+    updated();
     return false;
   }
 
@@ -570,14 +575,21 @@ public:
   }
 
 private:
-  void applyLangOptsToTarget() {
-    if (Target && InitializedLanguage) {
-      // Inform the target of the language options.
-      //
-      // FIXME: We shouldn't need to do this, the target should be immutable once
-      // created. This complexity should be lifted elsewhere.
-      Target->setForcedLangOptions(LangOpt);
-    }
+  void updated() {
+    if (!Target || !InitializedLanguage)
+      return;
+
+    // Inform the target of the language options.
+    //
+    // FIXME: We shouldn't need to do this, the target should be immutable once
+    // created. This complexity should be lifted elsewhere.
+    Target->setForcedLangOptions(LangOpt);
+
+    // Initialize the preprocessor.
+    PP.Initialize(*Target);
+
+    // Initialize the ASTContext
+    Context.InitBuiltinTypes(*Target);
   }
 };
 
@@ -642,6 +654,12 @@ const std::string &ASTUnit::getOriginalSourceFileName() {
   return OriginalSourceFile;
 }
 
+ASTDeserializationListener *ASTUnit::getDeserializationListener() {
+  if (WriterData)
+    return &WriterData->Writer;
+  return 0;
+}
+
 llvm::MemoryBuffer *ASTUnit::getBufferForFile(StringRef Filename,
                                               std::string *ErrorStr) {
   assert(FileMgr);
@@ -917,6 +935,10 @@ public:
     for (DeclGroupRef::iterator it = D.begin(), ie = D.end(); it != ie; ++it)
       handleTopLevelDecl(*it);
   }
+
+  virtual ASTDeserializationListener *GetASTDeserializationListener() {
+    return Unit.getDeserializationListener();
+  }
 };
 
 class TopLevelDeclTrackerAction : public ASTFrontendAction {
@@ -1929,6 +1951,7 @@ ASTUnit *ASTUnit::LoadFromCommandLine(const char **ArgBegin,
                                       bool AllowPCHWithCompilerErrors,
                                       bool SkipFunctionBodies,
                                       bool UserFilesAreVolatile,
+                                      bool ForSerialization,
                                       OwningPtr<ASTUnit> *ErrAST) {
   if (!Diags.getPtr()) {
     // No diagnostics engine was provided, so create our own diagnostics object
@@ -1992,6 +2015,8 @@ ASTUnit *ASTUnit::LoadFromCommandLine(const char **ArgBegin,
   AST->NumStoredDiagnosticsFromDriver = StoredDiagnostics.size();
   AST->StoredDiagnostics.swap(StoredDiagnostics);
   AST->Invocation = CI;
+  if (ForSerialization)
+    AST->WriterData.reset(new ASTWriterData());
   CI = 0; // Zero out now to ease cleanup during crash recovery.
   
   // Recover resources if we crash before exiting this method.
@@ -2506,20 +2531,31 @@ bool ASTUnit::Save(StringRef File) {
   return false;
 }
 
+static bool serializeUnit(ASTWriter &Writer,
+                          SmallVectorImpl<char> &Buffer,
+                          Sema &S,
+                          bool hasErrors,
+                          raw_ostream &OS) {
+  Writer.WriteAST(S, 0, std::string(), 0, "", hasErrors);
+
+  // Write the generated bitstream to "Out".
+  if (!Buffer.empty())
+    OS.write(Buffer.data(), Buffer.size());
+
+  return false;
+}
+
 bool ASTUnit::serialize(raw_ostream &OS) {
   bool hasErrors = getDiagnostics().hasErrorOccurred();
 
+  if (WriterData)
+    return serializeUnit(WriterData->Writer, WriterData->Buffer,
+                         getSema(), hasErrors, OS);
+
   SmallString<128> Buffer;
   llvm::BitstreamWriter Stream(Buffer);
   ASTWriter Writer(Stream);
-  // FIXME: Handle modules
-  Writer.WriteAST(getSema(), 0, std::string(), 0, "", hasErrors);
-  
-  // Write the generated bitstream to "Out".
-  if (!Buffer.empty())
-    OS.write((char *)&Buffer.front(), Buffer.size());
-
-  return false;
+  return serializeUnit(Writer, Buffer, getSema(), hasErrors, OS);
 }
 
 typedef ContinuousRangeMap<unsigned, int, 2> SLocRemap;