Merge pull request #70384 from cachemeifyoucan/eng/PR-119387650

[Caching] Encoding cache key for input file with index instead of path

Author: cachemeifyoucan

include/swift-c/DependencyScan/DependencyScan.h

@@ -25,7 +25,7 @@
 /// SWIFTSCAN_VERSION_MINOR should increase when there are API additions.
 /// SWIFTSCAN_VERSION_MAJOR is intended for "major" source/ABI breaking changes.
 #define SWIFTSCAN_VERSION_MAJOR 0
-#define SWIFTSCAN_VERSION_MINOR 6
+#define SWIFTSCAN_VERSION_MINOR 7
 
 SWIFTSCAN_BEGIN_DECLS
 
@@ -505,10 +505,26 @@ SWIFTSCAN_PUBLIC void swiftscan_cas_dispose(swiftscan_cas_t cas);
 /// swift input on the command-line by convention. Return \c CacheKey as string.
 /// If error happens, the error message is returned via `error` parameter, and
 /// caller needs to free the error message via `swiftscan_string_dispose`.
+/// This API is DEPRECATED and in favor of using
+/// `swiftscan_cache_compute_key_from_input_index`.
 SWIFTSCAN_PUBLIC swiftscan_string_ref_t
 swiftscan_cache_compute_key(swiftscan_cas_t cas, int argc, const char **argv,
                             const char *input, swiftscan_string_ref_t *error);
 
+/// Compute \c CacheKey for the outputs of a primary input file from a compiler
+/// invocation with command-line \c argc and \c argv and the index for the
+/// input. The index of the input is computed from the position of the input
+/// file from all input files. When primary input file is not available for
+/// compilation, e.g., using WMO, primary file is the first swift input on the
+/// command-line by convention. Return \c CacheKey as string. If error happens,
+/// the error message is returned via `error` parameter, and caller needs to
+/// free the error message via `swiftscan_string_dispose`.
+SWIFTSCAN_PUBLIC swiftscan_string_ref_t
+swiftscan_cache_compute_key_from_input_index(swiftscan_cas_t cas, int argc,
+                                             const char **argv,
+                                             unsigned input_index,
+                                             swiftscan_string_ref_t *error);
+
 /// Query the result of the compilation using the output cache key. \c globally
 /// suggests if the lookup should check remote cache if such operation exists.
 /// Returns the cached compilation of the result if found, or nullptr if output

include/swift/Frontend/CASOutputBackends.h

@@ -33,8 +33,7 @@ class SwiftCASOutputBackend : public llvm::vfs::OutputBackend {
                  llvm::Optional<llvm::vfs::OutputConfig> Config) override;
 
   virtual llvm::Error storeImpl(llvm::StringRef Path, llvm::StringRef Bytes,
-                                llvm::StringRef CorrespondingInput,
-                                file_types::ID OutputKind);
+                                unsigned InputIndex, file_types::ID OutputKind);
 
 private:
   file_types::ID getOutputFileType(llvm::StringRef Path) const;
@@ -47,7 +46,7 @@ class SwiftCASOutputBackend : public llvm::vfs::OutputBackend {
                         FrontendOptions::ActionType Action);
   ~SwiftCASOutputBackend();
 
-  llvm::Error storeCachedDiagnostics(llvm::StringRef InputFile,
+  llvm::Error storeCachedDiagnostics(unsigned InputIndex,
                                      llvm::StringRef Bytes);
 
 private:

include/swift/Frontend/CompileJobCacheKey.h

@@ -36,12 +36,13 @@ createCompileJobBaseCacheKey(llvm::cas::ObjectStore &CAS,
                              ArrayRef<const char *> Args);
 
 /// Compute CompileJobKey for the compiler outputs. The key for the output
-/// is computed from the base key for the compilation, the output kind and the
-/// input file path that is associated with this specific output.
+/// is computed from the base key for the compilation and the input file index
+/// which is the index for the input among all the input files (not just the
+/// output producing inputs).
 llvm::Expected<llvm::cas::ObjectRef>
 createCompileJobCacheKeyForOutput(llvm::cas::ObjectStore &CAS,
                                   llvm::cas::ObjectRef BaseKey,
-                                  StringRef ProducingInput);
+                                  unsigned InputIndex);
 } // namespace swift
 
 #endif

