Allow to store precompiled preambles in memory.

[ilya-biryukov]

Summary

These preambles are built by ASTUnit and clangd. Previously, preambles
were always stored on disk.

In-memory preambles are routed back to the compiler as virtual files in
a custom VFS.

Interface of ASTUnit does not allow to use in-memory preambles, as
ASTUnit::CodeComplete receives FileManager as a parameter, so we can't
change VFS used by the compiler inside the CodeComplete method.

A follow-up commit will update clangd in clang-tools-extra to use
in-memory preambles.

[sammccall]

I'm really happy you've made this work! I don't understand it enough to do a meaningful review (keen to learn if you have time for a walkthrough when back in the office).

[klimek]

include/clang/Frontend/PrecompiledPreamble.h

99–100

...ensure *the* preamble...

lib/Frontend/ASTUnit.cpp

1032

Since when are we using the /*ref*/ annotation?

lib/Frontend/PrecompiledPreamble.cpp

206

It looks like we should pass in the output stream, not the storage?

489

This looks a bit like we should push it into the PCHStorage.

[ilya-biryukov]

• Fixed comments.
• Removed /*ref*/ annotations.
• Removed unused "Storage" variable.
• Extract a helper function that properly sets up VFS to access the PCHStorage.

[ilya-biryukov]

lib/Frontend/ASTUnit.cpp

1032

Oh, sorry, that's my thing. Slipped into this change.
Removed those.

lib/Frontend/PrecompiledPreamble.cpp

206

We're not actually using the Storage variable, it's a leftover from previous versions. Removed it.

Or did you mean that we should pass in the output stream directly to PrecompilePreambleAction's constructor?

489

I've extracted a function here to make the code read simpler.
However, I placed it directly into the PrecompiledPreamble instead of PCHStorage to keep PCHStorage responsible for just one thing: managing the variant-like union.

[klimek]

lib/Frontend/PrecompiledPreamble.cpp

206

Yes, I'm generally looking at things that might be better to decide at a higher abstraction level and pass in, rather than having switches for behavior (like InMemStorage) all over the place. Generally, I think we should have a storage (PCHStorage sounds like it was the right abstraction, but perhaps that one currently has a bad name) and the things dealing with that storage shouldn't care whether it's in memory or on the file system.

489

It being called PCHStorage makes it sound like it handles the abstraction for how the preamble is stored. Given that the variant-like union is basically an interface with an implementation switch, I think all switching on it is also the responsibility of that class. Otherwise we'd need another abstraction on top of it?

[ilya-biryukov]

lib/Frontend/PrecompiledPreamble.cpp

206

It's a bit easier to share code with GeneratePCHAction this way in a way that would obviously works (as we call the same functions at the same points in the compilation pipeline, that is in CreateASTConsumer).

PCHStorage is not a public interface and PrecompiledPreamble is exactly the interface that let's you use it without caring whether PCHs are stored in memory or on disk. It also feels ok for it to have the storage-dependent code as part of its own implementation.

489

Abstraction on top of it is PrecompiledPreamble itself. And PCHStorage is just an implementation detail.
Does that sound reasonable?

[klimek]

LG

lib/Frontend/PrecompiledPreamble.cpp

45

erasedOnReboot seems weird for something that we don't want to have on the actual file system. Why do we even want a temp directory?

489

That makes more sense now, thx for explaining. It still seems a bit off to me, but I can't come up with a better solution, so LG.

[ilya-biryukov]

lib/Frontend/PrecompiledPreamble.cpp

45

Yeah, erasedOnReboot seems weird here.
What I really wanted is a filepath that's valid on all platforms and is on an existing drive for Windows, so the value of erasedOnReboot doesn't really matter. system_temp_directory might be the wrong thing to use for that in the first place. Is there a better alternative I'm missing?

[klimek]

lib/Frontend/PrecompiledPreamble.cpp

45

Why does it need to get an existing path / existing drive?

[ilya-biryukov]

lib/Frontend/PrecompiledPreamble.cpp

45

Oh, maybe it doesn't have to be. But I expect that a probability of something breaking is lower if we use an existing path.
Maybe I'm wrong and it will all work out just fine, but that seems to be a rather good compromise adding an extra bit of confidence that things will be working in the long-term.

[klimek]

lib/Frontend/PrecompiledPreamble.cpp

45

I'd on the contrary say that this is more likely to fail. If we don't have a real filesystem, or have a read-only view of a real file system, getting a temp dir might well fail, while we can easily overlay any path we want in memory.

[ilya-biryukov]

lib/Frontend/PrecompiledPreamble.cpp

45

That's true in general, but LLVM's system_temp_directory will return the default (/var/tmp or /tmp on Linux, C:\Temp on Windows) if it fails to find the real temp dir.
So in the very worst case we'll be using hard-coded paths.

On the other hand, if that's one of the possible behaviors, we'd probably want to not break in that case anyway. I'll update the patch to follow your suggestions. Thanks!

[ilya-biryukov]

• Use a hard-coded virtual path for in-memory PCHs instead of system_temp_directory.

[ilya-biryukov]

• Removed redundant #include.

[ilya-biryukov]

lib/Frontend/PrecompiledPreamble.cpp

45

This should do the trick.

[mgorny]

Please revert this commit. You've just broken all the stand-alone builds of clang.

cfe/trunk/lib/Frontend/PrecompiledPreamble.cpp

27 ↗

