[Frontend] Cache preamble-related data

[yvvan]

Summary

In case two translation units are created for the same file - reuse preamble data to reduce memory and save time on extra preamble generation.

[ilya-biryukov]

I would argue this should be handled by the clients instead. Adding global state and locking is complicated. (And ASTUnit is complicated enough).

What are the use-cases of creating multiple ASTUnit inside the same process for the same file? Which clients do that and why they can't have a single ASTUnit per file?

[yvvan]

@ilya-biryukov Sorry. I didn't have time to post comments here.
The usecase that we have is a supportive translation unit for code completion. Probably you use something similar in clangd not to wait for the TU to be reparsed after a change?
The gain from this change is both performance and memory but I don't have measurements under my hand right now. And of course they are only valid for this usecase.

[ilya-biryukov]

I agree that the optimization is compelling, we do also share preamble in clangd and run code completion in parallel with other actions that need an AST.
However, I would argue that a proper way to introduce the optimization would be to change interface of ASTUnit that would somehow allow reusing the PCHs between different ASTUnits.

Having the cache is fine, but I suggest we:

1. Make it explicit in the interface, e.g. ASTUnit ctor will accept an extra PreambleCache parameter.
2. Have it off by default.

Whenever ASTUnit needs to rebuild a preamble, it can look into the cache and check there's a cached preamble and whether it matches the one that we're building.
After the rebuild is finished, it will try to push the built preamble into the cache.
ASTUnit can store the shared reference to a const preamble, i.e. std::shared_ptr<const PrecompiledPreamble>. The cache itself will store the weak_ptr<const PrecompiledPreamble>, thus allowing the preambles to be removed when they're not used anymore.

It would allow to:

1. avoid global state in ASTUnit, the cache will have to be created somewhere up the stack and passed to ASTUnit explicitly.
2. avoid complicating ASTUnit with synchronization code, it can be moved to preamble cache.
3. keep changes to the ASTUnit minimal. The only function that needs to be changed is the one building the preamble. Given how complicated ASTUnit is already, I think that's a big win.

include/clang/Frontend/ASTUnit.h

private:
 
   /// Map from FileID to the file-level declarations that it contains.
   /// The files and decls are only local (and non-preamble) ones.
   FileDeclsTy FileDecls;
   
   /// The name of the original source file used to generate this ASTUnit.
   std::string OriginalSourceFile;
 
-  /// The set of diagnostics produced when creating the preamble.
-  SmallVector<StandaloneDiagnostic, 4> PreambleDiagnostics;
-
   /// The set of diagnostics produced when creating this
   /// translation unit.
   SmallVector<StoredDiagnostic, 4> StoredDiagnostics;
 
   /// The set of diagnostics produced when failing to parse, e.g. due
   /// to failure to load the PCH.
   SmallVector<StoredDiagnostic, 4> FailedParseDiagnostics;
 
   /// The number of stored diagnostics that come from the driver
   /// itself.
   ///
   /// Diagnostics that come from the driver are retained from one parse to
   /// the next.
   unsigned NumStoredDiagnosticsFromDriver = 0;
-  
-  /// Counter that determines when we want to try building a
-  /// precompiled preamble.
-  ///
-  /// If zero, we will never build a precompiled preamble. Otherwise,
-  /// it's treated as a counter that decrements each time we reparse
-  /// without the benefit of a precompiled preamble. When it hits 1,
-  /// we'll attempt to rebuild the precompiled header. This way, if
-  /// building the precompiled preamble fails, we won't try again for
-  /// some number of calls.
-  unsigned PreambleRebuildCounter = 0;
-
-  /// Cache pairs "filename - source location"
-  ///
-  /// Cache contains only source locations from preamble so it is
-  /// guaranteed that they stay valid when the SourceManager is recreated.
-  /// This cache is used when loading preamble to increase performance
-  /// of that loading. It must be cleared when preamble is recreated.
-  llvm::StringMap<SourceLocation> PreambleSrcLocCache;
-
-  /// The contents of the preamble.
-  llvm::Optional<PrecompiledPreamble> Preamble;
-
   /// When non-NULL, this is the buffer used to store the contents of
   /// the main file when it has been padded for use with the precompiled
   /// preamble.
   std::unique_ptr<llvm::MemoryBuffer> SavedMainFileBuffer;
 
   /// The number of warnings that occurred while parsing the preamble.
   ///
   /// This value will be used to restore the state of the \c DiagnosticsEngine
   /// object when re-using the precompiled preamble. Note that only the
   /// number of warnings matters, since we will not save the preamble
   /// when any errors are present.
   unsigned NumWarningsInPreamble = 0;
 
