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

Avoid path encoding difference (for example, real_path vs. path from symlink) by eliminating the path from cache key. Cache key is now encoded with the index of the input file from all the input files from the command-line, reguardless if those inputs will produce output or not. This is to ensure stable ordering even the batching is different. Add a new cache computation API that is preferred for using input index directly. Old API for cache key is deprecated but still updated to fallback to real_path comparsion if needed. As a result of swift scan API change, rename the feature in JSON file to avoid version confusion between swift-driver and libSwiftScan. rdar://119387650
2025-12-21 12:14:44 +01:00 · 2023-12-11 12:05:11 -08:00
parent 443cac20fb
commit 76bde39ee7
16 changed files with 246 additions and 87 deletions
--- a/include/swift-c/DependencyScan/DependencyScan.h
+++ b/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
--- a/include/swift/Frontend/CASOutputBackends.h
+++ b/include/swift/Frontend/CASOutputBackends.h
@@ -33,8 +33,7 @@ protected:
                 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 @@ public:
                        FrontendOptions::ActionType Action);
  ~SwiftCASOutputBackend();

-  llvm::Error storeCachedDiagnostics(llvm::StringRef InputFile,
+  llvm::Error storeCachedDiagnostics(unsigned InputIndex,
                                     llvm::StringRef Bytes);

 private:
--- a/include/swift/Frontend/CompileJobCacheKey.h
+++ b/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
--- a/include/swift/Frontend/FrontendInputsAndOutputs.h
+++ b/include/swift/Frontend/FrontendInputsAndOutputs.h
@@ -159,6 +159,8 @@ public:

  const InputFile &getFirstOutputProducingInput() const;

+  unsigned getIndexOfFirstOutputProducingInput() const;
+
  bool isInputPrimary(StringRef file) const;

  unsigned numberOfPrimaryInputsEndingWith(StringRef extension) const;
--- a/lib/DependencyScan/ScanDependencies.cpp
+++ b/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();

--- a/lib/DriverTool/swift_cache_tool_main.cpp
+++ b/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())
--- a/lib/Frontend/CASOutputBackends.cpp
+++ b/lib/Frontend/CASOutputBackends.cpp
@@ -68,7 +68,7 @@ public:
    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 @@ public:

    return std::make_unique<SwiftCASOutputFile>(
        ResolvedPath,
-        [this, InputFilename, OutputType](StringRef Path,
+        [this, InputIndex, OutputType](StringRef Path,
                                       StringRef Bytes) -> Error {
-          return storeImpl(Path, Bytes, InputFilename, OutputType);
+          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 @@ private:
  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,28 +148,34 @@ 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(),
+  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))
@@ -174,28 +183,28 @@ Error SwiftCASOutputBackend::Implementation::storeImpl(

  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;
--- a/lib/Frontend/CachedDiagnostics.cpp
+++ b/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) {
--- a/lib/Frontend/CachingUtils.cpp
+++ b/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);
+  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;
-    });
-  else
-    llvm::for_each(InputsAndOutputs.getAllInputs(), replayOutputFromInput);

  // 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;

--- a/lib/Frontend/CompileJobCacheKey.cpp
+++ b/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());
 }
--- a/lib/Frontend/FrontendInputsAndOutputs.cpp
+++ b/lib/Frontend/FrontendInputsAndOutputs.cpp
@@ -159,6 +159,10 @@ FrontendInputsAndOutputs::getFirstOutputProducingInput() const {
  return isWholeModule() ? firstInput() : firstPrimaryInput();
 }

+unsigned FrontendInputsAndOutputs::getIndexOfFirstOutputProducingInput() const {
+  return isWholeModule() ? 0 : PrimaryInputsInOrder[0];
+}
+
 bool FrontendInputsAndOutputs::isInputPrimary(StringRef file) const {
  return primaryInputNamed(file) != nullptr;
 }
--- a/lib/Option/features.json
+++ b/lib/Option/features.json
@@ -40,7 +40,7 @@
      "name": "ld-path-driver-option"
    },
    {
-      "name": "cache-compile-job"
+      "name": "compilation-caching"
    }
  ]
 }
--- a/test/CAS/swift-scan-test.swift
+++ b/test/CAS/swift-scan-test.swift
@@ -4,6 +4,12 @@
 // RUN:   -emit-module -emit-module-path %t/Test.swiftmodule -c -emit-dependencies -module-name Test -o %t/test.o -cas-path %t/cas \
 // RUN:   -allow-unstable-cache-key-for-testing > %t/key.casid