(On Diff #123190)

This is private LLVM header which is not available when doing stand-alone builds.

[ilya-biryukov]

In D39842#927578, @mgorny wrote:

Please revert this commit. You've just broken all the stand-alone builds of clang.

Sorry about that. Should be fixed in rL318514

include/clang/Frontend/FrontendActions.h

protected:
   bool hasASTFileSupport() const override { return false; }
 
   bool shouldEraseOutputFiles() override;
 
 public:
   /// \brief Compute the AST consumer arguments that will be used to
   /// create the PCHGenerator instance returned by CreateASTConsumer.
   ///
-  /// \returns true if an error occurred, false otherwise.
-  static std::unique_ptr<raw_pwrite_stream>
-  ComputeASTConsumerArguments(CompilerInstance &CI, StringRef InFile,
-                              std::string &Sysroot, std::string &OutputFile);
+  /// \returns false if an error occurred, true otherwise.
+  static bool ComputeASTConsumerArguments(CompilerInstance &CI,
+                                          std::string &Sysroot);
+
+  /// \brief Creates file to write the PCH into and returns a stream to write it
+  /// into. On error, returns null.
+  static std::unique_ptr<llvm::raw_pwrite_stream>
+  CreateOutputFile(CompilerInstance &CI, StringRef InFile,
+                   std::string &OutputFile);
 
   bool BeginSourceFileAction(CompilerInstance &CI) override;
 };
 
 class GenerateModuleAction : public ASTFrontendAction {
   virtual std::unique_ptr<raw_pwrite_stream>
   CreateOutputFile(CompilerInstance &CI, StringRef InFile) = 0;
 

include/clang/Frontend/PrecompiledPreamble.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_FRONTEND_PRECOMPILED_PREAMBLE_H
 #define LLVM_CLANG_FRONTEND_PRECOMPILED_PREAMBLE_H
 
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/Support/AlignOf.h"
 #include "llvm/Support/MD5.h"
 #include <memory>
 #include <system_error>
 #include <type_traits>
 
 namespace llvm {
 class MemoryBuffer;
 }
PreambleBounds ComputePreambleBounds(const LangOptions &LangOpts,
                                      unsigned MaxLines);
 
 class PreambleCallbacks;
 
 /// A class holding a PCH and all information to check whether it is valid to
 /// reuse the PCH for the subsequent runs. Use BuildPreamble to create PCH and
 /// CanReusePreamble + AddImplicitPreamble to make use of it.
 class PrecompiledPreamble {
-  class TempPCHFile;
+  class PCHStorage;
   struct PreambleFileHash;
 
 public:
   /// \brief Try to build PrecompiledPreamble for \p Invocation. See
   /// BuildPreambleError for possible error codes.
   ///
   /// \param Invocation Original CompilerInvocation with options to compile the
   /// file.
   ///
   /// \param MainFileBuffer Buffer with the contents of the main file.
   ///
   /// \param Bounds Bounds of the preamble, result of calling
   /// ComputePreambleBounds.
   ///
   /// \param Diagnostics Diagnostics engine to be used while building the
   /// preamble.
   ///
   /// \param VFS An instance of vfs::FileSystem to be used for file
   /// accesses.
   ///
   /// \param PCHContainerOps An instance of PCHContainerOperations.
   ///
+  /// \param StoreInMemory Store PCH in memory. If false, PCH will be stored in
+  /// a temporary file.
+  ///
   /// \param Callbacks A set of callbacks to be executed when building
   /// the preamble.
   static llvm::ErrorOr<PrecompiledPreamble>
   Build(const CompilerInvocation &Invocation,
         const llvm::MemoryBuffer *MainFileBuffer, PreambleBounds Bounds,
         DiagnosticsEngine &Diagnostics, IntrusiveRefCntPtr<vfs::FileSystem> VFS,
         std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-        PreambleCallbacks &Callbacks);
+        bool StoreInMemory, PreambleCallbacks &Callbacks);
 
   PrecompiledPreamble(PrecompiledPreamble &&) = default;
   PrecompiledPreamble &operator=(PrecompiledPreamble &&) = default;
 
   /// PreambleBounds used to build the preamble.
   PreambleBounds getBounds() const;
 
-  /// The temporary file path at which the preamble PCH was placed.
-  StringRef GetPCHPath() const { return PCHFile.getFilePath(); }
-
   /// Check whether PrecompiledPreamble can be reused for the new contents(\p
   /// MainFileBuffer) of the main file.
   bool CanReuse(const CompilerInvocation &Invocation,
                 const llvm::MemoryBuffer *MainFileBuffer, PreambleBounds Bounds,
                 vfs::FileSystem *VFS) const;
 
   /// Changes options inside \p CI to use PCH from this preamble. Also remaps
-  /// main file to \p MainFileBuffer.
+  /// main file to \p MainFileBuffer and updates \p VFS to ensure the preamble
+  /// is accessible.
   void AddImplicitPreamble(CompilerInvocation &CI,
+                           IntrusiveRefCntPtr<vfs::FileSystem> &VFS,
                            llvm::MemoryBuffer *MainFileBuffer) const;
 
 private:
-  PrecompiledPreamble(TempPCHFile PCHFile, std::vector<char> PreambleBytes,
+  PrecompiledPreamble(PCHStorage Storage, std::vector<char> PreambleBytes,
                       bool PreambleEndsAtStartOfLine,
                       llvm::StringMap<PreambleFileHash> FilesInPreamble);
 
   /// A temp file that would be deleted on destructor call. If destructor is not
   /// called for any reason, the file will be deleted at static objects'
   /// destruction.
   /// An assertion will fire if two TempPCHFiles are created with the same name,
   /// so it's not intended to be used outside preamble-handling.
private:
 
   private:
     void RemoveFileIfPresent();
 
   private:
     llvm::Optional<std::string> FilePath;
   };
 
+  class InMemoryPreamble {
+  public:
+    std::string Data;
+  };
+
+  class PCHStorage {
+  public:
+    enum class Kind { Empty, InMemory, TempFile };
+
+    PCHStorage() = default;
+    PCHStorage(TempPCHFile File);
+    PCHStorage(InMemoryPreamble Memory);
+
+    PCHStorage(const PCHStorage &) = delete;
+    PCHStorage &operator=(const PCHStorage &) = delete;
+
+    PCHStorage(PCHStorage &&Other);
+    PCHStorage &operator=(PCHStorage &&Other);
+
+    ~PCHStorage();
+
+    Kind getKind() const;
+
+    TempPCHFile &asFile();
+    const TempPCHFile &asFile() const;
+
+    InMemoryPreamble &asMemory();
+    const InMemoryPreamble &asMemory() const;
+
+  private:
+    void destroy();
+    void setEmpty();
+
+  private:
+    Kind StorageKind = Kind::Empty;
+    llvm::AlignedCharArrayUnion<TempPCHFile, InMemoryPreamble> Storage = {};
+  };
+
   /// Data used to determine if a file used in the preamble has been changed.
   struct PreambleFileHash {
     /// All files have size set.
     off_t Size = 0;
 
     /// Modification time is set for files that are on disk.  For memory
     /// buffers it is zero.
     time_t ModTime = 0;
friend bool operator==(const PreambleFileHash &LHS,
              LHS.MD5 == RHS.MD5;
     }
     friend bool operator!=(const PreambleFileHash &LHS,
                            const PreambleFileHash &RHS) {
       return !(LHS == RHS);
     }
   };
 
-  /// Manages the lifetime of temporary file that stores a PCH.
-  TempPCHFile PCHFile;
+  /// Sets up the PreprocessorOptions and changes VFS, so that PCH stored in \p
+  /// Storage is accessible to clang. This method is an implementation detail of
+  /// AddImplicitPreamble.
+  static void setupPreambleStorage(const PCHStorage &Storage,
+                                   PreprocessorOptions &PreprocessorOpts,
+                                   IntrusiveRefCntPtr<vfs::FileSystem> &VFS);
+
+  /// Manages the memory buffer or temporary file that stores the PCH.
+  PCHStorage Storage;
   /// Keeps track of the files that were used when computing the
   /// preamble, with both their buffer size and their modification time.
   ///
   /// If any of the files have changed from one compile to the next,
   /// the preamble must be thrown away.
   llvm::StringMap<PreambleFileHash> FilesInPreamble;
   /// The contents of the file that was used to precompile the preamble. Only
   /// contains first PreambleBounds::Size bytes. Used to compare if the relevant

lib/Frontend/ASTUnit.cpp

static void checkAndSanitizeDiags(SmallVectorImpl<StoredDiagnostic> &
   for (StoredDiagnostic &SD : StoredDiagnostics) {
     if (SD.getLocation().isValid()) {
       FullSourceLoc Loc(SD.getLocation(), SM);
       SD.setLocation(Loc);
     }
   }
 }
 
-static IntrusiveRefCntPtr<vfs::FileSystem> createVFSOverlayForPreamblePCH(
-    StringRef PCHFilename,
-    IntrusiveRefCntPtr<vfs::FileSystem> RealFS,
-    IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
-  // We want only the PCH file from the real filesystem to be available,
-  // so we create an in-memory VFS with just that and overlay it on top.
-  auto Buf = RealFS->getBufferForFile(PCHFilename);
-  if (!Buf)
-    return VFS;
-  IntrusiveRefCntPtr<vfs::InMemoryFileSystem>
-      PCHFS(new vfs::InMemoryFileSystem());
-  PCHFS->addFile(PCHFilename, 0, std::move(*Buf));
-  IntrusiveRefCntPtr<vfs::OverlayFileSystem>
-      Overlay(new vfs::OverlayFileSystem(VFS));
-  Overlay->pushOverlay(PCHFS);
-  return Overlay;
-}
-
 /// Parse the source file into a translation unit using the given compiler
 /// invocation, replacing the current translation unit.
 ///
 /// \returns True if a failure occurred that causes the ASTUnit not to
 /// contain any translation-unit information, false otherwise.
 bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
                     std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer,
                     IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   if (!Invocation)
     return true;
 
+  auto CCInvocation = std::make_shared<CompilerInvocation>(*Invocation);
+  if (OverrideMainBuffer) {
+    assert(Preamble &&
+           "No preamble was built, but OverrideMainBuffer is not null");
+    IntrusiveRefCntPtr<vfs::FileSystem> OldVFS = VFS;
+    Preamble->AddImplicitPreamble(*CCInvocation, VFS, OverrideMainBuffer.get());
+    if (OldVFS != VFS && FileMgr) {
+      assert(OldVFS == FileMgr->getVirtualFileSystem() &&
+             "VFS passed to Parse and VFS in FileMgr are different");
+      FileMgr = new FileManager(FileMgr->getFileSystemOpts(), VFS);
+    }
+  }
+
   // Create the compiler instance to use for building the AST.
   std::unique_ptr<CompilerInstance> Clang(
       new CompilerInstance(std::move(PCHContainerOps)));
   if (FileMgr && VFS) {
     assert(VFS == FileMgr->getVirtualFileSystem() &&
            "VFS passed to Parse and VFS in FileMgr are different");
   } else if (VFS) {
     Clang->setVirtualFileSystem(VFS);
   }
 
-  // Make sure we can access the PCH file even if we're using a VFS
-  if (!VFS && FileMgr)
-    VFS = FileMgr->getVirtualFileSystem();
-  IntrusiveRefCntPtr<vfs::FileSystem> RealFS = vfs::getRealFileSystem();
-  if (OverrideMainBuffer && VFS && RealFS && VFS != RealFS &&
-      !VFS->exists(Preamble->GetPCHPath())) {
-    // We have a slight inconsistency here -- we're using the VFS to
-    // read files, but the PCH was generated in the real file system.
-    VFS = createVFSOverlayForPreamblePCH(Preamble->GetPCHPath(), RealFS, VFS);
-    if (FileMgr) {
-      FileMgr = new FileManager(FileMgr->getFileSystemOpts(), VFS);
-      Clang->setFileManager(FileMgr.get());
-    }
-    else {
-      Clang->setVirtualFileSystem(VFS);
-    }
-  }
-
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<CompilerInstance>
     CICleanup(Clang.get());
 
-  Clang->setInvocation(std::make_shared<CompilerInvocation>(*Invocation));
+  Clang->setInvocation(CCInvocation);
   OriginalSourceFile = Clang->getFrontendOpts().Inputs[0].getFile();
     
   // Set up diagnostics, capturing any diagnostics that would
   // otherwise be dropped.
   Clang->setDiagnostics(&getDiagnostics());
   
   // Create the target instance.
   Clang->setTarget(TargetInfo::CreateTargetInfo(
bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   Clang->setFileManager(&getFileManager());
   
   // Create the source manager.
   Clang->setSourceManager(&getSourceManager());
   
   // If the main file has been overridden due to the use of a preamble,
   // make that override happen and introduce the preamble.
   if (OverrideMainBuffer) {
-    assert(Preamble && "No preamble was built, but OverrideMainBuffer is not null");
-    Preamble->AddImplicitPreamble(Clang->getInvocation(), OverrideMainBuffer.get());
-    
     // The stored diagnostic has the old source manager in it; update
     // the locations to refer into the new source manager. Since we've
     // been careful to make sure that the source manager's state
     // before and after are identical, so that we can reuse the source
     // location itself.
     checkAndSanitizeDiags(StoredDiagnostics, getSourceManager());
 
     // Keep track of the override buffer;
if (CaptureDiagnostics)
                       &NewPreambleDiagsStandalone);
 
     // We did not previously compute a preamble, or it can't be reused anyway.
     SimpleTimer PreambleTimer(WantTiming);
     PreambleTimer.setOutput("Precompiling preamble");
 
     llvm::ErrorOr<PrecompiledPreamble> NewPreamble = PrecompiledPreamble::Build(
         PreambleInvocationIn, MainFileBuffer.get(), Bounds, *Diagnostics, VFS,
-        PCHContainerOps, Callbacks);
+        PCHContainerOps, /*StoreInMemory=*/false, Callbacks);
     if (NewPreamble) {
       Preamble = std::move(*NewPreamble);
       PreambleRebuildCounter = 1;
     } else {
       switch (static_cast<BuildPreambleError>(NewPreamble.getError().value())) {
       case BuildPreambleError::CouldntCreateTempFile:
       case BuildPreambleError::PreambleIsEmpty:
         // Try again next time.
if (CompleteFileStatus) {
               PCHContainerOps, Inv, VFS, false, Line - 1);
       }
     }
   }
 
   // If the main file has been overridden due to the use of a preamble,
   // make that override happen and introduce the preamble.
   if (OverrideMainBuffer) {
-    assert(Preamble && "No preamble was built, but OverrideMainBuffer is not null");
-    Preamble->AddImplicitPreamble(Clang->getInvocation(), OverrideMainBuffer.get());
+    assert(Preamble &&
+           "No preamble was built, but OverrideMainBuffer is not null");
+
+    auto VFS = FileMgr.getVirtualFileSystem();
+    Preamble->AddImplicitPreamble(Clang->getInvocation(), VFS,
+                                  OverrideMainBuffer.get());
+    // FIXME: there is no way to update VFS if it was changed by
+    // AddImplicitPreamble as FileMgr is accepted as a parameter by this method.
+    // We use on-disk preambles instead and rely on FileMgr's VFS to ensure the
+    // PCH files are always readable.
     OwnedBuffers.push_back(OverrideMainBuffer.release());
   } else {
     PreprocessorOpts.PrecompiledPreambleBytes.first = 0;
     PreprocessorOpts.PrecompiledPreambleBytes.second = false;
   }
 
   // Disable the preprocessing record if modules are not enabled.
   if (!Clang->getLangOpts().Modules)

lib/Frontend/FrontendActions.cpp

 DeclContextPrintAction::CreateASTConsumer(CompilerInstance &CI,
                                           StringRef InFile) {
   return CreateDeclContextPrinter();
 }
 
 std::unique_ptr<ASTConsumer>
 GeneratePCHAction::CreateASTConsumer(CompilerInstance &CI, StringRef InFile) {
   std::string Sysroot;
+  if (!ComputeASTConsumerArguments(CI, /*ref*/ Sysroot))
+    return nullptr;
+
   std::string OutputFile;
   std::unique_ptr<raw_pwrite_stream> OS =
-      ComputeASTConsumerArguments(CI, InFile, Sysroot, OutputFile);
+      CreateOutputFile(CI, InFile, /*ref*/ OutputFile);
   if (!OS)
     return nullptr;
 
   if (!CI.getFrontendOpts().RelocatablePCH)
     Sysroot.clear();
 
   auto Buffer = std::make_shared<PCHBuffer>();
   std::vector<std::unique_ptr<ASTConsumer>> Consumers;
   Consumers.push_back(llvm::make_unique<PCHGenerator>(
                         CI.getPreprocessor(), OutputFile, Sysroot,
                         Buffer, CI.getFrontendOpts().ModuleFileExtensions,
       /*AllowASTWithErrors*/CI.getPreprocessorOpts().AllowPCHWithCompilerErrors,
                         /*IncludeTimestamps*/
                           +CI.getFrontendOpts().IncludeTimestamps));
   Consumers.push_back(CI.getPCHContainerWriter().CreatePCHContainerGenerator(
       CI, InFile, OutputFile, std::move(OS), Buffer));
 
   return llvm::make_unique<MultiplexConsumer>(std::move(Consumers));
 }
 
-std::unique_ptr<raw_pwrite_stream>
-GeneratePCHAction::ComputeASTConsumerArguments(CompilerInstance &CI,
-                                               StringRef InFile,
-                                               std::string &Sysroot,
-                                               std::string &OutputFile) {
+bool GeneratePCHAction::ComputeASTConsumerArguments(CompilerInstance &CI,
+                                                    std::string &Sysroot) {
   Sysroot = CI.getHeaderSearchOpts().Sysroot;
   if (CI.getFrontendOpts().RelocatablePCH && Sysroot.empty()) {
     CI.getDiagnostics().Report(diag::err_relocatable_without_isysroot);
-    return nullptr;
+    return false;
   }
 
+  return true;
+}
+
+std::unique_ptr<llvm::raw_pwrite_stream>
+GeneratePCHAction::CreateOutputFile(CompilerInstance &CI, StringRef InFile,
+                                    std::string &OutputFile) {
   // We use createOutputFile here because this is exposed via libclang, and we
   // must disable the RemoveFileOnSignal behavior.
   // We use a temporary to avoid race conditions.
   std::unique_ptr<raw_pwrite_stream> OS =
       CI.createOutputFile(CI.getFrontendOpts().OutputFile, /*Binary=*/true,
                           /*RemoveFileOnSignal=*/false, InFile,
                           /*Extension=*/"", /*useTemporary=*/true);
   if (!OS)

lib/Frontend/PrecompiledPreamble.cpp

 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/FrontendOptions.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Serialization/ASTWriter.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
+#include "llvm/Config/config.h"
 #include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/MutexGuard.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/Process.h"
 
+#include <utility>
+
 using namespace clang;
 
 namespace {
 
+StringRef getInMemoryPreamblePath() {
+#if defined(LLVM_ON_UNIX)
+  return "/__clang_tmp/___clang_inmemory_preamble___";
+#elif defined(LLVM_ON_WIN32)
+  return "C:\\__clang_tmp\\___clang_inmemory_preamble___";
+#else
+#warning "Unknown platform. Defaulting to UNIX-style paths for in-memory PCHs"
+  return "/__clang_tmp/___clang_inmemory_preamble___";
+#endif
+}
+
+IntrusiveRefCntPtr<vfs::FileSystem>
+createVFSOverlayForPreamblePCH(StringRef PCHFilename,
+                               std::unique_ptr<llvm::MemoryBuffer> PCHBuffer,
+                               IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
+  // We want only the PCH file from the real filesystem to be available,
+  // so we create an in-memory VFS with just that and overlay it on top.
+  IntrusiveRefCntPtr<vfs::InMemoryFileSystem> PCHFS(
+      new vfs::InMemoryFileSystem());
+  PCHFS->addFile(PCHFilename, 0, std::move(PCHBuffer));
+  IntrusiveRefCntPtr<vfs::OverlayFileSystem> Overlay(
+      new vfs::OverlayFileSystem(VFS));
+  Overlay->pushOverlay(PCHFS);
+  return Overlay;
+}
+
 /// Keeps a track of files to be deleted in destructor.
 class TemporaryFiles {
 public:
   // A static instance to be used by all clients.
   static TemporaryFiles &getInstance();
 
 private:
   // Disallow constructing the class directly.
public:
   }
 
 private:
   PreambleCallbacks &Callbacks;
 };
 
 class PrecompilePreambleAction : public ASTFrontendAction {
 public:
-  PrecompilePreambleAction(PreambleCallbacks &Callbacks)
-      : Callbacks(Callbacks) {}
+  PrecompilePreambleAction(std::string *InMemStorage,
+                           PreambleCallbacks &Callbacks)
+      : InMemStorage(InMemStorage), Callbacks(Callbacks) {}
 
   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                  StringRef InFile) override;
 
   bool hasEmittedPreamblePCH() const { return HasEmittedPreamblePCH; }
 
   void setEmittedPreamblePCH(ASTWriter &Writer) {
     this->HasEmittedPreamblePCH = true;
     Callbacks.AfterPCHEmitted(Writer);
   }
 
   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;
   PreambleCallbacks &Callbacks;
 };
 
 class PrecompilePreambleConsumer : public PCHGenerator {
 public:
   PrecompilePreambleConsumer(PrecompilePreambleAction &Action,
                              const Preprocessor &PP, StringRef isysroot,
                              std::unique_ptr<raw_ostream> Out)

 
 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 = llvm::make_unique<llvm::raw_string_ostream>(*InMemStorage);
+  } else {
     std::string OutputFile;
-  std::unique_ptr<raw_ostream> OS =
-      GeneratePCHAction::ComputeASTConsumerArguments(CI, InFile, Sysroot,
-                                                     OutputFile);
+    OS = GeneratePCHAction::CreateOutputFile(CI, InFile, OutputFile);
+  }
   if (!OS)
     return nullptr;
 
   if (!CI.getFrontendOpts().RelocatablePCH)
     Sysroot.clear();
 
   CI.getPreprocessor().addPPCallbacks(
       llvm::make_unique<PreambleMacroCallbacks>(Callbacks));
PreambleBounds clang::ComputePreambleBounds(const LangOptions &LangOpts,
                                             unsigned MaxLines) {
   return Lexer::ComputePreamble(Buffer->getBuffer(), LangOpts, MaxLines);
 }
 
 llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
     const CompilerInvocation &Invocation,
     const llvm::MemoryBuffer *MainFileBuffer, PreambleBounds Bounds,
     DiagnosticsEngine &Diagnostics, IntrusiveRefCntPtr<vfs::FileSystem> VFS,
-    std::shared_ptr<PCHContainerOperations> PCHContainerOps,
+    std::shared_ptr<PCHContainerOperations> PCHContainerOps, bool StoreInMemory,
     PreambleCallbacks &Callbacks) {
   assert(VFS && "VFS is null");
 
   if (!Bounds.Size)
     return BuildPreambleError::PreambleIsEmpty;
 
   auto PreambleInvocation = std::make_shared<CompilerInvocation>(Invocation);
   FrontendOptions &FrontendOpts = PreambleInvocation->getFrontendOpts();
   PreprocessorOptions &PreprocessorOpts =
       PreambleInvocation->getPreprocessorOpts();
 
+  llvm::Optional<TempPCHFile> TempFile;
+  if (!StoreInMemory) {
     // Create a temporary file for the precompiled preamble. In rare
     // circumstances, this can fail.
     llvm::ErrorOr<PrecompiledPreamble::TempPCHFile> PreamblePCHFile =
         PrecompiledPreamble::TempPCHFile::CreateNewPreamblePCHFile();
     if (!PreamblePCHFile)
       return BuildPreambleError::CouldntCreateTempFile;
+    TempFile = std::move(*PreamblePCHFile);
+  }
+
+  PCHStorage Storage = StoreInMemory ? PCHStorage(InMemoryPreamble())
+                                     : PCHStorage(std::move(*TempFile));
 
   // Save the preamble text for later; we'll need to compare against it for
   // subsequent reparses.
   std::vector<char> PreambleBytes(MainFileBuffer->getBufferStart(),
                                   MainFileBuffer->getBufferStart() +
                                       Bounds.Size);
   bool PreambleEndsAtStartOfLine = Bounds.PreambleEndsAtStartOfLine;
 
   // Tell the compiler invocation to generate a temporary precompiled header.
   FrontendOpts.ProgramAction = frontend::GeneratePCH;
-  // FIXME: Generate the precompiled header into memory?
-  FrontendOpts.OutputFile = PreamblePCHFile->getFilePath();
+  FrontendOpts.OutputFile = StoreInMemory ? getInMemoryPreamblePath()
+                                          : Storage.asFile().getFilePath();
   PreprocessorOpts.PrecompiledPreambleBytes.first = 0;
   PreprocessorOpts.PrecompiledPreambleBytes.second = false;
   // Inform preprocessor to record conditional stack when building the preamble.
   PreprocessorOpts.GeneratePreamble = true;
 
   // Create the compiler instance to use for building the precompiled preamble.
   std::unique_ptr<CompilerInstance> Clang(
       new CompilerInstance(std::move(PCHContainerOps)));
llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
   } 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(Callbacks));