-  /// A list of the serialization ID numbers for each of the top-level
-  /// declarations parsed within the precompiled preamble.
-  std::vector<serialization::DeclID> TopLevelDeclsInPreamble;
-  
   /// Whether we should be caching code-completion results.
   bool ShouldCacheCodeCompletionResults : 1;
 
   /// Whether to include brief documentation within the set of code
   /// completions cached.
   bool IncludeBriefCommentsInCodeCompletion : 1;
 
   /// True if non-system source files should be treated as volatile
struct CachedCodeCompletionResult {
     /// The type of a non-macro completion result, stored as a unique
     /// integer used by the string map of cached completion types.
     ///
     /// This value will be zero if the type is not known, or a unique value
     /// determined by the formatted type string. Se \c CachedCompletionTypes
     /// for more information.
     unsigned Type;
   };
-  
+  /// The cache to reuse the sharable data if there's more than 1 ASTUnit
+  /// for the same file.
+  class ASTUnitCache {
+    mutable unsigned RefCount = 0;
+
+  public:
+    ASTUnitCache() = default;
+    ASTUnitCache(const ASTUnitCache &) {}
+    void Retain() const { ++RefCount; }
+    void Release() const {
+      assert(RefCount > 0 && "Reference count is already zero.");
+      if (--RefCount == 0)
+        delete this;
+    }
+    unsigned getRefCount() const { return RefCount; }
+
+    /// The contents of the preamble.
+    llvm::Optional<PrecompiledPreamble> Preamble;
+
+    /// A list of the serialization ID numbers for each of the top-level
+    /// declarations parsed within the precompiled preamble.
+    std::vector<serialization::DeclID> TopLevelDeclsInPreamble;
+
+    /// The set of diagnostics produced when creating the preamble.
+    SmallVector<StandaloneDiagnostic, 4> PreambleDiagnostics;
+
+    /// Cache pairs "filename - source location"
+    ///
+    /// Cache contains only source locations from preamble so it is
+    /// guaranteed that they stay valid when the SourceManager is recreated.
+    /// This cache is used when loading preamble to increase performance
+    /// of that loading. It must be cleared when preamble is recreated.
+    llvm::StringMap<SourceLocation> PreambleSrcLocCache;
+
+    /// Counter that determines when we want to try building a
+    /// precompiled preamble.
+    ///
+    /// If zero, we will never build a precompiled preamble. Otherwise,
+    /// it's treated as a counter that decrements each time we reparse
+    /// without the benefit of a precompiled preamble. When it hits 1,
+    /// we'll attempt to rebuild the precompiled header. This way, if
+    /// building the precompiled preamble fails, we won't try again for
+    /// some number of calls.
+    unsigned PreambleRebuildCounter = 0;
+
+    /// A string hash of the top-level declaration and macro definition
+    /// names processed the last time that we reparsed the precompiled preamble.
+    ///
+    /// This hash value is used to determine when we need to refresh the
+    /// global code-completion cache after a rebuild of the precompiled preamble.
+    unsigned PreambleTopLevelHashValue = 0;
+  };
+
+  void initCache();
+  ASTUnitCache &getCache() { return ASTUnitCacheMap[getMainFileName()]; }
+  static std::recursive_mutex &getCacheMutex() { return CacheMutex; }
+
   /// Retrieve the mapping from formatted type names to unique type
   /// identifiers.
   llvm::StringMap<unsigned> &getCachedCompletionTypes() {
     return CachedCompletionTypes; 
   }
   
   /// Retrieve the allocator used to cache global code completions.
   std::shared_ptr<GlobalCodeCompletionAllocator>
   getCachedCompletionAllocator() {
     return CachedCompletionAllocator;
   }
 
   CodeCompletionTUInfo &getCodeCompletionTUInfo() {
     if (!CCTUInfo)
       CCTUInfo = llvm::make_unique<CodeCompletionTUInfo>(
           std::make_shared<GlobalCodeCompletionAllocator>());
     return *CCTUInfo;
   }
 
 private:
+  static std::recursive_mutex CacheMutex;
+  static llvm::StringMap<ASTUnitCache> ASTUnitCacheMap;
+
   /// Allocator used to store cached code completions.
   std::shared_ptr<GlobalCodeCompletionAllocator> CachedCompletionAllocator;
 
   std::unique_ptr<CodeCompletionTUInfo> CCTUInfo;
 
   /// The set of cached code-completion results.
   std::vector<CachedCodeCompletionResult> CachedCompletionResults;
