[CompilerCaching] Teach CachedDiagnostics to replay with prefix map

Teach CachedDiagnosticsProcessor to replay diagnostics for all consumers
with a different path using a path prefix map. This allows diagnostics
to be replayed for the actual path on disk after the scanner
canonicalized the path for compilation.

Author: cachemeifyoucan

include/swift/Frontend/FrontendOptions.h

@@ -148,6 +148,9 @@ class FrontendOptions {
   /// CacheKey for input file.
   std::string InputFileKey;
 
+  /// CacheReplay PrefixMap.
+  std::vector<std::string> CacheReplayPrefixMap;
+
   /// Number of retry opening an input file if the previous opening returns
   /// bad file descriptor error.
   unsigned BadFileDescriptorRetryCount = 0;

include/swift/Option/Options.td

@@ -1865,6 +1865,10 @@ def scanner_prefix_map_toolchain : Separate<["-"], "scanner-prefix-map-toolchain
   Flags<[NewDriverOnlyOption]>,
   HelpText<"Remap paths within toolchain directory reported by dependency scanner">, MetaVarName<"<path>">;
 
+def cache_replay_prefix_map: Separate<["-"], "cache-replay-prefix-map">,
+  Flags<[FrontendOption, NoDriverOption, CacheInvariant]>,
+  HelpText<"Remap paths when replaying outputs from cache">, MetaVarName<"<prefix=replacement>">;
+
 // END ONLY SUPPORTED IN NEW DRIVER
 
 

lib/Frontend/ArgsToFrontendOptionsConverter.cpp

@@ -270,6 +270,7 @@ bool ArgsToFrontendOptionsConverter::convert(
   Opts.CASFSRootIDs = Args.getAllArgValues(OPT_cas_fs);
   Opts.ClangIncludeTrees = Args.getAllArgValues(OPT_clang_include_tree_root);
   Opts.InputFileKey = Args.getLastArgValue(OPT_input_file_key);
+  Opts.CacheReplayPrefixMap = Args.getAllArgValues(OPT_cache_replay_prefix_map);
 
   if (Opts.EnableCaching && Opts.CASFSRootIDs.empty() &&
       Opts.ClangIncludeTrees.empty() &&

lib/Frontend/CachedDiagnostics.cpp

@@ -29,6 +29,7 @@
 #include "llvm/Support/Compression.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/PrefixMapper.h"
 #include "llvm/Support/SMLoc.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Support/YAMLTraits.h"
@@ -98,8 +99,9 @@ struct SerializedGeneratedFileInfo {
 };
 
 struct DiagnosticSerializer {
-  DiagnosticSerializer(llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS)
-      : SrcMgr(FS) {}
+  DiagnosticSerializer(llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS,
+                       llvm::PrefixMapper &Mapper)
+      : SrcMgr(FS), Mapper(Mapper) {}
 
   using ReplayFunc = llvm::function_ref<llvm::Error(const DiagnosticInfo &)>;
 
@@ -111,10 +113,11 @@ struct DiagnosticSerializer {
   static llvm::Error
   emitDiagnosticsFromCached(llvm::StringRef Buffer, SourceManager &SrcMgr,
                             DiagnosticEngine &Diags,
+                            llvm::PrefixMapper &Mapper,
                             const FrontendInputsAndOutputs &InAndOut) {
     // Create a new DiagnosticSerializer since this cannot be shared with a
     // serialization instance.
-    DiagnosticSerializer DS(SrcMgr.getFileSystem());
+    DiagnosticSerializer DS(SrcMgr.getFileSystem(), Mapper);
     DS.addInputsToSourceMgr(InAndOut);
     return DS.doEmitFromCached(Buffer, Diags);
   }
@@ -128,7 +131,7 @@ struct DiagnosticSerializer {
     // has references to input files to find subconsumer.
     auto addInputToSourceMgr = [&](const InputFile &Input) {
       if (Input.getFileName() != "-")
-        SrcMgr.getExternalSourceBufferID(Input.getFileName());
+        SrcMgr.getExternalSourceBufferID(remapFilePath(Input.getFileName()));
       return false;
     };
     InAndOut.forEachInputProducingSupplementaryOutput(addInputToSourceMgr);
@@ -163,6 +166,9 @@ struct DiagnosticSerializer {
   unsigned deserializeFile(const SerializedFile &File);
   llvm::Error deserializeVirtualFile(const SerializedVirtualFile &VF);
   llvm::Error deserializeGeneratedFileInfo(const SerializedGeneratedFileInfo &Info);
+  std::string remapFilePath(StringRef Path) {
+    return Mapper.mapToString(Path);
+  }
 
 public:
   std::vector<SerializedDiagnosticInfo> DiagInfos;
@@ -176,6 +182,7 @@ struct DiagnosticSerializer {
 
   // Serializing SourceManager.
   SourceManager SrcMgr;
+  llvm::PrefixMapper &Mapper;
 
   // Mapping of the FileID between SourceManager from CompilerInstance vs.
   // the serialized FileID in cached diagnostics. Lookup tables are
@@ -482,10 +489,19 @@ DiagnosticSerializer::deserializeFixIt(const SerializedFixIt &FI) {
 
 unsigned DiagnosticSerializer::deserializeFile(const SerializedFile &File) {
   assert(File.IncludeLoc.FileID == 0 && "IncludeLoc not supported yet");
-  return File.Content.empty()
-             ? SrcMgr.getExternalSourceBufferID(File.FileName)
-             : SrcMgr.addNewSourceBuffer(llvm::MemoryBuffer::getMemBufferCopy(
-                   File.Content, File.FileName));
+  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;
 }
 
 llvm::Error
@@ -495,8 +511,8 @@ DiagnosticSerializer::deserializeVirtualFile(const SerializedVirtualFile &VF) {
     return Range.takeError();
   unsigned Length = (const char *)Range->getEnd().getOpaquePointerValue() -
                     (const char *)Range->getStart().getOpaquePointerValue();
-  SrcMgr.createVirtualFile(Range->getStart(), VF.FileName, VF.LineOffset,
-                           Length);
+  auto FileName = remapFilePath(VF.FileName);
+  SrcMgr.createVirtualFile(Range->getStart(), FileName, VF.LineOffset, Length);
   return llvm::Error::success();
 }
 
@@ -625,7 +641,14 @@ class CachingDiagnosticsProcessor::Implementation
       : InstanceSourceMgr(Instance.getSourceMgr()),
         InAndOut(
             Instance.getInvocation().getFrontendOptions().InputsAndOutputs),
-        Diags(Instance.getDiags()) {}
+        Diags(Instance.getDiags()) {
+    SmallVector<llvm::MappedPrefix, 4> Prefixes;
+    llvm::MappedPrefix::transformJoinedIfValid(
+        Instance.getInvocation().getFrontendOptions().CacheReplayPrefixMap,
+        Prefixes);
+    Mapper.addRange(Prefixes);
+    Mapper.sort();
+  }
   ~Implementation() {}
 
   void startDiagnosticCapture() {
@@ -648,7 +671,7 @@ class CachingDiagnosticsProcessor::Implementation
 
   llvm::Error replayCachedDiagnostics(llvm::StringRef Buffer) {
     return DiagnosticSerializer::emitDiagnosticsFromCached(
-        Buffer, getDiagnosticSourceMgr(), Diags, InAndOut);
+        Buffer, getDiagnosticSourceMgr(), Diags, Mapper, InAndOut);
   }
 
   void handleDiagnostic(SourceManager &SM,
@@ -698,7 +721,7 @@ class CachingDiagnosticsProcessor::Implementation
     // file system changes on later diagnostics.
     if (!Serializer) {
       Serializer.reset(
-          new DiagnosticSerializer(InstanceSourceMgr.getFileSystem()));
+          new DiagnosticSerializer(InstanceSourceMgr.getFileSystem(), Mapper));
       Serializer->addInputsToSourceMgr(InAndOut);
     }
 
@@ -717,6 +740,7 @@ class CachingDiagnosticsProcessor::Implementation
   SourceManager &InstanceSourceMgr;
   const FrontendInputsAndOutputs &InAndOut;
   DiagnosticEngine &Diags;
+  llvm::PrefixMapper Mapper;
 
   llvm::unique_function<bool(StringRef)> serializedOutputCallback;
 

test/CAS/cached_diagnostics.swift

@@ -19,3 +19,10 @@
 
 // Verify the serialized diags have the right magic at the top.
 // CHECK-SERIALIZED: DIA
+
+/// Check path remapping.
+// RUN: %cache-tool -cas-path %t/cas -cache-tool-action render-diags %t/cache_key.json -- %target-swift-frontend -c -cache-compile-job -cas-path %t/cas -allow-unstable-cache-key-for-testing %s \
+// RUN:   -import-objc-header %S/Inputs/objc.h -emit-module -emit-module-path %t/test.swiftmodule -cache-replay-prefix-map %S=/^test  2>&1 | %FileCheck %s --check-prefix REMAP
+
+// REMAP: /^test/Inputs/objc.h:3:2: warning: warning in bridging header
+// REMAP: /^test/cached_diagnostics.swift:10:10: warning: this is a warning