+  Act.reset(new PrecompilePreambleAction(
+      StoreInMemory ? &Storage.asMemory().Data : nullptr, Callbacks));
   if (!Act->BeginSourceFile(*Clang.get(), Clang->getFrontendOpts().Inputs[0]))
     return BuildPreambleError::BeginSourceFileFailed;
 
   Act->Execute();
 
   // Run the callbacks.
   Callbacks.AfterExecute(*Clang);
 
if (time_t ModTime = File->getModificationTime()) {
                                                                ModTime);
     } else {
       llvm::MemoryBuffer *Buffer = SourceMgr.getMemoryBufferForFile(File);
       FilesInPreamble[File->getName()] =
           PrecompiledPreamble::PreambleFileHash::createForMemoryBuffer(Buffer);
     }
   }
 
-  return PrecompiledPreamble(
-      std::move(*PreamblePCHFile), std::move(PreambleBytes),
-      PreambleEndsAtStartOfLine, std::move(FilesInPreamble));
+  return PrecompiledPreamble(std::move(Storage), std::move(PreambleBytes),
+                             PreambleEndsAtStartOfLine,
+                             std::move(FilesInPreamble));
 }
 
 PreambleBounds PrecompiledPreamble::getBounds() const {
   return PreambleBounds(PreambleBytes.size(), PreambleEndsAtStartOfLine);
 }
 
 bool PrecompiledPreamble::CanReuse(const CompilerInvocation &Invocation,
                                    const llvm::MemoryBuffer *MainFileBuffer,
if (Status.getSize() != uint64_t(F.second.Size) ||
         llvm::sys::toTimeT(Status.getLastModificationTime()) !=
             F.second.ModTime)
       return false;
   }
   return true;
 }
 
 void PrecompiledPreamble::AddImplicitPreamble(
-    CompilerInvocation &CI, llvm::MemoryBuffer *MainFileBuffer) const {
+    CompilerInvocation &CI, IntrusiveRefCntPtr<vfs::FileSystem> &VFS,
+    llvm::MemoryBuffer *MainFileBuffer) const {
+  assert(VFS && "VFS must not be null");
+
   auto &PreprocessorOpts = CI.getPreprocessorOpts();
 
+  // Remap main file to point to MainFileBuffer.
+  auto MainFilePath = CI.getFrontendOpts().Inputs[0].getFile();
+  PreprocessorOpts.addRemappedFile(MainFilePath, MainFileBuffer);
+
   // Configure ImpicitPCHInclude.
   PreprocessorOpts.PrecompiledPreambleBytes.first = PreambleBytes.size();
   PreprocessorOpts.PrecompiledPreambleBytes.second = PreambleEndsAtStartOfLine;
-  PreprocessorOpts.ImplicitPCHInclude = PCHFile.getFilePath();
   PreprocessorOpts.DisablePCHValidation = true;
 
-  // Remap main file to point to MainFileBuffer.
-  auto MainFilePath = CI.getFrontendOpts().Inputs[0].getFile();
-  PreprocessorOpts.addRemappedFile(MainFilePath, MainFileBuffer);
+  setupPreambleStorage(Storage, PreprocessorOpts, VFS);
 }
 
 PrecompiledPreamble::PrecompiledPreamble(
-    TempPCHFile PCHFile, std::vector<char> PreambleBytes,
+    PCHStorage Storage, std::vector<char> PreambleBytes,
     bool PreambleEndsAtStartOfLine,
     llvm::StringMap<PreambleFileHash> FilesInPreamble)