-  
   /// A mapping from the formatted type name to a unique number for that
   /// type, which is used for type equality comparisons.
   llvm::StringMap<unsigned> CachedCompletionTypes;
   
   /// A string hash of the top-level declaration and macro definition 
   /// names processed the last time that we reparsed the file.
   ///
   /// This hash value is used to determine when we need to refresh the 
   /// global code-completion cache.
   unsigned CompletionCacheTopLevelHashValue = 0;
 
-  /// A string hash of the top-level declaration and macro definition 
-  /// names processed the last time that we reparsed the precompiled preamble.
-  ///
-  /// This hash value is used to determine when we need to refresh the 
-  /// global code-completion cache after a rebuild of the precompiled preamble.
-  unsigned PreambleTopLevelHashValue = 0;
-
   /// The current hash value for the top-level declaration and macro
   /// definition names
   unsigned CurrentTopLevelHashValue = 0;
   
   /// Bit used by CIndex to mark when a translation unit may be in an
   /// inconsistent state, and is not safe to free.
   unsigned UnsafeToFree : 1;
 
public:
   StringRef getMainFileName() const;
 
   /// If this ASTUnit came from an AST file, returns the filename for it.
   StringRef getASTFileName() const;
 
   using top_level_iterator = std::vector<Decl *>::iterator;
 
   top_level_iterator top_level_begin() {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
     assert(!isMainFileAST() && "Invalid call for AST based ASTUnit!");
-    if (!TopLevelDeclsInPreamble.empty())
+    if (!ASTUnitCacheMap[getMainFileName()].TopLevelDeclsInPreamble.empty())
       RealizeTopLevelDeclsFromPreamble();
     return TopLevelDecls.begin();
   }
 
   top_level_iterator top_level_end() {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
     assert(!isMainFileAST() && "Invalid call for AST based ASTUnit!");
-    if (!TopLevelDeclsInPreamble.empty())
+    if (!ASTUnitCacheMap[getMainFileName()].TopLevelDeclsInPreamble.empty())
       RealizeTopLevelDeclsFromPreamble();
     return TopLevelDecls.end();
   }
 
   std::size_t top_level_size() const {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
     assert(!isMainFileAST() && "Invalid call for AST based ASTUnit!");
-    return TopLevelDeclsInPreamble.size() + TopLevelDecls.size();
+    return ASTUnitCacheMap[getMainFileName()].TopLevelDeclsInPreamble.size() + TopLevelDecls.size();
   }
 
   bool top_level_empty() const {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
     assert(!isMainFileAST() && "Invalid call for AST based ASTUnit!");
-    return TopLevelDeclsInPreamble.empty() && TopLevelDecls.empty();
+    return ASTUnitCacheMap[getMainFileName()].TopLevelDeclsInPreamble.empty() && TopLevelDecls.empty();
   }
 
   /// Add a new top-level declaration.
   void addTopLevelDecl(Decl *D) {
     TopLevelDecls.push_back(D);
   }
 
   /// Add a new local file-level declaration.
cached_completion_iterator cached_completion_begin() {
     return CachedCompletionResults.begin();
   }
 
   cached_completion_iterator cached_completion_end() {
     return CachedCompletionResults.end();
   }
 
   unsigned cached_completion_size() const { 
-    return CachedCompletionResults.size(); 
+    return CachedCompletionResults.size();
   }
 
   /// Returns an iterator range for the local preprocessing entities
   /// of the local Preprocessor, if this is a parsed source file, or the loaded
   /// preprocessing entities of the primary module if this is an AST file.
   llvm::iterator_range<PreprocessingRecord::iterator>
   getLocalPreprocessingEntities() const;
 

lib/Frontend/ASTUnit.cpp

getBufferForFileHandlingRemapping(const CompilerInvocation &Invocation,
 
   if (BufferOwner)
     return BufferOwner;
   if (!Buffer)
     return nullptr;
   return llvm::MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(), FilePath);
 }
 