include/swift/Frontend/FrontendInputsAndOutputs.h

@@ -159,6 +159,8 @@ class FrontendInputsAndOutputs {
 
   const InputFile &getFirstOutputProducingInput() const;
 
+  unsigned getIndexOfFirstOutputProducingInput() const;
+
   bool isInputPrimary(StringRef file) const;
 
   unsigned numberOfPrimaryInputsEndingWith(StringRef extension) const;

lib/DependencyScan/ScanDependencies.cpp

@@ -179,7 +179,9 @@ updateModuleCacheKey(ModuleDependencyInfo &depInfo,
   if (cache.getScanService().hasPathMapping())
     InputPath = cache.getScanService().remapPath(InputPath);
 
-  auto key = createCompileJobCacheKeyForOutput(CAS, *base, InputPath);
+  // Module compilation commands always have only one input and the input
+  // index is always 0.
+  auto key = createCompileJobCacheKeyForOutput(CAS, *base, /*InputIndex=*/0);
   if (!key)
     return key.takeError();
 

lib/DriverTool/swift_cache_tool_main.cpp

@@ -284,10 +284,10 @@ int SwiftCacheToolInvocation::printOutputKeys() {
 
   std::vector<OutputEntry> OutputKeys;
   bool hasError = false;
-  auto addFromInputFile = [&](const InputFile &Input) {
+  auto addFromInputFile = [&](const InputFile &Input, unsigned InputIndex) {
     auto InputPath = Input.getFileName();
     auto OutputKey =
-        createCompileJobCacheKeyForOutput(CAS, *BaseKey, InputPath);
+        createCompileJobCacheKeyForOutput(CAS, *BaseKey, InputIndex);
     if (!OutputKey) {
       llvm::errs() << "cannot create cache key for " << InputPath << ": "
                    << toString(OutputKey.takeError()) << "\n";
@@ -310,9 +310,10 @@ int SwiftCacheToolInvocation::printOutputKeys() {
               Outputs.emplace_back(file_types::getTypeName(ID), File);
             });
   };
-  llvm::for_each(
-      Invocation.getFrontendOptions().InputsAndOutputs.getAllInputs(),
-      addFromInputFile);
+  auto AllInputs =
+      Invocation.getFrontendOptions().InputsAndOutputs.getAllInputs();
+  for (unsigned Index = 0; Index < AllInputs.size(); ++Index)
+    addFromInputFile(AllInputs[Index], Index);
 
   // Add diagnostics file.
   if (!OutputKeys.empty())

lib/Frontend/CASOutputBackends.cpp

@@ -68,7 +68,7 @@ class SwiftCASOutputBackend::Implementation {
     auto ProducingInput = OutputToInputMap.find(ResolvedPath);
     assert(ProducingInput != OutputToInputMap.end() && "Unknown output file");
 
-    std::string InputFilename = ProducingInput->second.first.getFileName();
+    unsigned InputIndex = ProducingInput->second.first;
     auto OutputType = ProducingInput->second.second;
 
     // Uncached output kind.
@@ -77,18 +77,18 @@ class SwiftCASOutputBackend::Implementation {
 
     return std::make_unique<SwiftCASOutputFile>(
         ResolvedPath,
-        [this, InputFilename, OutputType](StringRef Path,
-                                          StringRef Bytes) -> Error {
-          return storeImpl(Path, Bytes, InputFilename, OutputType);
+        [this, InputIndex, OutputType](StringRef Path,
+                                       StringRef Bytes) -> Error {
+          return storeImpl(Path, Bytes, InputIndex, OutputType);
         });
   }
 
   void initBackend(const FrontendInputsAndOutputs &InputsAndOutputs);
 
-  Error storeImpl(StringRef Path, StringRef Bytes, StringRef CorrespondingInput,
+  Error storeImpl(StringRef Path, StringRef Bytes, unsigned InputIndex,
                   file_types::ID OutputKind);
 
-  Error finalizeCacheKeysFor(StringRef Input);
+  Error finalizeCacheKeysFor(unsigned InputIndex);
 
 private:
   friend class SwiftCASOutputBackend;
@@ -98,8 +98,11 @@ class SwiftCASOutputBackend::Implementation {
   const FrontendInputsAndOutputs &InputsAndOutputs;
   FrontendOptions::ActionType Action;
 
-  StringMap<std::pair<const InputFile &, file_types::ID>> OutputToInputMap;
-  StringMap<DenseMap<file_types::ID, ObjectRef>> OutputRefs;
+  // Map from output path to the input index and output kind.
+  StringMap<std::pair<unsigned, file_types::ID>> OutputToInputMap;
+
+  // A vector of output refs where the index is the input index.
+  SmallVector<DenseMap<file_types::ID, ObjectRef>> OutputRefs;
 };
 
 SwiftCASOutputBackend::SwiftCASOutputBackend(
@@ -127,14 +130,14 @@ file_types::ID SwiftCASOutputBackend::getOutputFileType(StringRef Path) const {
 }
 
 Error SwiftCASOutputBackend::storeImpl(StringRef Path, StringRef Bytes,
-                                       StringRef CorrespondingInput,
+                                       unsigned InputIndex,
                                        file_types::ID OutputKind) {
-  return Impl.storeImpl(Path, Bytes, CorrespondingInput, OutputKind);
+  return Impl.storeImpl(Path, Bytes, InputIndex, OutputKind);
 }
 
-Error SwiftCASOutputBackend::storeCachedDiagnostics(StringRef InputFile,
+Error SwiftCASOutputBackend::storeCachedDiagnostics(unsigned InputIndex,
                                                     StringRef Bytes) {
-  return storeImpl("<cached-diagnostics>", Bytes, InputFile,
+  return storeImpl("<cached-diagnostics>", Bytes, InputIndex,
                    file_types::ID::TY_CachedDiagnostics);
 }
 
@@ -145,57 +148,63 @@ void SwiftCASOutputBackend::Implementation::initBackend(
   // input it actually comes from. Maybe the solution is just not to cache
   // any commands write output to `-`.
   file_types::ID mainOutputType = InputsAndOutputs.getPrincipalOutputType();
-  auto addInput = [&](const InputFile &Input) {
+  auto addInput = [&](const InputFile &Input, unsigned Index) {
     if (!Input.outputFilename().empty())
       OutputToInputMap.insert(
-          {Input.outputFilename(), {Input, mainOutputType}});
+          {Input.outputFilename(), {Index, mainOutputType}});
     Input.getPrimarySpecificPaths()
         .SupplementaryOutputs.forEachSetOutputAndType(
             [&](const std::string &Out, file_types::ID ID) {
               if (!file_types::isProducedFromDiagnostics(ID))
-                OutputToInputMap.insert({Out, {Input, ID}});
+                OutputToInputMap.insert({Out, {Index, ID}});
             });
   };
-  llvm::for_each(InputsAndOutputs.getAllInputs(), addInput);
+
+  for (unsigned idx = 0; idx < InputsAndOutputs.getAllInputs().size(); ++idx)
+    addInput(InputsAndOutputs.getAllInputs()[idx], idx);
 
   // FIXME: Cached diagnostics is associated with the first output producing
   // input file.
-  OutputToInputMap.insert({"<cached-diagnostics>",
-                           {InputsAndOutputs.getFirstOutputProducingInput(),
-                            file_types::TY_CachedDiagnostics}});
+  OutputToInputMap.insert(
+      {"<cached-diagnostics>",
+       {InputsAndOutputs.getIndexOfFirstOutputProducingInput(),
+        file_types::TY_CachedDiagnostics}});
+
+  // Resize the output refs to hold all inputs.
+  OutputRefs.resize(InputsAndOutputs.getAllInputs().size());
 }
 
 Error SwiftCASOutputBackend::Implementation::storeImpl(
-    StringRef Path, StringRef Bytes, StringRef CorrespondingInput,
+    StringRef Path, StringRef Bytes, unsigned InputIndex,
     file_types::ID OutputKind) {
   Optional<ObjectRef> BytesRef;
   if (Error E = CAS.storeFromString(None, Bytes).moveInto(BytesRef))
     return E;
 
   LLVM_DEBUG(llvm::dbgs() << "DEBUG: producing CAS output of type \'"
                           << file_types::getTypeName(OutputKind)
-                          << "\' for input \'" << CorrespondingInput << "\': \'"
+                          << "\' for input \'" << InputIndex << "\': \'"
                           << CAS.getID(*BytesRef).toString() << "\'\n";);
 
-  OutputRefs[CorrespondingInput].insert({OutputKind, *BytesRef});
+  OutputRefs[InputIndex].insert({OutputKind, *BytesRef});
 
-  return finalizeCacheKeysFor(CorrespondingInput);
+  return finalizeCacheKeysFor(InputIndex);
 }
 
 Error SwiftCASOutputBackend::Implementation::finalizeCacheKeysFor(
-    StringRef Input) {
-  auto Entry = OutputRefs.find(Input);
-  assert(Entry != OutputRefs.end() && "Unexpected input");
+    unsigned InputIndex) {
+  auto ProducedOutputs = OutputRefs[InputIndex];
+  assert(!ProducedOutputs.empty() && "Expect outputs for this input");
 
   // If not all outputs for the input are emitted, return.
   if (!llvm::all_of(OutputToInputMap, [&](auto &E) {
-        return (E.second.first.getFileName() != Input ||
-                Entry->second.count(E.second.second));
+        return (E.second.first != InputIndex ||
+                ProducedOutputs.count(E.second.second));
       }))
     return Error::success();
 
   std::vector<std::pair<file_types::ID, ObjectRef>> OutputsForInput;
-  llvm::for_each(Entry->second, [&OutputsForInput](auto E) {
+  llvm::for_each(ProducedOutputs, [&OutputsForInput](auto E) {
     OutputsForInput.emplace_back(E.first, E.second);
   });
   // Sort to a stable ordering for deterministic output cache object.
@@ -237,14 +246,14 @@ Error SwiftCASOutputBackend::Implementation::finalizeCacheKeysFor(
       return Err;
   }
 
-  auto CacheKey = createCompileJobCacheKeyForOutput(CAS, BaseKey, Input);
+  auto CacheKey = createCompileJobCacheKeyForOutput(CAS, BaseKey, InputIndex);
   if (!CacheKey)
     return CacheKey.takeError();
 
-  LLVM_DEBUG(llvm::dbgs() << "DEBUG: writing cache entry for input \'" << Input
-                          << "\': \'" << CAS.getID(*CacheKey).toString()
-                          << "\' => \'" << CAS.getID(*Result).toString()
-                          << "\'\n";);
+  LLVM_DEBUG(llvm::dbgs() << "DEBUG: writing cache entry for input \'"
+                          << InputIndex << "\': \'"
+                          << CAS.getID(*CacheKey).toString() << "\' => \'"
+                          << CAS.getID(*Result).toString() << "\'\n";);
 
   if (auto E = Cache.put(CAS.getID(*CacheKey), CAS.getID(*Result)))
     return E;

lib/Frontend/CachedDiagnostics.cpp

@@ -777,8 +777,7 @@ CachingDiagnosticsProcessor::CachingDiagnosticsProcessor(
     auto Err = Instance.getCASOutputBackend().storeCachedDiagnostics(
         Instance.getInvocation()
             .getFrontendOptions()
-            .InputsAndOutputs.getFirstOutputProducingInput()
-            .getFileName(),
+            .InputsAndOutputs.getIndexOfFirstOutputProducingInput(),
         Content);
 
     if (Err) {

lib/Frontend/CachingUtils.cpp

@@ -145,10 +145,12 @@ bool replayCachedCompilerOutputs(
   Optional<OutputEntry> DiagnosticsOutput;
 
   auto replayOutputsForInputFile = [&](const std::string &InputPath,
+                                       unsigned InputIndex,
                                        const DenseMap<file_types::ID,
                                                       std::string> &Outputs) {
     auto lookupFailed = [&CanReplayAllOutput] { CanReplayAllOutput = false; };
-    auto OutputKey = createCompileJobCacheKeyForOutput(CAS, BaseKey, InputPath);
+    auto OutputKey =
+        createCompileJobCacheKeyForOutput(CAS, BaseKey, InputIndex);
 
     if (!OutputKey) {
       Diag.diagnose(SourceLoc(), diag::error_cas,
@@ -199,7 +201,8 @@ bool replayCachedCompilerOutputs(
     }
   };
 
-  auto replayOutputFromInput = [&](const InputFile &Input) {
+  auto replayOutputFromInput = [&](const InputFile &Input,
+                                   unsigned InputIndex) {
     auto InputPath = Input.getFileName();
     DenseMap<file_types::ID, std::string> Outputs;
     if (!Input.outputFilename().empty())
@@ -225,27 +228,33 @@ bool replayCachedCompilerOutputs(
     Outputs.try_emplace(file_types::ID::TY_CachedDiagnostics,
                         "<cached-diagnostics>");
 
-    return replayOutputsForInputFile(InputPath, Outputs);
+    return replayOutputsForInputFile(InputPath, InputIndex, Outputs);
   };
 
+  auto AllInputs = InputsAndOutputs.getAllInputs();
   // If there are primary inputs, look up only the primary input files.
   // Otherwise, prepare to do cache lookup for all inputs.
-  if (InputsAndOutputs.hasPrimaryInputs())
-    InputsAndOutputs.forEachPrimaryInput([&](const InputFile &File) {
-      replayOutputFromInput(File);
-      return false;
-    });
-  else
-    llvm::for_each(InputsAndOutputs.getAllInputs(), replayOutputFromInput);
+  for (unsigned Index = 0; Index < AllInputs.size(); ++Index) {
+    const auto &Input = AllInputs[Index];
+    if (InputsAndOutputs.hasPrimaryInputs() && !Input.isPrimary())
+      continue;
+
+    replayOutputFromInput(Input, Index);
+  }
+
+  if (!CanReplayAllOutput)
+    return false;
 
   // If there is not diagnostic output, this is a job that produces no output
   // and only diagnostics, like `typecheck-module-from-interface`, look up
   // diagnostics from first file.
   if (!DiagnosticsOutput)
     replayOutputsForInputFile(
-        InputsAndOutputs.getFirstOutputProducingInput().getFileName(),
+        "<cached-diagnostics>",
+        InputsAndOutputs.getIndexOfFirstOutputProducingInput(),
         {{file_types::ID::TY_CachedDiagnostics, "<cached-diagnostics>"}});
 
+  // Check again to make sure diagnostics is fetched successfully.
   if (!CanReplayAllOutput)
     return false;
 

lib/Frontend/CompileJobCacheKey.cpp

@@ -23,6 +23,7 @@
 #include "llvm/CAS/ObjectStore.h"
 #include "llvm/Option/ArgList.h"
 #include "llvm/Option/OptTable.h"
+#include "llvm/Support/EndianStream.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 
@@ -84,13 +85,16 @@ llvm::Expected<llvm::cas::ObjectRef> swift::createCompileJobBaseCacheKey(
     return Out.takeError();
 }
 
-llvm::Expected<llvm::cas::ObjectRef> swift::createCompileJobCacheKeyForOutput(
-    llvm::cas::ObjectStore &CAS, llvm::cas::ObjectRef BaseKey,
-    StringRef ProducingInput) {
-  SmallString<256> OutputInfo;
+llvm::Expected<llvm::cas::ObjectRef>
+swift::createCompileJobCacheKeyForOutput(llvm::cas::ObjectStore &CAS,
+                                         llvm::cas::ObjectRef BaseKey,
+                                         unsigned InputIndex) {
+  std::string InputInfo;
+  llvm::raw_string_ostream OS(InputInfo);
 
-  // CacheKey is the producting input + the base key.
-  OutputInfo.append(ProducingInput);
+  // CacheKey is the index of the producting input + the base key.
+  // Encode the unsigned value as little endian in the field.
+  llvm::support::endian::write<uint32_t>(OS, InputIndex, llvm::support::little);
 
-  return CAS.storeFromString({BaseKey}, OutputInfo);
+  return CAS.storeFromString({BaseKey}, OS.str());
 }