-    : PCHFile(std::move(PCHFile)), FilesInPreamble(std::move(FilesInPreamble)),
+    : Storage(std::move(Storage)), FilesInPreamble(std::move(FilesInPreamble)),
       PreambleBytes(std::move(PreambleBytes)),
-      PreambleEndsAtStartOfLine(PreambleEndsAtStartOfLine) {}
+      PreambleEndsAtStartOfLine(PreambleEndsAtStartOfLine) {
+  assert(this->Storage.getKind() != PCHStorage::Kind::Empty);
+}
 
 llvm::ErrorOr<PrecompiledPreamble::TempPCHFile>
 PrecompiledPreamble::TempPCHFile::CreateNewPreamblePCHFile() {
   // FIXME: This is a hack so that we can override the preamble file during
   // crash-recovery testing, which is the only case where the preamble files
   // are not necessarily cleaned up.
   const char *TmpFile = ::getenv("CINDEXTEST_PREAMBLE_FILE");
   if (TmpFile)
     return TempPCHFile::createFromCustomPath(TmpFile);
   return TempPCHFile::createInSystemTempDir("preamble", "pch");
 }
 
 llvm::ErrorOr<PrecompiledPreamble::TempPCHFile>
 PrecompiledPreamble::TempPCHFile::createInSystemTempDir(const Twine &Prefix,
                                                         StringRef Suffix) {
   llvm::SmallString<64> File;
   // Using a version of createTemporaryFile with a file descriptor guarantees
   // that we would never get a race condition in a multi-threaded setting (i.e.,
   // multiple threads getting the same temporary path).
   int FD;
-  auto EC = llvm::sys::fs::createTemporaryFile(Prefix, Suffix, /*ref*/ FD,
-                                               /*ref*/ File);
+  auto EC = llvm::sys::fs::createTemporaryFile(Prefix, Suffix, FD, File);
   if (EC)
     return EC;
   // We only needed to make sure the file exists, close the file right away.
   llvm::sys::Process::SafelyCloseFileDescriptor(FD);
   return TempPCHFile(std::move(File).str());
 }
 
 llvm::ErrorOr<PrecompiledPreamble::TempPCHFile>