+std::recursive_mutex ASTUnit::CacheMutex;
+llvm::StringMap<ASTUnit::ASTUnitCache> ASTUnit::ASTUnitCacheMap;
+
 struct ASTUnit::ASTWriterData {
   SmallString<128> Buffer;
   llvm::BitstreamWriter Stream;
   ASTWriter Writer;
 
   ASTWriterData(MemoryBufferCache &PCMCache)
       : Stream(Buffer), Writer(Stream, Buffer, PCMCache, {}) {}
 };

 }
 
 ASTUnit::~ASTUnit() {
   // If we loaded from an AST file, balance out the BeginSourceFile call.
   if (MainFileIsAST && getDiagnostics().getClient()) {
     getDiagnostics().getClient()->EndSourceFile();
   }
 
+  std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+  StringRef MainFilePath = getMainFileName();
+  if (ASTUnitCacheMap.find(MainFilePath) == ASTUnitCacheMap.end())
+    return;
+
+  if (ASTUnitCacheMap[MainFilePath].getRefCount() == 1)
+    ASTUnitCacheMap.erase(MainFilePath);
+  else
+    ASTUnitCacheMap[MainFilePath].Release();
+
   clearFileLevelDecls();
 
   // Free the buffers associated with remapped files. We are required to
   // perform this operation here because we explicitly request that the
   // compiler instance *not* free these buffers for each invocation of the
   // parser.
   if (Invocation && OwnsRemappedFileBuffers) {
     PreprocessorOptions &PPOpts = Invocation->getPreprocessorOpts();
     for (const auto &RB : PPOpts.RemappedFileBuffers)
       delete RB.second;
   }
   
-  ClearCachedCompletionResults();
-
   if (getenv("LIBCLANG_OBJTRACKING"))
     fprintf(stderr, "--- %u translation units\n", --ActiveASTUnitObjects);
 }
 
+void ASTUnit::initCache() {
+  std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+  StringRef MainFilePath = getMainFileName();
+
+  ASTUnitCacheMap[MainFilePath].Retain();
+}
+
 void ASTUnit::setPreprocessor(std::shared_ptr<Preprocessor> PP) {
   this->PP = std::move(PP);
 }
 
 void ASTUnit::enableSourceFileDiagnostics() {
   assert(getDiagnostics().getClient() && Ctx &&
       "Bad context for source file");
   getDiagnostics().getClient()->BeginSourceFile(Ctx->getLangOpts(), PP.get());
if (ToLoad >= LoadEverything) {
     AST->TheSema.reset(new Sema(PP, *AST->Ctx, *AST->Consumer));
     AST->TheSema->Initialize();
     AST->Reader->InitializeSema(*AST->TheSema);
   }
 
   // Tell the diagnostic client that we have started a source file.
   AST->getDiagnostics().getClient()->BeginSourceFile(PP.getLangOpts(), &PP);
 
+  AST->initCache();
+
   return AST;
 }
 
 /// Add the given macro to the hash of all top-level entities.
 static void AddDefinedMacroToHash(const Token &MacroNameTok, unsigned &Hash) {
   Hash = llvm::djbHash(MacroNameTok.getIdentifierInfo()->getName(), Hash);
 }
 

 /// 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);
+  auto MainFilePath = getMainFileName();
   if (OverrideMainBuffer) {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+    auto &Preamble = ASTUnitCacheMap[MainFilePath].Preamble;
     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);
bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   }
 
   ResetForParse();
 
   SourceMgr = new SourceManager(getDiagnostics(), *FileMgr,
                                 UserFilesAreVolatile);
   if (!OverrideMainBuffer) {
     checkAndRemoveNonDriverDiags(StoredDiagnostics);
-    TopLevelDeclsInPreamble.clear();
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+    ASTUnitCacheMap[MainFilePath].TopLevelDeclsInPreamble.clear();
   }
 
   // Create a file manager object to provide access to and cache the filesystem.
   Clang->setFileManager(&getFileManager());
 
   // Create the source manager.
   Clang->setSourceManager(&getSourceManager());
 
bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
 
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<TopLevelDeclTrackerAction>
     ActCleanup(Act.get());
 
   if (!Act->BeginSourceFile(*Clang.get(), Clang->getFrontendOpts().Inputs[0]))
     goto error;
 
-  if (SavedMainFileBuffer)
-    TranslateStoredDiagnostics(getFileManager(), getSourceManager(),
-                               PreambleDiagnostics, StoredDiagnostics);
-  else
-    PreambleSrcLocCache.clear();
+  {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+    if (SavedMainFileBuffer) {
+      TranslateStoredDiagnostics(
+          getFileManager(), getSourceManager(),
+          ASTUnitCacheMap[MainFilePath].PreambleDiagnostics, StoredDiagnostics);
+    } else {
+      ASTUnitCacheMap[MainFilePath].PreambleSrcLocCache.clear();
+    }
+  }
 
   if (!Act->Execute())
     goto error;
 
   transferASTDataFromCompilerInstance(*Clang);
 
   Act->EndSourceFile();
 
