[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.

Qt Creator creates two translation units for each file.
With this patch they try to share the same preamble which
in most of the cases prevents double generation by checking
the bounds of the existing preamble and reusing it if it's valid.
In the worst case preamble is still generated twice and we get
slightly worse performance than before but I never saw such case.
This is a cherry pick from commit 7e438f61d9caa617562cd650bf6578be64be94ac

Change-Id: I9c1a5fe8c433cec9a5d75d1bded97064346b946c
Reviewed-by: Nikolai Kosjar <nikolai.kosjar@qt.io>

2 files changed, 231 insertions, 133 deletions

diff --git a/include/clang/Frontend/ASTUnit.h b/include/clang/Frontend/ASTUnit.h
index ba14596546..412a9778bd 100644
--- a/include/clang/Frontend/ASTUnit.h
+++ b/include/clang/Frontend/ASTUnit.h
@@ -180,9 +180,6 @@ private:
   /// 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;
@@ -198,28 +195,6 @@ private:
   /// 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.
@@ -233,10 +208,6 @@ private:
   /// 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;
 
@@ -297,6 +268,63 @@ public:
     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() {
@@ -317,6 +345,9 @@ public:
   }
 
 private:
+  static std::recursive_mutex             CacheMutex;
+  static llvm::StringMap<ASTUnitCache>    ASTUnitCacheMap;
+
   /// Allocator used to store cached code completions.
   std::shared_ptr<GlobalCodeCompletionAllocator> CachedCompletionAllocator;
 
@@ -336,13 +367,6 @@ private:
   /// 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;
@@ -496,27 +520,31 @@ public:
   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.
diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index e4c313fed3..98c428d95d 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -214,6 +214,9 @@ getBufferForFileHandlingRemapping(const CompilerInvocation &Invocation,
   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;
@@ -253,6 +256,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
@@ -271,6 +284,13 @@ ASTUnit::~ASTUnit() {
     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);
 }
@@ -841,6 +861,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;
 }
 
@@ -1079,7 +1101,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;
@@ -1147,7 +1172,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.
@@ -1180,11 +1206,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;
@@ -1295,38 +1326,42 @@ 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();
-      PreambleSrcLocCache.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].PreambleSrcLocCache.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)
@@ -1357,54 +1392,61 @@ 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:
         // 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) {
@@ -1480,6 +1522,8 @@ ASTUnit::create(std::shared_ptr<CompilerInvocation> CI,
                                      UserFilesAreVolatile);
   AST->PCMCache = new MemoryBufferCache;
 
+  AST->initCache();
+
   return AST;
 }
 
@@ -1510,8 +1554,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
@@ -1637,6 +1683,8 @@ bool ASTUnit::LoadFromCompilerInvocation(
 
   assert(VFS && "VFS is null");
 
+  initCache();
+
   // We'll manage file buffers ourselves.
   Invocation->getPreprocessorOpts().RetainRemappedFileBuffers = true;
   Invocation->getFrontendOpts().DisableFree = false;
@@ -1645,7 +1693,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();
@@ -1823,9 +1873,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();
@@ -2122,11 +2178,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.LoadExternal = Consumer.loadExternal();
   CodeCompleteOpts.IncludeFixIts = Consumer.includeFixIts();
@@ -2208,42 +2265,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.
@@ -2341,13 +2402,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())
@@ -2486,6 +2552,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;
 
@@ -2507,6 +2575,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;