void PrecompiledPreamble::TempPCHFile::RemoveFileIfPresent() {
   }
 }
 
 llvm::StringRef PrecompiledPreamble::TempPCHFile::getFilePath() const {
   assert(FilePath && "TempPCHFile doesn't have a FilePath. Had it been moved?");
   return *FilePath;
 }
 
+PrecompiledPreamble::PCHStorage::PCHStorage(TempPCHFile File)
+    : StorageKind(Kind::TempFile) {
+  new (&asFile()) TempPCHFile(std::move(File));
+}
+
+PrecompiledPreamble::PCHStorage::PCHStorage(InMemoryPreamble Memory)
+    : StorageKind(Kind::InMemory) {
+  new (&asMemory()) InMemoryPreamble(std::move(Memory));
+}
+
+PrecompiledPreamble::PCHStorage::PCHStorage(PCHStorage &&Other) : PCHStorage() {
+  *this = std::move(Other);
+}
+
+PrecompiledPreamble::PCHStorage &PrecompiledPreamble::PCHStorage::
+operator=(PCHStorage &&Other) {
+  destroy();
+
+  StorageKind = Other.StorageKind;
+  switch (StorageKind) {
+  case Kind::Empty:
+    // do nothing;
+    break;
+  case Kind::TempFile:
+    new (&asFile()) TempPCHFile(std::move(Other.asFile()));
+    break;
+  case Kind::InMemory:
+    new (&asMemory()) InMemoryPreamble(std::move(Other.asMemory()));
+    break;
+  }
+
+  Other.setEmpty();
+  return *this;
+}
+
+PrecompiledPreamble::PCHStorage::~PCHStorage() { destroy(); }
+
+PrecompiledPreamble::PCHStorage::Kind
+PrecompiledPreamble::PCHStorage::getKind() const {
+  return StorageKind;
+}
+
+PrecompiledPreamble::TempPCHFile &PrecompiledPreamble::PCHStorage::asFile() {
+  assert(getKind() == Kind::TempFile);
+  return *reinterpret_cast<TempPCHFile *>(Storage.buffer);
+}
+
+const PrecompiledPreamble::TempPCHFile &
+PrecompiledPreamble::PCHStorage::asFile() const {
+  return const_cast<PCHStorage *>(this)->asFile();
+}
+
+PrecompiledPreamble::InMemoryPreamble &
+PrecompiledPreamble::PCHStorage::asMemory() {
+  assert(getKind() == Kind::InMemory);
+  return *reinterpret_cast<InMemoryPreamble *>(Storage.buffer);
+}
+
+const PrecompiledPreamble::InMemoryPreamble &
+PrecompiledPreamble::PCHStorage::asMemory() const {
+  return const_cast<PCHStorage *>(this)->asMemory();
+}
+
+void PrecompiledPreamble::PCHStorage::destroy() {
+  switch (StorageKind) {
+  case Kind::Empty:
+    return;
+  case Kind::TempFile:
+    asFile().~TempPCHFile();
+    return;
+  case Kind::InMemory:
+    asMemory().~InMemoryPreamble();
+    return;
+  }
+}
+
+void PrecompiledPreamble::PCHStorage::setEmpty() {
+  destroy();
+  StorageKind = Kind::Empty;
+}
+
 PrecompiledPreamble::PreambleFileHash
 PrecompiledPreamble::PreambleFileHash::createForFile(off_t Size,
                                                      time_t ModTime) {
   PreambleFileHash Result;
   Result.Size = Size;
   Result.ModTime = ModTime;
   Result.MD5 = {};
   return Result;
 }
 
 PrecompiledPreamble::PreambleFileHash
 PrecompiledPreamble::PreambleFileHash::createForMemoryBuffer(
     const llvm::MemoryBuffer *Buffer) {
   PreambleFileHash Result;
   Result.Size = Buffer->getBufferSize();
   Result.ModTime = 0;
 
   llvm::MD5 MD5Ctx;
   MD5Ctx.update(Buffer->getBuffer().data());
   MD5Ctx.final(Result.MD5);
 
   return Result;
 }
 