+// RUN: %swift-scan-test -action compute_cache_key_from_index -cas-path %t/cas -input 0 -- %target-swift-frontend -cache-compile-job -Rcache-compile-job %s \
+// RUN:   -emit-module -emit-module-path %t/Test.swiftmodule -c -emit-dependencies -module-name Test -o %t/test.o -cas-path %t/cas \
+// RUN:   -allow-unstable-cache-key-for-testing > %t/key1.casid
+
+// RUN: diff %t/key.casid %t/key1.casid
+
 // RUN: not %swift-scan-test -action cache_query -id @%t/key.casid -cas-path %t/cas 2>&1 | %FileCheck %s --check-prefix=CHECK-QUERY-NOT-FOUND

 // RUN: %target-swift-frontend -cache-compile-job -Rcache-compile-job %s -emit-module -emit-module-path %t/Test.swiftmodule -c -emit-dependencies \
--- a/tools/libSwiftScan/SwiftCaching.cpp
+++ b/tools/libSwiftScan/SwiftCaching.cpp
@@ -38,7 +38,9 @@
 #include "llvm/CAS/CASReference.h"
 #include "llvm/CAS/ObjectStore.h"
 #include "llvm/Support/Allocator.h"
+#include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/PrefixMapper.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/VirtualOutputBackend.h"
@@ -73,19 +75,19 @@ struct SwiftCachedCompilationHandle {
  SwiftCachedCompilationHandle(llvm::cas::ObjectRef Key,
                               llvm::cas::ObjectRef Output,
                               clang::cas::CompileJobCacheResult &&Result,
-                               llvm::StringRef Input, SwiftScanCAS &CAS)
-      : Key(Key), Output(Output), CorrespondingInput(Input), Result(Result),
+                               unsigned InputIndex, SwiftScanCAS &CAS)
+      : Key(Key), Output(Output), InputIndex(InputIndex), Result(Result),
        DB(CAS) {}
  SwiftCachedCompilationHandle(llvm::cas::ObjectRef Key,
                               llvm::cas::ObjectRef Output,
                               swift::cas::CompileJobCacheResult &&Result,
-                               llvm::StringRef Input, SwiftScanCAS &CAS)
-      : Key(Key), Output(Output), CorrespondingInput(Input), Result(Result),
+                               unsigned InputIndex, SwiftScanCAS &CAS)
+      : Key(Key), Output(Output), InputIndex(InputIndex), Result(Result),
        DB(CAS) {}

  llvm::cas::ObjectRef Key;
  llvm::cas::ObjectRef Output;
-  std::string CorrespondingInput;
+  unsigned InputIndex;
  std::variant<swift::cas::CompileJobCacheResult,
               clang::cas::CompileJobCacheResult>
      Result;
@@ -201,10 +203,67 @@ computeCacheKey(llvm::cas::ObjectStore &CAS, llvm::ArrayRef<const char *> Args,
  if (!BaseKey)
    return BaseKey.takeError();

-  auto Key = swift::createCompileJobCacheKeyForOutput(CAS, *BaseKey, InputPath);
+  // Parse the arguments to figure out the index for the input.
+  swift::CompilerInvocation Invocation;
+  swift::SourceManager SourceMgr;
+  swift::DiagnosticEngine Diags(SourceMgr);
+  llvm::SmallString<128> workingDirectory;
+  llvm::sys::fs::current_path(workingDirectory);
+  llvm::SmallVector<std::unique_ptr<llvm::MemoryBuffer>, 4>
+      configurationFileBuffers;
+
+  std::string MainExecutablePath = llvm::sys::fs::getMainExecutable(
+      "swift-frontend", (void *)swiftscan_cache_replay_compilation);
+
+  // Drop the `-frontend` option if it is passed.
+  if (llvm::StringRef(Args.front()) == "-frontend")
+    Args = Args.drop_front();
+
+  if (Invocation.parseArgs(Args, Diags, &configurationFileBuffers,
+                           workingDirectory, MainExecutablePath))
+    return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                   "Argument parsing failed");
+
+  auto computeKey = [&](unsigned Index) -> llvm::Expected<std::string> {
+    auto Key = swift::createCompileJobCacheKeyForOutput(CAS, *BaseKey, Index);
    if (!Key)
      return Key.takeError();
+    return CAS.getID(*Key).toString();
+  };
+  auto AllInputs =
+      Invocation.getFrontendOptions().InputsAndOutputs.getAllInputs();
+  // First pass, check for path equal.
+  for (unsigned Idx = 0; Idx < AllInputs.size(); ++Idx) {
+    if (AllInputs[Idx].getFileName() == InputPath)
+      return computeKey(Idx);
+  }

