[Frontend] Cache preamble-related data

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

Change-Id: I9c1a5fe8c433cec9a5d75d1bded97064346b946c
Reviewed-by: Erik Verbruggen <erik.verbruggen@qt.io>

2 files changed, 237 insertions, 145 deletions

diff --git a/include/clang/Frontend/ASTUnit.h b/include/clang/Frontend/ASTUnit.h
index 09aeab5b6f..896cb31680 100644
--- a/include/clang/Frontend/ASTUnit.h
+++ b/include/clang/Frontend/ASTUnit.h
@@ -166,9 +166,6 @@ private:
   /// The name of the original source file used to generate this ASTUnit.
   std::string OriginalSourceFile;
 
-  /// \brief The set of diagnostics produced when creating the preamble.
-  SmallVector<StandaloneDiagnostic, 4> PreambleDiagnostics;
-
   /// \brief The set of diagnostics produced when creating this
   /// translation unit.
   SmallVector<StoredDiagnostic, 4> StoredDiagnostics;
@@ -183,30 +180,8 @@ private:
   /// Diagnostics that come from the driver are retained from one parse to
   /// the next.
   unsigned NumStoredDiagnosticsFromDriver;
-  
-  /// \brief 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;
-
-  /// \brief 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 preambule to increase performance
-  /// of that loading. It must be cleared when preamble is recreated.
-  llvm::StringMap<SourceLocation> PreambleSrcLocCache;
 
 private:
-  /// The contents of the preamble.
-  llvm::Optional<PrecompiledPreamble> Preamble;
-
   /// \brief 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.
@@ -220,10 +195,6 @@ private:
   /// when any errors are present.
   unsigned NumWarningsInPreamble;
 
-  /// \brief A list of the serialization ID numbers for each of the top-level
-  /// declarations parsed within the precompiled preamble.
-  std::vector<serialization::DeclID> TopLevelDeclsInPreamble;
-  
   /// \brief Whether we should be caching code-completion results.
   bool ShouldCacheCodeCompletionResults : 1;
 
@@ -283,7 +254,64 @@ public:
     /// for more information.
     unsigned Type;
   };
-  
+
+  /// \brief 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; }
+
+    /// \brief The contents of the preamble.
+    llvm::Optional<PrecompiledPreamble> Preamble;
+
+    /// \brief A list of the serialization ID numbers for each of the top-level
+    /// declarations parsed within the precompiled preamble.
+    std::vector<serialization::DeclID> TopLevelDeclsInPreamble;
+
+    /// \brief The set of diagnostics produced when creating the preamble.
+    SmallVector<StandaloneDiagnostic, 4> PreambleDiagnostics;
+
+    /// \brief 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 preambule to increase performance
+    /// of that loading. It must be cleared when preamble is recreated.
+    llvm::StringMap<SourceLocation> PreambleSrcLocCache;
+
+    /// \brief 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;
+
+    /// \brief 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; }
+
   /// \brief Retrieve the mapping from formatted type names to unique type
   /// identifiers.
   llvm::StringMap<unsigned> &getCachedCompletionTypes() { 
@@ -304,6 +332,9 @@ public:
   }
 
 private:
+  static std::recursive_mutex             CacheMutex;
+  static llvm::StringMap<ASTUnitCache>    ASTUnitCacheMap;
+
   /// \brief Allocator used to store cached code completions.
   std::shared_ptr<GlobalCodeCompletionAllocator> CachedCompletionAllocator;
 
@@ -311,24 +342,17 @@ private:
 
   /// \brief The set of cached code-completion results.
   std::vector<CachedCodeCompletionResult> CachedCompletionResults;
-  
+
   /// \brief 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;
-  
-  /// \brief A string hash of the top-level declaration and macro definition 
+
+  /// \brief 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 
+  /// This hash value is used to determine when we need to refresh the
   /// global code-completion cache.
-  unsigned CompletionCacheTopLevelHashValue;
-
-  /// \brief 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;
+  unsigned CompletionCacheTopLevelHashValue = 0;
 
   /// \brief The current hash value for the top-level declaration and macro
   /// definition names
@@ -474,27 +498,31 @@ public:
   typedef std::vector<Decl *>::iterator top_level_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();
   }
 
   /// \brief Add a new top-level declaration.
@@ -589,7 +617,7 @@ public:
   }
 
   unsigned cached_completion_size() const { 
-    return CachedCompletionResults.size(); 
+    return CachedCompletionResults.size();
   }
 
   /// \brief Returns an iterator range for the local preprocessing entities
diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index 1ec4502ba5..b1a6583714 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -161,6 +161,9 @@ getBufferForFileHandlingRemapping(const CompilerInvocation &Invocation,
 }
 }
 