+void PrecompiledPreamble::setupPreambleStorage(
+    const PCHStorage &Storage, PreprocessorOptions &PreprocessorOpts,
+    IntrusiveRefCntPtr<vfs::FileSystem> &VFS) {
+  if (Storage.getKind() == PCHStorage::Kind::TempFile) {
+    const TempPCHFile &PCHFile = Storage.asFile();
+    PreprocessorOpts.ImplicitPCHInclude = PCHFile.getFilePath();
+
+    // Make sure we can access the PCH file even if we're using a VFS
+    IntrusiveRefCntPtr<vfs::FileSystem> RealFS = vfs::getRealFileSystem();
+    auto PCHPath = PCHFile.getFilePath();
+    if (VFS == RealFS || VFS->exists(PCHPath))
+      return;
+    auto Buf = RealFS->getBufferForFile(PCHPath);
+    if (!Buf) {
+      // We can't read the file even from RealFS, this is clearly an error,
+      // but we'll just leave the current VFS as is and let clang's code
+      // figure out what to do with missing PCH.
+      return;
+    }
+
+    // We have a slight inconsistency here -- we're using the VFS to
+    // read files, but the PCH was generated in the real file system.
+    VFS = createVFSOverlayForPreamblePCH(PCHPath, std::move(*Buf), VFS);
+  } else {
+    assert(Storage.getKind() == PCHStorage::Kind::InMemory);
+    // For in-memory preamble, we have to provide a VFS overlay that makes it
+    // accessible.
+    StringRef PCHPath = getInMemoryPreamblePath();
+    PreprocessorOpts.ImplicitPCHInclude = PCHPath;
+
+    // FIMXE(ibiryukov): Preambles can be large. We should allow shared access
+    // to the preamble data instead of copying it here.
+    auto Buf = llvm::MemoryBuffer::getMemBufferCopy(Storage.asMemory().Data);
+    VFS = createVFSOverlayForPreamblePCH(PCHPath, std::move(Buf), VFS);
+  }
+}
+
 void PreambleCallbacks::AfterExecute(CompilerInstance &CI) {}
 void PreambleCallbacks::AfterPCHEmitted(ASTWriter &Writer) {}
 void PreambleCallbacks::HandleTopLevelDecl(DeclGroupRef DG) {}
 void PreambleCallbacks::HandleMacroDefined(const Token &MacroNameTok,
                                            const MacroDirective *MD) {}
 
 std::error_code clang::make_error_code(BuildPreambleError Error) {
   return std::error_code(static_cast<int>(Error), BuildPreambleErrorCategory());