+  // If not found, slow second iteration with real_path.
+  llvm::SmallString<256> InputRealPath;
+  llvm::sys::fs::real_path(InputPath, InputRealPath, true);
+  for (unsigned Idx = 0; Idx < AllInputs.size(); ++Idx) {
+    llvm::SmallString<256> TestRealPath;
+    llvm::sys::fs::real_path(AllInputs[Idx].getFileName(), TestRealPath, true);
+    if (InputRealPath == TestRealPath)
+      return computeKey(Idx);
+  }
+
+  return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                 "requested input not found from invocation");
+}
+
+static llvm::Expected<std::string>
+computeCacheKeyFromIndex(llvm::cas::ObjectStore &CAS,
+                         llvm::ArrayRef<const char *> Args,
+                         unsigned InputIndex) {
+  auto BaseKey = swift::createCompileJobBaseCacheKey(CAS, Args);
+  if (!BaseKey)
+    return BaseKey.takeError();
+
+  auto Key =
+      swift::createCompileJobCacheKeyForOutput(CAS, *BaseKey, InputIndex);
+  if (!Key)
+    return Key.takeError();
  return CAS.getID(*Key).toString();
 }

@@ -225,6 +284,26 @@ swiftscan_cache_compute_key(swiftscan_cas_t cas, int argc, const char **argv,
  return swift::c_string_utils::create_clone(ID->c_str());
 }

