[Caching] Don't rely on FileSystem when replaying diagnostics

Stop relying on file system to provide source buffer for diagnostics
when replying. This avoids initializing CASFS which is quite expensive.
Now cached diagnostics contains CASID for the file buffer so it can
initialize its own source manager without relying on the underlying file
system.

rdar://128423393

Author: cachemeifyoucan

lib/Frontend/CachedDiagnostics.cpp

@@ -24,11 +24,14 @@
 #include "swift/Frontend/FrontendInputsAndOutputs.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/TinyPtrVector.h"
+#include "llvm/CAS/ObjectStore.h"
 #include "llvm/Support/Compression.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/PrefixMapper.h"
 #include "llvm/Support/SMLoc.h"
 #include "llvm/Support/VirtualFileSystem.h"
@@ -82,7 +85,8 @@ struct SerializedDiagnosticInfo {
 struct SerializedFile {
   std::string FileName;
   SerializedSourceLoc IncludeLoc = SerializedSourceLoc();
-  StringRef Content;
+  bool IsFileBacked;
+  std::string Content;
 };
 
 struct SerializedVirtualFile {
@@ -100,8 +104,8 @@ struct SerializedGeneratedFileInfo {
 
 struct DiagnosticSerializer {
   DiagnosticSerializer(llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS,
-                       llvm::PrefixMapper &Mapper)
-      : SrcMgr(FS), Mapper(Mapper) {}
+                       llvm::PrefixMapper &Mapper, llvm::cas::ObjectStore &CAS)
+      : SrcMgr(FS), Mapper(Mapper), CAS(CAS) {}
 
   using ReplayFunc = llvm::function_ref<llvm::Error(const DiagnosticInfo &)>;
 
@@ -111,31 +115,36 @@ struct DiagnosticSerializer {
   llvm::Error serializeEmittedDiagnostics(llvm::raw_ostream &os);
 
   static llvm::Error
-  emitDiagnosticsFromCached(llvm::StringRef Buffer, SourceManager &SrcMgr,
-                            DiagnosticEngine &Diags,
-                            llvm::PrefixMapper &Mapper,
+  emitDiagnosticsFromCached(llvm::StringRef Buffer,
+                            DiagnosticEngine &Diags, llvm::PrefixMapper &Mapper,
+                            llvm::cas::ObjectStore &CAS,
                             const FrontendInputsAndOutputs &InAndOut) {
     // Create a new DiagnosticSerializer since this cannot be shared with a
-    // serialization instance.
-    DiagnosticSerializer DS(SrcMgr.getFileSystem(), Mapper);
-    DS.addInputsToSourceMgr(InAndOut);
+    // serialization instance. Using an empty in-memory file system as
+    // underlying file system because the replay logic should not touch file
+    // system.
+    auto FS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+    DiagnosticSerializer DS(std::move(FS), Mapper, CAS);
     return DS.doEmitFromCached(Buffer, Diags);
   }
 
   SourceManager &getSourceMgr() { return SrcMgr; }
 
-  void addInputsToSourceMgr(const FrontendInputsAndOutputs &InAndOut) {
+  void addInputsToSourceMgr(SourceManager &SM,
+                            const FrontendInputsAndOutputs &InAndOut) {
     // Extract all the input file names so they can be added to the source
     // manager when replaying the diagnostics. All input files are needed even
     // they don't contain diagnostics because FileSpecificDiagConsumer need
     // has references to input files to find subconsumer.
     auto addInputToSourceMgr = [&](const InputFile &Input) {
-      if (Input.getFileName() != "-")
-        SrcMgr.getExternalSourceBufferID(remapFilePath(Input.getFileName()));
+      auto Path = remapFilePath(Input.getFileName());
+      SrcMgr.getExternalSourceBufferID(Path);
+      auto Idx = SM.getExternalSourceBufferID(Path);
+      if (Idx != 0)
+        getFileIDFromBufferID(SM, Idx);
       return false;
     };
-    InAndOut.forEachInputProducingSupplementaryOutput(addInputToSourceMgr);
-    InAndOut.forEachNonPrimaryInput(addInputToSourceMgr);
+    InAndOut.forEachInput(addInputToSourceMgr);
   }
 
 private:
@@ -169,7 +178,7 @@ struct DiagnosticSerializer {
                                         DiagnosticStorage &, ReplayFunc);
 
   // Deserialize File and return the bufferID in serializing SourceManager.
-  unsigned deserializeFile(const SerializedFile &File);
+  llvm::Expected<unsigned> deserializeFile(const SerializedFile &File);
   llvm::Error deserializeVirtualFile(const SerializedVirtualFile &VF);
   llvm::Error deserializeGeneratedFileInfo(const SerializedGeneratedFileInfo &Info);
   std::string remapFilePath(StringRef Path) {
@@ -190,6 +199,9 @@ struct DiagnosticSerializer {
   SourceManager SrcMgr;
   llvm::PrefixMapper &Mapper;
 
+  // CAS for file system backing.
+  llvm::cas::ObjectStore &CAS;
+
   // Mapping of the FileID between SourceManager from CompilerInstance vs.
   // the serialized FileID in cached diagnostics. Lookup tables are
   // per-SourceManager to handle diagnostics from all sub-instances which
@@ -247,7 +259,8 @@ struct MappingTraits<SerializedFile> {
   static void mapping(IO &io, SerializedFile &F) {
     io.mapRequired("Name", F.FileName);
     io.mapOptional("IncludeLoc", F.IncludeLoc, SerializedSourceLoc());
-    io.mapOptional("Content", F.Content, StringRef());
+    io.mapOptional("FileBacked", F.IsFileBacked, false);
+    io.mapRequired("Content", F.Content);
   }
 };
 
@@ -302,25 +315,28 @@ void DiagnosticSerializer::handleDiagnostic(SourceManager &SM,
 
 unsigned DiagnosticSerializer::getFileIDFromBufferID(SourceManager &SM,
                                                      unsigned Idx) {
-  auto &Buf = SM.getLLVMSourceMgr().getBufferInfo(Idx);
-  auto Filename = Buf.Buffer->getBufferIdentifier();
-  bool IsFSBacked = SM.getFileSystem()->exists(Filename);
-
   // See if the file is already constructed.
   auto &Allocated = FileMapper[&SM];
   auto ID = Allocated.find(Idx);
   if (ID != Allocated.end())
     return ID->second;
 
+  auto &Buf = SM.getLLVMSourceMgr().getBufferInfo(Idx);
+  auto Filename = Buf.Buffer->getBufferIdentifier();
+  bool IsFileBacked = SM.getFileSystem()->exists(Filename);
+
   // Construct and add to files. If there is an IncludeLoc, the file from
   // IncludeLoc is added before current file.
   assert(CurrentFileID == Files.size() && "File index mismatch");
-  StringRef FileContent = IsFSBacked ? StringRef() : Buf.Buffer->getBuffer();
+
+  StringRef FileContent = Buf.Buffer->getBuffer();
   SerializedFile File = {Filename.str(),
                          convertSourceLoc(SM, SourceLoc(Buf.IncludeLoc)),
-                         FileContent};
+                         IsFileBacked, IsFileBacked ? "" : FileContent.str()};
+
   // Add file to serializing source manager.
-  FileMapper[&SrcMgr].insert({CurrentFileID, deserializeFile(File)});
+  unsigned NewIdx = SrcMgr.addMemBufferCopy(Buf.Buffer.get());
+  FileMapper[&SrcMgr].insert({CurrentFileID, NewIdx});
 
   Files.emplace_back(std::move(File));
   Allocated.insert({Idx, ++CurrentFileID});
@@ -495,21 +511,26 @@ DiagnosticSerializer::deserializeFixIt(const SerializedFixIt &FI) {
   return DiagnosticInfo::FixIt(*Range, FI.Text, {});
 }
 
-unsigned DiagnosticSerializer::deserializeFile(const SerializedFile &File) {
+llvm::Expected<unsigned>
+DiagnosticSerializer::deserializeFile(const SerializedFile &File) {
   assert(File.IncludeLoc.FileID == 0 && "IncludeLoc not supported yet");
   auto FileName = remapFilePath(File.FileName);
-  if (File.Content.empty() && FileName == File.FileName)
-    return SrcMgr.getExternalSourceBufferID(FileName);
-
-  std::unique_ptr<llvm::MemoryBuffer> Content;
-  if (!File.Content.empty())
-    Content = llvm::MemoryBuffer::getMemBufferCopy(File.Content, FileName);
-  else if (auto InputFileOrErr = swift::vfs::getFileOrSTDIN(
-               *SrcMgr.getFileSystem(), File.FileName))
-    Content = llvm::MemoryBuffer::getMemBufferCopy(
-        (*InputFileOrErr)->getBuffer(), FileName);
-
-  return Content ? SrcMgr.addNewSourceBuffer(std::move(Content)) : 0u;
+
+  if (File.IsFileBacked) {
+    auto ID = CAS.parseID(File.Content);
+    if (!ID)
+      return ID.takeError();
+
+    auto Proxy = CAS.getProxy(*ID);
+    if (!Proxy)
+      return Proxy.takeError();
+
+    auto Content = Proxy->getMemoryBuffer(FileName);
+    return SrcMgr.addNewSourceBuffer(std::move(Content));
+  }
+
+  auto Content = llvm::MemoryBuffer::getMemBufferCopy(File.Content, FileName);
+  return SrcMgr.addNewSourceBuffer(std::move(Content));
 }
 
 llvm::Error
@@ -599,6 +620,33 @@ llvm::Error DiagnosticSerializer::deserializeDiagnosticInfo(
 
 llvm::Error
 DiagnosticSerializer::serializeEmittedDiagnostics(llvm::raw_ostream &os) {
+  // Convert all file backed source file into CASIDs.
+  for (auto &File : Files) {
+    if (!File.IsFileBacked || !File.Content.empty())
+      continue;
+
+    auto Ref =
+        SrcMgr.getFileSystem()->getObjectRefForFileContent(File.FileName);
+    if (!Ref)
+      return llvm::createFileError(File.FileName, Ref.getError());
+
+    if (*Ref) {
+      auto CASID = CAS.getID(**Ref);
+      File.Content = CASID.toString();
+      continue;
+    }
+
+    // Probably a file system that is not CAS based. Ingest the buffer.
+    auto Buf = SrcMgr.getFileSystem()->getBufferForFile(File.FileName);
+    if (!Buf)
+      return llvm::createFileError(File.FileName, Buf.getError());
+
+    auto BufRef = CAS.storeFromString({}, (*Buf)->getBuffer());
+    if (!BufRef)
+      return llvm::createFileError(File.FileName, BufRef.takeError());
+    File.Content = CAS.getID(*BufRef).toString();
+  }
+
   llvm::yaml::Output yout(os);
   yout << *this;
   return llvm::Error::success();
@@ -616,8 +664,10 @@ llvm::Error DiagnosticSerializer::doEmitFromCached(llvm::StringRef Buffer,
   unsigned ID = 0;
   for (auto &File : Files) {
     assert(File.IncludeLoc.FileID == 0 && "IncludeLoc not supported yet");
-    unsigned Idx = deserializeFile(File);
-    FileMapper[&SrcMgr].insert({ID++, Idx});
+    auto Idx = deserializeFile(File);
+    if (!Idx)
+      return Idx.takeError();
+    FileMapper[&SrcMgr].insert({ID++, *Idx});
   }
 
   for (auto &VF : VFiles) {
@@ -651,7 +701,7 @@ class CachingDiagnosticsProcessor::Implementation
       : InstanceSourceMgr(Instance.getSourceMgr()),
         InAndOut(
             Instance.getInvocation().getFrontendOptions().InputsAndOutputs),
-        Diags(Instance.getDiags()) {
+        Diags(Instance.getDiags()), CAS(*Instance.getSharedCASInstance()) {
     SmallVector<llvm::MappedPrefix, 4> Prefixes;
     llvm::MappedPrefix::transformJoinedIfValid(
         Instance.getInvocation().getFrontendOptions().CacheReplayPrefixMap,
@@ -681,7 +731,7 @@ class CachingDiagnosticsProcessor::Implementation
 
   llvm::Error replayCachedDiagnostics(llvm::StringRef Buffer) {
     return DiagnosticSerializer::emitDiagnosticsFromCached(
-        Buffer, getDiagnosticSourceMgr(), Diags, Mapper, InAndOut);
+        Buffer, Diags, Mapper, CAS, InAndOut);
   }
 
   void handleDiagnostic(SourceManager &SM,
@@ -691,7 +741,7 @@ class CachingDiagnosticsProcessor::Implementation
            "Caching for a different file system");
     Serializer.handleDiagnostic(SM, Info, [&](const DiagnosticInfo &Info) {
       for (auto *Diag : OrigConsumers)
-        Diag->handleDiagnostic(getDiagnosticSourceMgr(), Info);
+        Diag->handleDiagnostic(Serializer.getSourceMgr(), Info);
       return llvm::Error::success();
     });
   }
@@ -718,10 +768,6 @@ class CachingDiagnosticsProcessor::Implementation
   }
 
 private:
-  SourceManager &getDiagnosticSourceMgr() {
-    return getSerializer().getSourceMgr();
-  }
-
   DiagnosticSerializer &getSerializer() {
     // If the DiagnosticSerializer is not setup, create it. It cannot
     // be created on the creation of CachingDiagnosticsProcessor because the
@@ -730,9 +776,9 @@ class CachingDiagnosticsProcessor::Implementation
     // compiler instance on the first diagnostics and assert if the underlying
     // file system changes on later diagnostics.
     if (!Serializer) {
-      Serializer.reset(
-          new DiagnosticSerializer(InstanceSourceMgr.getFileSystem(), Mapper));
-      Serializer->addInputsToSourceMgr(InAndOut);
+      Serializer.reset(new DiagnosticSerializer(
+          InstanceSourceMgr.getFileSystem(), Mapper, CAS));
+      Serializer->addInputsToSourceMgr(InstanceSourceMgr, InAndOut);
     }
 
     return *Serializer;
@@ -751,6 +797,7 @@ class CachingDiagnosticsProcessor::Implementation
   const FrontendInputsAndOutputs &InAndOut;
   DiagnosticEngine &Diags;
   llvm::PrefixMapper Mapper;
+  llvm::cas::ObjectStore &CAS;
 
   llvm::unique_function<bool(StringRef)> serializedOutputCallback;
 

test/CAS/path_remap.swift

@@ -0,0 +1,51 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %target-swift-frontend -scan-dependencies -module-name Test -module-cache-path %t/clang-module-cache -O \
+// RUN:   -disable-implicit-string-processing-module-import -disable-implicit-concurrency-module-import -parse-stdlib \
+// RUN:   %t/main.swift -o %t/deps.json -swift-version 5 -cache-compile-job -cas-path %t/cas -I %t/include \
+// RUN:   -scanner-prefix-map %swift_src_root=/^src -scanner-prefix-map %t=/^tmp
+
+// RUN: %{python} %S/Inputs/BuildCommandExtractor.py %t/deps.json A > %t/A.cmd
+// RUN: %swift_frontend_plain @%t/A.cmd
+
+// RUN: %{python} %S/Inputs/GenerateExplicitModuleMap.py %t/deps.json > %t/map.json
+// RUN: llvm-cas --cas %t/cas --make-blob --data %t/map.json > %t/map.casid
+// RUN: %{python} %S/Inputs/BuildCommandExtractor.py %t/deps.json Test > %t/MyApp.cmd
+
+// RUN: %target-swift-frontend \
+// RUN:   -c -o %t/main.o -cache-compile-job -cas-path %t/cas \
+// RUN:   -swift-version 5 -disable-implicit-swift-modules \
+// RUN:   -disable-implicit-string-processing-module-import -disable-implicit-concurrency-module-import -parse-stdlib \
+// RUN:   -module-name Test -explicit-swift-module-map-file @%t/map.casid \
+// RUN:   -cache-replay-prefix-map /^src=%swift_src_root -cache-replay-prefix-map /^tmp=%t \
+// RUN:   /^tmp/main.swift @%t/MyApp.cmd
+
+// RUN: %swift-scan-test -action compute_cache_key_from_index -cas-path %t/cas -input 0 -- \
+// RUN:   %target-swift-frontend \
+// RUN:   -c -o %t/main.o -cache-compile-job -cas-path %t/cas \
+// RUN:   -swift-version 5 -disable-implicit-swift-modules \
+// RUN:   -disable-implicit-string-processing-module-import -disable-implicit-concurrency-module-import -parse-stdlib \
+// RUN:   -module-name Test -explicit-swift-module-map-file @%t/map.casid \
+// RUN:   -cache-replay-prefix-map /^src=%swift_src_root -cache-replay-prefix-map /^tmp=%t \
+// RUN:   /^tmp/main.swift @%t/MyApp.cmd > %t/key.casid
+
+// RUN: %swift-scan-test -action replay_result -cas-path %t/cas -id @%t/key.casid -- \
+// RUN:   %target-swift-frontend \
+// RUN:   -c -o %t/main.o -cache-compile-job -cas-path %t/cas \
+// RUN:   -swift-version 5 -disable-implicit-swift-modules \
+// RUN:   -disable-implicit-string-processing-module-import -disable-implicit-concurrency-module-import -parse-stdlib \
+// RUN:   -module-name Test -explicit-swift-module-map-file @%t/map.casid \
+// RUN:   -cache-replay-prefix-map /^src=%swift_src_root -cache-replay-prefix-map /^tmp=%t \
+// RUN:   /^tmp/main.swift @%t/MyApp.cmd
+
+//--- main.swift
+import A
+
+#warning("This is a warning")
+
+//--- include/A.swiftinterface
+// swift-interface-format-version: 1.0
+// swift-module-flags: -module-name A -O -disable-implicit-string-processing-module-import -disable-implicit-concurrency-module-import -parse-stdlib -user-module-version 1.0
+public func b() { }
+