if (!MainFileBuffer)
     return nullptr;
 
   PreambleBounds Bounds =
       ComputePreambleBounds(*PreambleInvocationIn.getLangOpts(),
                             MainFileBuffer.get(), MaxLines);
   if (!Bounds.Size)
     return nullptr;
 
+  {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+    auto &Preamble = ASTUnitCacheMap[MainFilePath].Preamble;
     if (Preamble) {
       if (Preamble->CanReuse(PreambleInvocationIn, MainFileBuffer.get(), Bounds,
                              VFS.get())) {
         // Okay! We can re-use the precompiled preamble.
 
         // Set the state of the diagnostic object to mimic its state
         // after parsing the preamble.
         getDiagnostics().Reset();
         ProcessWarningOptions(getDiagnostics(),
                               PreambleInvocationIn.getDiagnosticOpts());
         getDiagnostics().setNumWarnings(NumWarningsInPreamble);
 
-      PreambleRebuildCounter = 1;
+        ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter = 1;
         return MainFileBuffer;
       } else {
         Preamble.reset();
-      PreambleDiagnostics.clear();
-      TopLevelDeclsInPreamble.clear();
-      PreambleSrcLocCache.clear();
-      PreambleRebuildCounter = 1;
+        ASTUnitCacheMap[MainFilePath].PreambleDiagnostics.clear();
+        ASTUnitCacheMap[MainFilePath].TopLevelDeclsInPreamble.clear();
+        ASTUnitCacheMap[MainFilePath].PreambleSrcLocCache.clear();
+        ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter = 1;
       }
     }
 
     // If the preamble rebuild counter > 1, it's because we previously
     // failed to build a preamble and we're not yet ready to try
     // again. Decrement the counter and return a failure.
-  if (PreambleRebuildCounter > 1) {
-    --PreambleRebuildCounter;
+    if (ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter > 1) {
+      --ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter;
       return nullptr;
     }
+  }
 
-  assert(!Preamble && "No Preamble should be stored at that point");
   // If we aren't allowed to rebuild the precompiled preamble, just
   // return now.
   if (!AllowRebuild)
     return nullptr;
 
   SmallVector<StandaloneDiagnostic, 4> NewPreambleDiagsStandalone;
   SmallVector<StoredDiagnostic, 4> NewPreambleDiags;
   ASTUnitPreambleCallbacks Callbacks;
ASTUnitPreambleCallbacks Callbacks;
 
     llvm::ErrorOr<PrecompiledPreamble> NewPreamble = PrecompiledPreamble::Build(
         PreambleInvocationIn, MainFileBuffer.get(), Bounds, *Diagnostics, VFS,
         PCHContainerOps, /*StoreInMemory=*/false, Callbacks);
 
     PreambleInvocationIn.getFrontendOpts().SkipFunctionBodies =
         PreviousSkipFunctionBodies;
 
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
     if (NewPreamble) {
-      Preamble = std::move(*NewPreamble);
-      PreambleRebuildCounter = 1;
+      ASTUnitCacheMap[MainFilePath].Preamble = std::move(*NewPreamble);
+      ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter = 1;
+      assert(ASTUnitCacheMap[MainFilePath].Preamble && "Preamble wasn't built");
     } else {
       switch (static_cast<BuildPreambleError>(NewPreamble.getError().value())) {
       case BuildPreambleError::CouldntCreateTempFile:
       case BuildPreambleError::PreambleIsEmpty:
         // Try again next time.
-        PreambleRebuildCounter = 1;
+        ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter = 1;
         return nullptr;
       case BuildPreambleError::CouldntCreateTargetInfo:
       case BuildPreambleError::BeginSourceFileFailed:
       case BuildPreambleError::CouldntEmitPCH:
         // These erros are more likely to repeat, retry after some period.
-        PreambleRebuildCounter = DefaultPreambleRebuildInterval;
+        ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter = DefaultPreambleRebuildInterval;
         return nullptr;
       }
       llvm_unreachable("unexpected BuildPreambleError");
     }
   }
 
-  assert(Preamble && "Preamble wasn't built");
-
   TopLevelDecls.clear();
