[Frontend] avoid copy of PCH data when PrecompiledPreamble stores it in memory

Instead of unconditionally copying the PCHBuffer into an ostream which can be
backed either by a string or a file, just make the PCHBuffer itself the
in-memory storage.

Differential Revision: https://reviews.llvm.org/D124180

clang/lib/Frontend/PrecompiledPreamble.cpp

@@ -234,9 +234,10 @@ private:
 
 class PrecompilePreambleAction : public ASTFrontendAction {
 public:
-  PrecompilePreambleAction(std::string *InMemStorage,
+  PrecompilePreambleAction(std::shared_ptr<PCHBuffer> Buffer, bool WritePCHFile,
                            PreambleCallbacks &Callbacks)
-      : InMemStorage(InMemStorage), Callbacks(Callbacks) {}
+      : Buffer(std::move(Buffer)), WritePCHFile(WritePCHFile),
+        Callbacks(Callbacks) {}
 
   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                  StringRef InFile) override;
@@ -244,6 +245,12 @@ public:
   bool hasEmittedPreamblePCH() const { return HasEmittedPreamblePCH; }
 
   void setEmittedPreamblePCH(ASTWriter &Writer) {
+    if (FileOS) {
+      *FileOS << Buffer->Data;
+      // Make sure it hits disk now.
+      FileOS->flush();
+    }
+
     this->HasEmittedPreamblePCH = true;
     Callbacks.AfterPCHEmitted(Writer);
   }
@@ -262,7 +269,9 @@ private:
   friend class PrecompilePreambleConsumer;
 
   bool HasEmittedPreamblePCH = false;
-  std::string *InMemStorage;
+  std::shared_ptr<PCHBuffer> Buffer;
+  bool WritePCHFile; // otherwise the PCH is written into the PCHBuffer only.
+  std::unique_ptr<llvm::raw_pwrite_stream> FileOS; // null if in-memory
   PreambleCallbacks &Callbacks;
 };
 
@@ -272,12 +281,11 @@ public:
                              const Preprocessor &PP,
                              InMemoryModuleCache &ModuleCache,
                              StringRef isysroot,
-                             std::unique_ptr<raw_ostream> Out)
-      : PCHGenerator(PP, ModuleCache, "", isysroot,
-                     std::make_shared<PCHBuffer>(),
+                             std::shared_ptr<PCHBuffer> Buffer)
+      : PCHGenerator(PP, ModuleCache, "", isysroot, std::move(Buffer),
                      ArrayRef<std::shared_ptr<ModuleFileExtension>>(),
                      /*AllowASTWithErrors=*/true),
-        Action(Action), Out(std::move(Out)) {}
+        Action(Action) {}
 
   bool HandleTopLevelDecl(DeclGroupRef DG) override {
     Action.Callbacks.HandleTopLevelDecl(DG);
@@ -288,15 +296,6 @@ public:
     PCHGenerator::HandleTranslationUnit(Ctx);
     if (!hasEmittedPCH())
       return;
-
-    // Write the generated bitstream to "Out".
-    *Out << getPCH();
-    // Make sure it hits disk now.
-    Out->flush();
-    // Free the buffer.
-    llvm::SmallVector<char, 0> Empty;
-    getPCH() = std::move(Empty);
-
     Action.setEmittedPreamblePCH(getWriter());
   }
 
@@ -306,7 +305,6 @@ public:
 
 private:
   PrecompilePreambleAction &Action;
-  std::unique_ptr<raw_ostream> Out;
 };
 
 std::unique_ptr<ASTConsumer>
@@ -316,21 +314,18 @@ PrecompilePreambleAction::CreateASTConsumer(CompilerInstance &CI,
   if (!GeneratePCHAction::ComputeASTConsumerArguments(CI, Sysroot))
     return nullptr;
 
-  std::unique_ptr<llvm::raw_ostream> OS;
-  if (InMemStorage) {
-    OS = std::make_unique<llvm::raw_string_ostream>(*InMemStorage);
-  } else {
-    std::string OutputFile;
-    OS = GeneratePCHAction::CreateOutputFile(CI, InFile, OutputFile);
+  if (WritePCHFile) {
+    std::string OutputFile; // unused
+    FileOS = GeneratePCHAction::CreateOutputFile(CI, InFile, OutputFile);
+    if (!FileOS)
+      return nullptr;
   }
-  if (!OS)
-    return nullptr;
 
   if (!CI.getFrontendOpts().RelocatablePCH)
     Sysroot.clear();
 
   return std::make_unique<PrecompilePreambleConsumer>(
-      *this, CI.getPreprocessor(), CI.getModuleCache(), Sysroot, std::move(OS));
+      *this, CI.getPreprocessor(), CI.getModuleCache(), Sysroot, Buffer);
 }
 
 template <class T> bool moveOnNoError(llvm::ErrorOr<T> Val, T &Output) {
@@ -356,9 +351,9 @@ public:
     S->File = std::move(File);
     return S;
   }
-  static std::unique_ptr<PCHStorage> inMemory() {
+  static std::unique_ptr<PCHStorage> inMemory(std::shared_ptr<PCHBuffer> Buf) {
     std::unique_ptr<PCHStorage> S(new PCHStorage());
-    S->Memory.emplace();
+    S->Memory = std::move(Buf);
     return S;
   }
 
@@ -376,11 +371,7 @@ public:
   }
   llvm::StringRef memoryContents() const {
     assert(getKind() == Kind::InMemory);
-    return *Memory;
-  }
-  std::string &memoryBufferForWrite() {
-    assert(getKind() == Kind::InMemory);
-    return *Memory;
+    return StringRef(Memory->Data.data(), Memory->Data.size());
   }
 
 private:
@@ -388,7 +379,7 @@ private:
   PCHStorage(const PCHStorage &) = delete;
   PCHStorage &operator=(const PCHStorage &) = delete;
 
-  llvm::Optional<std::string> Memory;
+  std::shared_ptr<PCHBuffer> Memory;
   std::unique_ptr<TempPCHFile> File;
 };
 
@@ -411,9 +402,10 @@ llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
   PreprocessorOptions &PreprocessorOpts =
       PreambleInvocation->getPreprocessorOpts();
 
+  std::shared_ptr<PCHBuffer> Buffer = std::make_shared<PCHBuffer>();
   std::unique_ptr<PCHStorage> Storage;
   if (StoreInMemory) {
-    Storage = PCHStorage::inMemory();
+    Storage = PCHStorage::inMemory(Buffer);
   } else {
     // Create a temporary file for the precompiled preamble. In rare
     // circumstances, this can fail.
@@ -495,9 +487,10 @@ llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
                                      PreambleInputBuffer.release());
   }
 
-  std::unique_ptr<PrecompilePreambleAction> Act;
-  Act.reset(new PrecompilePreambleAction(
-      StoreInMemory ? &Storage->memoryBufferForWrite() : nullptr, Callbacks));
+  auto Act = std::make_unique<PrecompilePreambleAction>(
+      std::move(Buffer),
+      /*WritePCHFile=*/Storage->getKind() == PCHStorage::Kind::TempFile,
+      Callbacks);
   if (!Act->BeginSourceFile(*Clang.get(), Clang->getFrontendOpts().Inputs[0]))
     return BuildPreambleError::BeginSourceFileFailed;
 
@@ -527,6 +520,7 @@ llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
 
   if (!Act->hasEmittedPreamblePCH())
     return BuildPreambleError::CouldntEmitPCH;
+  Act.reset(); // Frees the PCH buffer frees, unless Storage keeps it in memory.
 
   // Keep track of all of the files that the source manager knows about,
   // so we can verify whether they have changed or not.