+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) {
+  std::vector<const char *> Compilation;
+  for (int i = 0; i < argc; ++i)
+    Compilation.push_back(argv[i]);
+
+  auto ID =
+      computeCacheKeyFromIndex(unwrap(cas)->getCAS(), Compilation, input_index);
+  if (!ID) {
+    *error =
+        swift::c_string_utils::create_clone(toString(ID.takeError()).c_str());
+    return swift::c_string_utils::create_null();
+  }
+  *error = swift::c_string_utils::create_null();
+  return swift::c_string_utils::create_clone(ID->c_str());
+}
+
 // Create a non-owning string ref that is used in call backs.
 static swiftscan_string_ref_t createNonOwningString(llvm::StringRef Str) {
  if (Str.empty())
@@ -259,6 +338,9 @@ createCachedCompilation(SwiftScanCAS &CAS, const llvm::cas::CASID &ID,
    return KeyProxy.takeError();
  auto Input = KeyProxy->getData();

+  unsigned Index =
+      llvm::support::endian::read<uint32_t, llvm::support::little,
+                                  llvm::support::unaligned>(Input.data());
  {
    swift::cas::CompileJobResultSchema Schema(CAS.getCAS());
    if (Schema.isRootNode(*Proxy)) {
@@ -266,7 +348,7 @@ createCachedCompilation(SwiftScanCAS &CAS, const llvm::cas::CASID &ID,
      if (!Result)
        return Result.takeError();
      return new SwiftCachedCompilationHandle(KeyProxy->getRef(), *Ref,
-                                              std::move(*Result), Input, CAS);
+                                              std::move(*Result), Index, CAS);
    }
  }
  {
@@ -276,7 +358,7 @@ createCachedCompilation(SwiftScanCAS &CAS, const llvm::cas::CASID &ID,
      if (!Result)
        return Result.takeError();
      return new SwiftCachedCompilationHandle(KeyProxy->getRef(), *Ref,
-                                              std::move(*Result), Input, CAS);
+                                              std::move(*Result), Index, CAS);
    }
  }
  return createUnsupportedSchemaError();
@@ -759,13 +841,10 @@ static llvm::Error replayCompilation(SwiftScanReplayInstance &Instance,
  auto &InputsAndOutputs =
      Instance.Invocation.getFrontendOptions().InputsAndOutputs;
  auto AllInputs = InputsAndOutputs.getAllInputs();
-  auto Input = llvm::find_if(AllInputs, [&](const InputFile &Input) {
-    return Input.getFileName() == Comp.CorrespondingInput;
-  });
-  if (Input == AllInputs.end())
+  if (Comp.InputIndex >= AllInputs.size())
    return createStringError(inconvertibleErrorCode(),
-                             "InputFile \"" + Comp.CorrespondingInput +
-                                 "\" is not part of the compilation");
+                             "InputFile index too large for compilation");
+  const auto &Input = AllInputs[Comp.InputIndex];

  // Setup DiagnosticsConsumers.
  // FIXME: Reduce code duplication against `performFrontend()` and add support
@@ -793,10 +872,10 @@ static llvm::Error replayCompilation(SwiftScanReplayInstance &Instance,

  // Collect the file that needs to write.
  DenseMap<file_types::ID, std::string> Outputs;
-  if (!Input->outputFilename().empty())
+  if (!Input.outputFilename().empty())
    Outputs.try_emplace(InputsAndOutputs.getPrincipalOutputType(),
-                        Input->outputFilename());
-  Input->getPrimarySpecificPaths().SupplementaryOutputs.forEachSetOutputAndType(
+                        Input.outputFilename());
+  Input.getPrimarySpecificPaths().SupplementaryOutputs.forEachSetOutputAndType(
      [&](const std::string &File, file_types::ID ID) {
        if (file_types::isProducedFromDiagnostics(ID))
          return;
--- a/tools/libSwiftScan/libSwiftScan.exports
+++ b/tools/libSwiftScan/libSwiftScan.exports
@@ -85,6 +85,7 @@ swiftscan_cas_options_set_plugin_path
 swiftscan_cas_options_set_plugin_option
 swiftscan_cas_store
 swiftscan_cache_compute_key
+swiftscan_cache_compute_key_from_input_index
 swiftscan_cache_query
 swiftscan_cache_query_async
 swiftscan_cached_compilation_get_num_outputs
--- a/tools/swift-scan-test/swift-scan-test.cpp
+++ b/tools/swift-scan-test/swift-scan-test.cpp
@@ -17,6 +17,7 @@
 #include "swift-c/DependencyScan/DependencyScan.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/FileTypes.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/StringSaver.h"
@@ -26,6 +27,7 @@ using namespace llvm;
 namespace {
 enum Actions {
  compute_cache_key,
+  compute_cache_key_from_index,
  cache_query,
  replay_result,
 };
@@ -35,11 +37,13 @@ llvm::cl::opt<std::string> CASPath("cas-path", llvm::cl::desc("<path>"),
                                   llvm::cl::cat(Category));
 llvm::cl::opt<std::string> CASID("id", llvm::cl::desc("<casid>"),
                                 llvm::cl::cat(Category));
-llvm::cl::opt<std::string> Input("input", llvm::cl::desc("<file>"),
+llvm::cl::opt<std::string> Input("input", llvm::cl::desc("<file|index>"),
                                 llvm::cl::cat(Category));
 llvm::cl::opt<Actions>
    Action("action", llvm::cl::desc("<action>"),
           llvm::cl::values(clEnumVal(compute_cache_key, "compute cache key"),
+                            clEnumVal(compute_cache_key_from_index,
+                                      "compute cache key from index"),
                            clEnumVal(cache_query, "cache query"),
                            clEnumVal(replay_result, "replay result")),
           llvm::cl::cat(Category));
@@ -78,6 +82,27 @@ static int action_compute_cache_key(swiftscan_cas_t cas, StringRef input,
  return EXIT_SUCCESS;
 }

+static int
+action_compute_cache_key_from_index(swiftscan_cas_t cas, StringRef index,
+                                    std::vector<const char *> &Args) {
+  unsigned inputIndex = 0;
+  if (!to_integer(index, inputIndex)) {
+    llvm::errs() << "-input is not a number for compute_cache_key_from_index\n";
+    return EXIT_FAILURE;
+  }
+
+  swiftscan_string_ref_t err_msg;
+  auto key = swiftscan_cache_compute_key_from_input_index(
+      cas, Args.size(), Args.data(), inputIndex, &err_msg);
+  if (key.length == 0)
+    return printError(err_msg);
+
+  llvm::outs() << toString(key) << "\n";
+  swiftscan_string_dispose(key);
+
+  return EXIT_SUCCESS;
+}
+
 static int print_cached_compilation(swiftscan_cached_compilation_t comp,
                                    const char *key) {
  auto numOutput = swiftscan_cached_compilation_get_num_outputs(comp);
@@ -194,6 +219,8 @@ int main(int argc, char *argv[]) {
  switch (Action) {
  case compute_cache_key:
    return action_compute_cache_key(cas, Input, Args);
+  case compute_cache_key_from_index:
+    return action_compute_cache_key_from_index(cas, Input, Args);
  case cache_query:
    return action_cache_query(cas, CASID.c_str());
  case replay_result: