[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/include/clang-c/Index.h

@@ -1035,13 +1035,15 @@ enum CXTranslationUnit_Flags {
    * code-completion operations.
    */
   CXTranslationUnit_CacheCompletionResults = 0x08,
+
   /**
-   * \brief DEPRECATED: Enable precompiled preambles in C++.
+   * \brief Used to indicate that the translation unit will be serialized with
+   * \c clang_saveTranslationUnit.
    *
-   * Note: this is a *temporary* option that is available only while
-   * we are testing C++ precompiled preamble support. It is deprecated.
+   * This option is typically used when parsing a header with the intent of
+   * producing a precompiled header.
    */
-  CXTranslationUnit_CXXPrecompiledPreamble = 0x10,
+  CXTranslationUnit_ForSerialization = 0x10,
 
   /**
    * \brief DEPRECATED: Enabled chained precompiled preambles in C++.

clang/include/clang/Frontend/ASTUnit.h

@@ -56,6 +56,7 @@ class Preprocessor;
 class SourceManager;
 class TargetInfo;
 class ASTFrontendAction;
+class ASTDeserializationListener;
 
 /// \brief Utility class for loading a ASTContext from an AST file.
 ///
@@ -71,6 +72,9 @@ private:
   IntrusiveRefCntPtr<ASTContext>        Ctx;
   ASTReader *Reader;
 
+  struct ASTWriterData;
+  OwningPtr<ASTWriterData> WriterData;
+
   FileSystemOptions FileSystemOpts;
 
   /// \brief The AST consumer that received information about the translation
@@ -468,6 +472,8 @@ public:
 
   const std::string &getOriginalSourceFileName();
 
+  ASTDeserializationListener *getDeserializationListener();
+
   /// \brief Add a temporary file that the ASTUnit depends on.
   ///
   /// This file will be erased when the ASTUnit is destroyed.
@@ -773,6 +779,7 @@ public:
                                       bool AllowPCHWithCompilerErrors = false,
                                       bool SkipFunctionBodies = false,
                                       bool UserFilesAreVolatile = false,
+                                      bool ForSerialization = false,
                                       OwningPtr<ASTUnit> *ErrAST = 0);
   
   /// \brief Reparse the source files using the same command-line options that

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;

clang/lib/Serialization/ASTWriter.cpp

@@ -3423,8 +3423,8 @@ void ASTWriter::WriteASTCore(Sema &SemaRef, MemorizeStatCalls *StatCalls,
   // Write the remaining AST contents.
   RecordData Record;
   Stream.EnterSubblock(AST_BLOCK_ID, 5);
-  WriteMetadata(Context, isysroot, OutputFile);
   WriteLanguageOptions(Context.getLangOpts());
+  WriteMetadata(Context, isysroot, OutputFile);
   if (StatCalls && isysroot.empty())
     WriteStatCache(*StatCalls);
 

clang/tools/c-index-test/c-index-test.c

@@ -2383,6 +2383,9 @@ static CXIdxClientFile index_importedASTFile(CXClientData client_data,
     printf(" | name: \"%s\"", clang_getCString(name));
     printf(" | isImplicit: %d\n", info->isImplicit);
     clang_disposeString(name);
+  } else {
+    // PCH file, the rest are not relevant.
+    printf("\n");
   }
 
   return (CXIdxClientFile)info->file;
@@ -3055,7 +3058,8 @@ int write_pch_file(const char *filename, int argc, const char *argv[]) {
                                   argc - num_unsaved_files,
                                   unsaved_files,
                                   num_unsaved_files,
-                                  CXTranslationUnit_Incomplete);
+                                  CXTranslationUnit_Incomplete |
+                                    CXTranslationUnit_ForSerialization);
   if (!TU) {
     fprintf(stderr, "Unable to load translation unit!\n");
     free_remapped_files(unsaved_files, num_unsaved_files);

clang/tools/libclang/CIndex.cpp

@@ -2556,6 +2556,7 @@ static void clang_parseTranslationUnit_Impl(void *UserData) {
   bool IncludeBriefCommentsInCodeCompletion
     = options & CXTranslationUnit_IncludeBriefCommentsInCodeCompletion;
   bool SkipFunctionBodies = options & CXTranslationUnit_SkipFunctionBodies;
+  bool ForSerialization = options & CXTranslationUnit_ForSerialization;
 
   // Configure the diagnostics.
   DiagnosticOptions DiagOpts;
@@ -2643,6 +2644,7 @@ static void clang_parseTranslationUnit_Impl(void *UserData) {
                                  /*AllowPCHWithCompilerErrors=*/true,
                                  SkipFunctionBodies,
                                  /*UserFilesAreVolatile=*/true,
+                                 ForSerialization,
                                  &ErrUnit));
 
   if (NumErrors != Diags->getClient()->getNumErrors()) {