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

[sammccall]

Summary

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.

[thakis]

I think this might've broken check-clangd: http://45.33.8.238/linux/74320/step_9.txt

The output isn't very useful though, I'm guessing @yln's recent gtest runner change it's not getting printed.

[thakis]

Sorry, @ychen, not @yln (D122251)

[sammccall]

In D124180#3465304, @thakis wrote:

I think this might've broken check-clangd: http://45.33.8.238/linux/74320/step_9.txt

The output isn't very useful though, I'm guessing @yln's recent gtest runner change it's not getting printed.

Thanks, I've reverted this change and fortunately it reproduces locally (I had only run the clang tests, oops).

I'm really concerned about the test runner change though: it means bots yield essentially no information from crashes (no test name, no stack trace).
So if a crash occurs in a bot config but not locally, there's no clear way to debug it.

clang/lib/Frontend/PrecompiledPreamble.cpp

TempPCHFile(std::string FilePath) : FilePath(std::move(FilePath)) {
     TemporaryFiles::getInstance().addFile(this->FilePath);
   }
 
   std::string FilePath;
 };
 
 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;
 
   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);
   }
 
   bool BeginSourceFileAction(CompilerInstance &CI) override {
     assert(CI.getLangOpts().CompilingPCH);
     return ASTFrontendAction::BeginSourceFileAction(CI);
   }
 
   bool shouldEraseOutputFiles() override { return !hasEmittedPreamblePCH(); }
   bool hasCodeCompletionSupport() const override { return false; }
   bool hasASTFileSupport() const override { return false; }
   TranslationUnitKind getTranslationUnitKind() override { return TU_Prefix; }
 
 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;
 };
 
 class PrecompilePreambleConsumer : public PCHGenerator {
 public:
   PrecompilePreambleConsumer(PrecompilePreambleAction &Action,
                              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);
     return true;
   }
 
   void HandleTranslationUnit(ASTContext &Ctx) override {
     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());
   }
 
   bool shouldSkipFunctionBody(Decl *D) override {
     return Action.Callbacks.shouldSkipFunctionBody(D);
   }
 
 private:
   PrecompilePreambleAction &Action;
-  std::unique_ptr<raw_ostream> Out;
 };
 
 std::unique_ptr<ASTConsumer>
 PrecompilePreambleAction::CreateASTConsumer(CompilerInstance &CI,
                                             StringRef InFile) {
   std::string Sysroot;
   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 (!OS)
+  if (WritePCHFile) {
+    std::string OutputFile; // unused
+    FileOS = GeneratePCHAction::CreateOutputFile(CI, InFile, OutputFile);
+    if (!FileOS)
       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) {
   if (!Val)
     return false;
   Output = std::move(*Val);
   return true;
 }

 class PrecompiledPreamble::PCHStorage {
 public:
   static std::unique_ptr<PCHStorage> file(std::unique_ptr<TempPCHFile> File) {
     assert(File);
     std::unique_ptr<PCHStorage> S(new PCHStorage());
     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;
   }
 
   enum class Kind { InMemory, TempFile };
   Kind getKind() const {
     if (Memory)
       return Kind::InMemory;
     if (File)
       return Kind::TempFile;
     llvm_unreachable("Neither Memory nor File?");
   }
   llvm::StringRef filePath() const {
     assert(getKind() == Kind::TempFile);
     return File->getFilePath();
   }
   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:
   PCHStorage() = default;
   PCHStorage(const PCHStorage &) = delete;
   PCHStorage &operator=(const PCHStorage &) = delete;
 
-  llvm::Optional<std::string> Memory;
+  std::shared_ptr<PCHBuffer> Memory;
   std::unique_ptr<TempPCHFile> File;
 };
 
 PrecompiledPreamble::~PrecompiledPreamble() = default;
 PrecompiledPreamble::PrecompiledPreamble(PrecompiledPreamble &&) = default;
 PrecompiledPreamble &
 PrecompiledPreamble::operator=(PrecompiledPreamble &&) = default;
 
 llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
     const CompilerInvocation &Invocation,
     const llvm::MemoryBuffer *MainFileBuffer, PreambleBounds Bounds,
     DiagnosticsEngine &Diagnostics,
     IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
     std::shared_ptr<PCHContainerOperations> PCHContainerOps, bool StoreInMemory,
     PreambleCallbacks &Callbacks) {
   assert(VFS && "VFS is null");
 
   auto PreambleInvocation = std::make_shared<CompilerInvocation>(Invocation);
   FrontendOptions &FrontendOpts = PreambleInvocation->getFrontendOpts();
   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.
     std::unique_ptr<TempPCHFile> PreamblePCHFile = TempPCHFile::create();
     if (!PreamblePCHFile)
       return BuildPreambleError::CouldntCreateTempFile;
     Storage = PCHStorage::file(std::move(PreamblePCHFile));
   }
if (PreprocessorOpts.RetainRemappedFileBuffers) {
     PreprocessorOpts.addRemappedFile(MainFilePath, PreambleInputBuffer.get());
   } else {
     // In that case, remapped buffer will be deleted by CompilerInstance on
     // BeginSourceFile, so we call release() to avoid double deletion.
     PreprocessorOpts.addRemappedFile(MainFilePath,
                                      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;
 
   // Performed after BeginSourceFile to ensure Clang->Preprocessor can be
   // referenced in the callback.
   Callbacks.BeforeExecute(*Clang);
 
   std::unique_ptr<PPCallbacks> DelegatedPPCallbacks =
llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
 
   // Run the callbacks.
   Callbacks.AfterExecute(*Clang);
 
   Act->EndSourceFile();
 
   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.
   llvm::StringMap<PrecompiledPreamble::PreambleFileHash> FilesInPreamble;
 
   SourceManager &SourceMgr = Clang->getSourceManager();
   for (auto &Filename : PreambleDepCollector->getDependencies()) {
     auto FileOrErr = Clang->getFileManager().getFile(Filename);