+std::recursive_mutex ASTUnit::CacheMutex;
+llvm::StringMap<ASTUnit::ASTUnitCache> ASTUnit::ASTUnitCacheMap;
+
 struct ASTUnit::ASTWriterData {
   SmallString<128> Buffer;
   llvm::BitstreamWriter Stream;
@@ -192,14 +195,11 @@ ASTUnit::ASTUnit(bool _MainFileIsAST)
     TUKind(TU_Complete), WantTiming(getenv("LIBCLANG_TIMING")),
     OwnsRemappedFileBuffers(true),
     NumStoredDiagnosticsFromDriver(0),
-    PreambleRebuildCounter(0),
     NumWarningsInPreamble(0),
     ShouldCacheCodeCompletionResults(false),
     IncludeBriefCommentsInCodeCompletion(false), UserFilesAreVolatile(false),
-    CompletionCacheTopLevelHashValue(0),
-    PreambleTopLevelHashValue(0),
     CurrentTopLevelHashValue(0),
-    UnsafeToFree(false) { 
+    UnsafeToFree(false) {
   if (getenv("LIBCLANG_OBJTRACKING"))
     fprintf(stderr, "+++ %u translation units\n", ++ActiveASTUnitObjects);
 }
@@ -210,6 +210,16 @@ ASTUnit::~ASTUnit() {
     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
@@ -221,13 +231,18 @@ ASTUnit::~ASTUnit() {
     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);
 }
@@ -794,6 +809,8 @@ std::unique_ptr<ASTUnit> ASTUnit::LoadFromASTFile(
   // Tell the diagnostic client that we have started a source file.
   AST->getDiagnostics().getClient()->BeginSourceFile(PP.getLangOpts(), &PP);
 
+  AST->initCache();
+
   return AST;
 }
 
@@ -1027,7 +1044,10 @@ bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
     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;
@@ -1095,7 +1115,8 @@ bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
                                 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.
@@ -1128,11 +1149,16 @@ bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   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;
@@ -1244,37 +1270,41 @@ ASTUnit::getMainBufferWithPrecompiledPreamble(
   if (!Bounds.Size)
     return nullptr;
 
-  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;
-      return MainFileBuffer;
-    } else {
-      Preamble.reset();
-      PreambleDiagnostics.clear();
-      TopLevelDeclsInPreamble.clear();
-      PreambleRebuildCounter = 1;
+  {
+    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);
+
+        ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter = 1;
+        return MainFileBuffer;
+      } else {
+        Preamble.reset();
+        ASTUnitCacheMap[MainFilePath].PreambleDiagnostics.clear();
+        ASTUnitCacheMap[MainFilePath].TopLevelDeclsInPreamble.clear();
+        ASTUnitCacheMap[MainFilePath].PreambleRebuildCounter = 1;
+      }
     }
-  }
+    assert(!Preamble && "No Preamble should be stored at that point");
 
-  // 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;
-    return nullptr;
+    // 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 (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)
@@ -1305,55 +1335,62 @@ ASTUnit::getMainBufferWithPrecompiledPreamble(
     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:
       case BuildPreambleError::CouldntCreateVFSOverlay:
         // 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 (serialization::DeclID TopLevelDecl : TopLevelDeclsInPreamble) {
@@ -1431,6 +1468,8 @@ ASTUnit::create(std::shared_ptr<CompilerInvocation> CI,
                                      UserFilesAreVolatile);
   AST->PCMCache = new MemoryBufferCache;
 
+  AST->initCache();
+
   return AST;
 }
 
@@ -1461,8 +1500,10 @@ ASTUnit *ASTUnit::LoadFromCompilerInvocationAction(
   }
   AST->OnlyLocalDecls = OnlyLocalDecls;
   AST->CaptureDiagnostics = CaptureDiagnostics;
-  if (PrecompilePreambleAfterNParses > 0)
-    AST->PreambleRebuildCounter = PrecompilePreambleAfterNParses;
+  if (PrecompilePreambleAfterNParses > 0) {
+    std::lock_guard<std::recursive_mutex> Lock(ASTUnit::getCacheMutex());
+    AST->getCache().PreambleRebuildCounter = PrecompilePreambleAfterNParses;
+  }
   AST->TUKind = Action ? Action->getTranslationUnitKind() : TU_Complete;
   AST->ShouldCacheCodeCompletionResults = CacheCodeCompletionResults;
   AST->IncludeBriefCommentsInCodeCompletion
@@ -1588,6 +1629,8 @@ bool ASTUnit::LoadFromCompilerInvocation(
 
   assert(VFS && "VFS is null");
 
+  initCache();
+
   // We'll manage file buffers ourselves.
   Invocation->getPreprocessorOpts().RetainRemappedFileBuffers = true;
   Invocation->getFrontendOpts().DisableFree = false;
@@ -1596,7 +1639,9 @@ bool ASTUnit::LoadFromCompilerInvocation(
 
   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();
@@ -1777,10 +1822,15 @@ bool ASTUnit::Reparse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   // 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();
@@ -2076,11 +2126,12 @@ void ASTUnit::CodeComplete(
   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.IncludeFixIts = Consumer.includeFixIts();
 
@@ -2161,42 +2212,46 @@ void ASTUnit::CodeComplete(
   // point is within the main file, after the end of the precompiled
   // preamble.
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
-  if (Preamble) {
-    std::string CompleteFilePath(File);
-
-    auto VFS = FileMgr.getVirtualFileSystem();
-    auto CompleteFileStatus = VFS->status(CompleteFilePath);
-    if (CompleteFileStatus) {
-      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);
+  {
+    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();
+
+        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 &&
-           "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;
+    // 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(ASTUnitCacheMap[MainFilePath].Preamble &&
+             "No preamble was built, but OverrideMainBuffer is not null");
+
+      auto VFS = FileMgr.getVirtualFileSystem();
+      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.
+      OwnedBuffers.push_back(OverrideMainBuffer.release());
+    } else {
+      PreprocessorOpts.PrecompiledPreambleBytes.first = 0;
+      PreprocessorOpts.PrecompiledPreambleBytes.second = false;
+    }
   }
 
   // Disable the preprocessing record if modules are not enabled.
@@ -2294,13 +2349,18 @@ void ASTUnit::TranslateStoredDiagnostics(
     if (!FE)
       continue;
     SourceLocation FileLoc;
-    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();
+    {
+      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())
@@ -2439,6 +2499,8 @@ SourceLocation ASTUnit::mapLocationFromPreamble(SourceLocation Loc) const {
   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;
 
@@ -2460,6 +2522,8 @@ SourceLocation ASTUnit::mapLocationToPreamble(SourceLocation Loc) const {
   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;