[Caching] Re-associate diagnostics cache key with InputFile

Change how cached diagnostics are stored inside the CAS. It used to be
stored as a standalone entry for a frontend invocation in the cache and
now it is switched to be associated with input files, stored together
with other outputs like object files, etc.

This enables cleaner Cache Replay APIs and future cached diagnostics
that can be splitted up by file contribution.

include/swift/Frontend/FrontendInputsAndOutputs.h

@@ -157,6 +157,8 @@ public:
   /// instead of just answering "batch" if there is more than one primary.
   std::string getStatsFileMangledInputName() const;
 
+  const InputFile &getFirstOutputProducingInput() const;
+
   bool isInputPrimary(StringRef file) const;
 
   unsigned numberOfPrimaryInputsEndingWith(StringRef extension) const;

lib/DriverTool/swift_cache_tool_main.cpp

@@ -50,9 +50,8 @@ enum class SwiftCacheToolAction {
 
 struct OutputEntry {
   std::string InputPath;
-  std::string OutputPath;
-  std::string OutputKind;
   std::string CacheKey;
+  std::vector<std::pair<std::string, std::string>> Outputs;
 };
 
 enum ID {
@@ -285,35 +284,30 @@ int SwiftCacheToolInvocation::printOutputKeys() {
 
   std::vector<OutputEntry> OutputKeys;
   bool hasError = false;
-  auto addOutputKey = [&](StringRef InputPath, file_types::ID OutputKind,
+  auto addFromInputFile = [&](const InputFile &Input) {
-                          StringRef OutputPath) {
+    auto InputPath = Input.getFileName();
     auto OutputKey =
         createCompileJobCacheKeyForOutput(CAS, *BaseKey, InputPath);
     if (!OutputKey) {
-      llvm::errs() << "cannot create cache key for " << OutputPath << ": "
+      llvm::errs() << "cannot create cache key for " << InputPath << ": "
                    << toString(OutputKey.takeError()) << "\n";
       hasError = true;
     }
     OutputKeys.emplace_back(
-        OutputEntry{InputPath.str(), OutputPath.str(),
+        OutputEntry{InputPath, CAS.getID(*OutputKey).toString(), {}});
-                    file_types::getTypeName(OutputKind).str(),
+    auto &Outputs = OutputKeys.back().Outputs;
-                    CAS.getID(*OutputKey).toString()});
-  };
-  auto addFromInputFile = [&](const InputFile &Input) {
-    auto InputPath = Input.getFileName();
     if (!Input.outputFilename().empty())
-      addOutputKey(InputPath,
+      Outputs.emplace_back(file_types::getTypeName(
-                   Invocation.getFrontendOptions()
+                               Invocation.getFrontendOptions()
-                       .InputsAndOutputs.getPrincipalOutputType(),
+                                   .InputsAndOutputs.getPrincipalOutputType()),
-                   Input.outputFilename());
+                           Input.outputFilename());
     Input.getPrimarySpecificPaths()
         .SupplementaryOutputs.forEachSetOutputAndType(
             [&](const std::string &File, file_types::ID ID) {
               // Dont print serialized diagnostics.
               if (file_types::isProducedFromDiagnostics(ID))
                 return;
-
+              Outputs.emplace_back(file_types::getTypeName(ID), File);
-              addOutputKey(InputPath, ID, File);
             });
   };
   llvm::for_each(
@@ -321,8 +315,10 @@ int SwiftCacheToolInvocation::printOutputKeys() {
       addFromInputFile);
 
   // Add diagnostics file.
-  addOutputKey("<cached-diagnostics>", file_types::ID::TY_CachedDiagnostics,
+  if (!OutputKeys.empty())
-               "<cached-diagnostics>");
+    OutputKeys.front().Outputs.emplace_back(
+        file_types::getTypeName(file_types::ID::TY_CachedDiagnostics),
+        "<cached-diagnostics>");
 
   if (hasError)
     return 1;
@@ -331,10 +327,16 @@ int SwiftCacheToolInvocation::printOutputKeys() {
   Out.array([&] {
     for (const auto &E : OutputKeys) {
       Out.object([&] {
-        Out.attribute("OutputPath", E.OutputPath);
-        Out.attribute("OutputKind", E.OutputKind);
         Out.attribute("Input", E.InputPath);
         Out.attribute("CacheKey", E.CacheKey);
+        Out.attributeArray("Outputs", [&] {
+          for (const auto &OutEntry : E.Outputs) {
+            Out.object([&] {
+              Out.attribute("Kind", OutEntry.first);
+              Out.attribute("Path", OutEntry.second);
+            });
+          }
+        });
       });
     }
   });

lib/Frontend/CASOutputBackends.cpp

@@ -157,6 +157,12 @@ void SwiftCASOutputBackend::Implementation::initBackend(
             });
   };
   llvm::for_each(InputsAndOutputs.getAllInputs(), addInput);
+
+  // FIXME: Cached diagnostics is associated with the first output producing
+  // input file.
+  OutputToInputMap.insert({"<cached-diagnostics>",
+                           {InputsAndOutputs.getFirstOutputProducingInput(),
+                            file_types::TY_CachedDiagnostics}});
 }
 
 Error SwiftCASOutputBackend::Implementation::storeImpl(

lib/Frontend/CachedDiagnostics.cpp

@@ -772,10 +772,14 @@ CachingDiagnosticsProcessor::CachingDiagnosticsProcessor(
     }
 
     StringRef Content = Compression.empty() ? Output : toStringRef(Compression);
-    // Store CachedDiagnostics in the CAS/Cache. There is no real associated
+    // Store CachedDiagnostics in the CAS/Cache.
-    // inputs.
+    // FIXME: Currently associated with first output producing input file.
     auto Err = Instance.getCASOutputBackend().storeCachedDiagnostics(
-        "<cached-diagnostics>", Content);
+        Instance.getInvocation()
+            .getFrontendOptions()
+            .InputsAndOutputs.getFirstOutputProducingInput()
+            .getFileName(),
+        Content);
 
     if (Err) {
       Instance.getDiags().diagnose(SourceLoc(), diag::error_cas,

lib/Frontend/CachingUtils.cpp

@@ -209,19 +209,15 @@ bool replayCachedCompilerOutputs(
               Outputs.try_emplace(ID, File);
             });
 
-    // Nothing to replay.
+    // Add cached diagnostic entry for lookup. Output path doesn't matter here.
-    if (Outputs.empty())
+    Outputs.try_emplace(file_types::ID::TY_CachedDiagnostics,
-      return;
+                        "<cached-diagnostics>");
 
     return replayOutputsForInputFile(InputPath, Outputs);
   };
 
   llvm::for_each(InputsAndOutputs.getAllInputs(), replayOutputFromInput);
 
-  replayOutputsForInputFile(
-      "<cached-diagnostics>",
-      {{file_types::ID::TY_CachedDiagnostics, "<cached-diagnostics>"}});
-
   if (!CanReplayAllOutput)
     return false;
 

lib/Frontend/FrontendInputsAndOutputs.cpp

@@ -150,6 +150,15 @@ std::string FrontendInputsAndOutputs::getStatsFileMangledInputName() const {
   return isWholeModule() ? "all" : firstPrimaryInput().getFileName();
 }
 
+const InputFile &
+FrontendInputsAndOutputs::getFirstOutputProducingInput() const {
+  // Get the first input file that produces the output file. That is currently
+  // used to compute with input should the cached diagnostics be associated
+  // with. The first output producing input file is the first input if using
+  // whole module, or first primary input if not using whole module.
+  return isWholeModule() ? firstInput() : firstPrimaryInput();
+}
+
 bool FrontendInputsAndOutputs::isInputPrimary(StringRef file) const {
   return primaryInputNamed(file) != nullptr;
 }

test/CAS/Inputs/ExtractOutputKey.py

@@ -10,8 +10,9 @@ output_path = sys.argv[2]
 
 
 with open(input_json, 'r') as file:
-    outputs = json.load(file)
+    entries = json.load(file)
-    for output in outputs:
+    for entry in entries:
-        if output['OutputPath'] != output_path:
+        for output in entry["Outputs"]:
-            continue
+            if output['Path'] != output_path:
-        print(output['CacheKey'])
+                continue
+            print(entry['CacheKey'])

test/CAS/cache_key_compute.swift

@@ -50,32 +50,34 @@
 // RUN:   %target-swift-frontend -cache-compile-job %s -emit-module -c -emit-dependencies \
 // RUN:   -emit-tbd -emit-tbd-path %t/test.tbd -o %t/test.o -allow-unstable-cache-key-for-testing | %FileCheck %s --check-prefix=CHECK --check-prefix=PLUGIN
 
-// CHECK: test.o
+// CHECK: "Input": "{{.*}}{{/|\\}}cache_key_compute.swift"
-// CHECK-NEXT: "OutputKind": "object"
-// CHECK-NEXT: "Input"
 // CHECK-NEXT: "CacheKey"
 // PLUGIN-SAME: myfirst-llvmcas://
 
-// CHECK: test.swiftmodule
+// CHECK-NEXT: "Outputs": [
-// CHECK-NEXT: "OutputKind": "swiftmodule"
+// CHECK-NEXT: {
-// CHECK-NEXT: "Input"
+// CHECK-NEXT: "Kind": "object",
-// CHECK-NEXT: "CacheKey"
+// CHECK-NEXT: "Path":
-// PLUGIN-SAME: myfirst-llvmcas://
+// CHECK-SAME: test.o
+// CHECK-NEXT: },
+// CHECK-NEXT: {
+// CHECK-NEXT: "Kind": "swiftmodule",
+// CHECK-NEXT: "Path":
+// CHECK-SAME: test.swiftmodule
+// CHECK-NEXT: },
+// CHECK-NEXT: {
+// CHECK-NEXT: "Kind": "dependencies",
+// CHECK-NEXT: "Path":
+// CHECK-SAME: test.d
+// CHECK-NEXT: },
+// CHECK-NEXT: {
+// CHECK-NEXT: "Kind": "tbd",
+// CHECK-NEXT: "Path":
+// CHECK-SAME: test.tbd
+// CHECK-NEXT: },
+// CHECK-NEXT: {
+// CHECK-NEXT: "Kind": "cached-diagnostics",
+// CHECK-NEXT: "Path": "<cached-diagnostics>"
+// CHECK-NEXT: }
+// CHECK-NEXT: ]
 
-// CHECK: test.d
-// CHECK-NEXT: "OutputKind": "dependencies"
-// CHECK-NEXT: "Input"
-// CHECK-NEXT: "CacheKey"
-// PLUGIN-SAME: myfirst-llvmcas://
-
-// CHECK: test.tbd
-// CHECK-NEXT: "OutputKind": "tbd"
-// CHECK-NEXT: "Input"
-// CHECK-NEXT: "CacheKey"
-// PLUGIN-SAME: myfirst-llvmcas://
-
-// CHECK: <cached-diagnostics>
-// CHECK-NEXT: "OutputKind": "cached-diagnostics"
-// CHECK-NEXT: "Input": "<cached-diagnostics>"
-// CHECK-NEXT: "CacheKey"
-// PLUGIN-SAME: myfirst-llvmcas://