-  TopLevelDeclsInPreamble = Callbacks.takeTopLevelDeclIDs();
-  PreambleTopLevelHashValue = Callbacks.getHash();
+
+  unsigned NewPreambleTopLevelHashValue = Callbacks.getHash();
 
   NumWarningsInPreamble = getDiagnostics().getNumWarnings();
 
   checkAndRemoveNonDriverDiags(NewPreambleDiags);
   StoredDiagnostics = std::move(NewPreambleDiags);
-  PreambleDiagnostics = std::move(NewPreambleDiagsStandalone);
 
+  std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+  ASTUnitCacheMap[MainFilePath].TopLevelDeclsInPreamble = Callbacks.takeTopLevelDeclIDs();
+  ASTUnitCacheMap[MainFilePath].PreambleDiagnostics = std::move(NewPreambleDiagsStandalone);
+
+  auto &PreambleTopLevelHashValue = ASTUnitCacheMap[MainFilePath].PreambleTopLevelHashValue;
   // If the hash of top-level entities differs from the hash of the top-level
   // entities the last time we rebuilt the preamble, clear out the completion
   // cache.
-  if (CurrentTopLevelHashValue != PreambleTopLevelHashValue) {
+  if (NewPreambleTopLevelHashValue != PreambleTopLevelHashValue) {
     CompletionCacheTopLevelHashValue = 0;
-    PreambleTopLevelHashValue = CurrentTopLevelHashValue;
+    PreambleTopLevelHashValue = NewPreambleTopLevelHashValue;
   }
 
   return MainFileBuffer;
 }
 
 void ASTUnit::RealizeTopLevelDeclsFromPreamble() {
-  assert(Preamble && "Should only be called when preamble was built");
+  auto MainFilePath = getMainFileName();
+  std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+  assert(ASTUnitCacheMap[MainFilePath].Preamble && "Should only be called when preamble was built");
 
   std::vector<Decl *> Resolved;
+  auto &TopLevelDeclsInPreamble = ASTUnitCacheMap[MainFilePath].TopLevelDeclsInPreamble;
   Resolved.reserve(TopLevelDeclsInPreamble.size());
   ExternalASTSource &Source = *getASTContext().getExternalSource();
   for (const auto TopLevelDecl : TopLevelDeclsInPreamble) {
     // Resolve the declaration ID to an actual declaration, possibly
     // deserializing the declaration in the process.
     if (Decl *D = Source.GetExternalDecl(TopLevelDecl))
       Resolved.push_back(D);
   }
ASTUnit::create(std::shared_ptr<CompilerInvocation> CI,
   AST->FileSystemOpts = CI->getFileSystemOpts();
   AST->Invocation = std::move(CI);
   AST->FileMgr = new FileManager(AST->FileSystemOpts, VFS);
   AST->UserFilesAreVolatile = UserFilesAreVolatile;
   AST->SourceMgr = new SourceManager(AST->getDiagnostics(), *AST->FileMgr,
                                      UserFilesAreVolatile);
   AST->PCMCache = new MemoryBufferCache;
 
+  AST->initCache();
+
   return AST;
 }
 
 ASTUnit *ASTUnit::LoadFromCompilerInvocationAction(
     std::shared_ptr<CompilerInvocation> CI,
     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
     IntrusiveRefCntPtr<DiagnosticsEngine> Diags, FrontendAction *Action,
     ASTUnit *Unit, bool Persistent, StringRef ResourceFilesPath,
ASTUnit *ASTUnit::LoadFromCompilerInvocationAction(
   }
 
   if (!ResourceFilesPath.empty()) {
     // Override the resources path.
     CI->getHeaderSearchOpts().ResourceDir = ResourceFilesPath;
   }
   AST->OnlyLocalDecls = OnlyLocalDecls;
   AST->CaptureDiagnostics = CaptureDiagnostics;
-  if (PrecompilePreambleAfterNParses > 0)
-    AST->PreambleRebuildCounter = PrecompilePreambleAfterNParses;
+  if (PrecompilePreambleAfterNParses > 0) {
+    std::lock_guard<std::recursive_mutex> Lock(AST->getCacheMutex());
+    AST->getCache().PreambleRebuildCounter = PrecompilePreambleAfterNParses;
+  }
   AST->TUKind = Action ? Action->getTranslationUnitKind() : TU_Complete;
   AST->ShouldCacheCodeCompletionResults = CacheCodeCompletionResults;
   AST->IncludeBriefCommentsInCodeCompletion
     = IncludeBriefCommentsInCodeCompletion;
 
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<ASTUnit>
     ASTUnitCleanup(OwnAST.get());
bool ASTUnit::LoadFromCompilerInvocation(
     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
     unsigned PrecompilePreambleAfterNParses,
     IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   if (!Invocation)
     return true;
 
   assert(VFS && "VFS is null");
 
+  initCache();
+
   // We'll manage file buffers ourselves.
   Invocation->getPreprocessorOpts().RetainRemappedFileBuffers = true;
   Invocation->getFrontendOpts().DisableFree = false;
   getDiagnostics().Reset();
   ProcessWarningOptions(getDiagnostics(), Invocation->getDiagnosticOpts());
 
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
   if (PrecompilePreambleAfterNParses > 0) {
-    PreambleRebuildCounter = PrecompilePreambleAfterNParses;
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+    ASTUnitCacheMap[getMainFileName()].PreambleRebuildCounter =
+        PrecompilePreambleAfterNParses;
     OverrideMainBuffer =
         getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation, VFS);
     getDiagnostics().Reset();
     ProcessWarningOptions(getDiagnostics(), Invocation->getDiagnosticOpts());
   }
 
   SimpleTimer ParsingTimer(WantTiming);
   ParsingTimer.setOutput("Parsing " + getMainFileName());
bool ASTUnit::Reparse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   for (const auto &RemappedFile : RemappedFiles) {
     Invocation->getPreprocessorOpts().addRemappedFile(RemappedFile.first,
                                                       RemappedFile.second);
   }
 
   // If we have a preamble file lying around, or if we might try to
   // build a precompiled preamble, do so now.
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
-  if (Preamble || PreambleRebuildCounter > 0)
-    OverrideMainBuffer =
-        getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation, VFS);
+  StringRef MainFilePath = getMainFileName();
+  {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+    if (ASTUnitCacheMap[MainFilePath].Preamble ||
+        ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter > 0) {
+      OverrideMainBuffer = getMainBufferWithPrecompiledPreamble(
+          PCHContainerOps, *Invocation, VFS);
+    }
+  }
 
   // Clear out the diagnostics state.
   FileMgr.reset();
   getDiagnostics().Reset();
   ProcessWarningOptions(getDiagnostics(), Invocation->getDiagnosticOpts());
   if (OverrideMainBuffer)
     getDiagnostics().setNumWarnings(NumWarningsInPreamble);
 
void ASTUnit::CodeComplete(
   CompletionTimer.setOutput("Code completion @ " + File + ":" +
                             Twine(Line) + ":" + Twine(Column));
 
   auto CCInvocation = std::make_shared<CompilerInvocation>(*Invocation);
 
   FrontendOptions &FrontendOpts = CCInvocation->getFrontendOpts();
   CodeCompleteOptions &CodeCompleteOpts = FrontendOpts.CodeCompleteOpts;
   PreprocessorOptions &PreprocessorOpts = CCInvocation->getPreprocessorOpts();
+  StringRef MainFilePath = getMainFileName();
 
   CodeCompleteOpts.IncludeMacros = IncludeMacros &&
                                    CachedCompletionResults.empty();
-  CodeCompleteOpts.IncludeCodePatterns = IncludeCodePatterns;
   CodeCompleteOpts.IncludeGlobals = CachedCompletionResults.empty();
+  CodeCompleteOpts.IncludeCodePatterns = IncludeCodePatterns;
   CodeCompleteOpts.IncludeBriefComments = IncludeBriefComments;
   CodeCompleteOpts.LoadExternal = Consumer.loadExternal();
   CodeCompleteOpts.IncludeFixIts = Consumer.includeFixIts();
 
   assert(IncludeBriefComments == this->IncludeBriefCommentsInCodeCompletion);
 
   FrontendOpts.CodeCompletionAt.FileName = File;
   FrontendOpts.CodeCompletionAt.Line = Line;
AugmentedCodeCompleteConsumer *AugmentedConsumer
     = new AugmentedCodeCompleteConsumer(*this, Consumer, CodeCompleteOpts);
   Clang->setCodeCompletionConsumer(AugmentedConsumer);
 
   // If we have a precompiled preamble, try to use it. We only allow
   // the use of the precompiled preamble if we're if the completion
   // point is within the main file, after the end of the precompiled
   // preamble.
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
-  if (Preamble) {
+  {
+    std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+    if (ASTUnitCacheMap[MainFilePath].Preamble) {
       std::string CompleteFilePath(File);
 
       auto VFS = FileMgr.getVirtualFileSystem();
       auto CompleteFileStatus = VFS->status(CompleteFilePath);
       if (CompleteFileStatus) {
-      llvm::sys::fs::UniqueID CompleteFileID = CompleteFileStatus->getUniqueID();
+        llvm::sys::fs::UniqueID CompleteFileID =
+            CompleteFileStatus->getUniqueID();
 
         std::string MainPath(OriginalSourceFile);
         auto MainStatus = VFS->status(MainPath);
         if (MainStatus) {
           llvm::sys::fs::UniqueID MainID = MainStatus->getUniqueID();
           if (CompleteFileID == MainID && Line > 1)
             OverrideMainBuffer = getMainBufferWithPrecompiledPreamble(
                 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 &&
+      assert(ASTUnitCacheMap[MainFilePath].Preamble &&
              "No preamble was built, but OverrideMainBuffer is not null");
 
       auto VFS = FileMgr.getVirtualFileSystem();
-    Preamble->AddImplicitPreamble(Clang->getInvocation(), VFS,
-                                  OverrideMainBuffer.get());
+      ASTUnitCacheMap[MainFilePath].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.
+      // 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)
     PreprocessorOpts.DetailedRecord = false;
 
   std::unique_ptr<SyntaxOnlyAction> Act;
   Act.reset(new SyntaxOnlyAction);
   if (Act->BeginSourceFile(*Clang.get(), Clang->getFrontendOpts().Inputs[0])) {
void ASTUnit::TranslateStoredDiagnostics(
   for (const auto &SD : Diags) {
     // Rebuild the StoredDiagnostic.
     if (SD.Filename.empty())
       continue;
     const FileEntry *FE = FileMgr.getFile(SD.Filename);
     if (!FE)
       continue;
     SourceLocation FileLoc;
+    {
+      std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+      llvm::StringMap<SourceLocation> &PreambleSrcLocCache =
+          ASTUnitCacheMap[getMainFileName()].PreambleSrcLocCache;
       auto ItFileID = PreambleSrcLocCache.find(SD.Filename);
       if (ItFileID == PreambleSrcLocCache.end()) {
         FileID FID = SrcMgr.translateFile(FE);
         FileLoc = SrcMgr.getLocForStartOfFile(FID);
         PreambleSrcLocCache[SD.Filename] = FileLoc;
       } else {
         FileLoc = ItFileID->getValue();
       }
+    }
 
     if (FileLoc.isInvalid())
       continue;
     SourceLocation L = FileLoc.getLocWithOffset(SD.LocOffset);
     FullSourceLoc Loc(L, SrcMgr);
 
     SmallVector<CharSourceRange, 4> Ranges;
     Ranges.reserve(SD.Ranges.size());

 /// If \arg Loc is a loaded location from the preamble, returns
 /// the corresponding local location of the main file, otherwise it returns
 /// \arg Loc.
 SourceLocation ASTUnit::mapLocationFromPreamble(SourceLocation Loc) const {
   FileID PreambleID;
   if (SourceMgr)
     PreambleID = SourceMgr->getPreambleFileID();
 
+  std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+  auto &Preamble = ASTUnitCacheMap[getMainFileName()].Preamble;
   if (Loc.isInvalid() || !Preamble || PreambleID.isInvalid())
     return Loc;
 
   unsigned Offs;
   if (SourceMgr->isInFileID(Loc, PreambleID, &Offs) && Offs < Preamble->getBounds().Size) {
     SourceLocation FileLoc
         = SourceMgr->getLocForStartOfFile(SourceMgr->getMainFileID());
     return FileLoc.getLocWithOffset(Offs);
   }
 
   return Loc;
 }
 
 /// If \arg Loc is a local location of the main file but inside the
 /// preamble chunk, returns the corresponding loaded location from the
 /// preamble, otherwise it returns \arg Loc.
 SourceLocation ASTUnit::mapLocationToPreamble(SourceLocation Loc) const {
   FileID PreambleID;
   if (SourceMgr)
     PreambleID = SourceMgr->getPreambleFileID();
 
+  std::lock_guard<std::recursive_mutex> Lock(CacheMutex);
+  auto &Preamble = ASTUnitCacheMap[getMainFileName()].Preamble;
   if (Loc.isInvalid() || !Preamble || PreambleID.isInvalid())
     return Loc;
 
   unsigned Offs;
   if (SourceMgr->isInFileID(Loc, SourceMgr->getMainFileID(), &Offs) &&
       Offs < Preamble->getBounds().Size) {
     SourceLocation FileLoc = SourceMgr->getLocForStartOfFile(PreambleID);
     return FileLoc.getLocWithOffset(Offs);