[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.

Author: cachemeifyoucan

Diff for: include/swift/Frontend/FrontendInputsAndOutputs.h

@@ -157,6 +157,8 @@ class FrontendInputsAndOutputs {
   /// 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;

Diff for: 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,44 +284,41 @@ int SwiftCacheToolInvocation::printOutputKeys() {
 
   std::vector<OutputEntry> OutputKeys;
   bool hasError = false;
-  auto addOutputKey = [&](StringRef InputPath, file_types::ID OutputKind,
-                          StringRef OutputPath) {
+  auto addFromInputFile = [&](const InputFile &Input) {
+    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(),
-                    file_types::getTypeName(OutputKind).str(),
-                    CAS.getID(*OutputKey).toString()});
-  };
-  auto addFromInputFile = [&](const InputFile &Input) {
-    auto InputPath = Input.getFileName();
+        OutputEntry{InputPath, CAS.getID(*OutputKey).toString(), {}});
+    auto &Outputs = OutputKeys.back().Outputs;
     if (!Input.outputFilename().empty())
-      addOutputKey(InputPath,
-                   Invocation.getFrontendOptions()
-                       .InputsAndOutputs.getPrincipalOutputType(),
-                   Input.outputFilename());
+      Outputs.emplace_back(file_types::getTypeName(
+                               Invocation.getFrontendOptions()
+                                   .InputsAndOutputs.getPrincipalOutputType()),
+                           Input.outputFilename());
     Input.getPrimarySpecificPaths()
         .SupplementaryOutputs.forEachSetOutputAndType(
             [&](const std::string &File, file_types::ID ID) {
               // Dont print serialized diagnostics.
               if (file_types::isProducedFromDiagnostics(ID))
                 return;
-
-              addOutputKey(InputPath, ID, File);
+              Outputs.emplace_back(file_types::getTypeName(ID), File);
             });
   };
   llvm::for_each(
       Invocation.getFrontendOptions().InputsAndOutputs.getAllInputs(),
       addFromInputFile);
 
   // Add diagnostics file.
-  addOutputKey("<cached-diagnostics>", file_types::ID::TY_CachedDiagnostics,
-               "<cached-diagnostics>");
+  if (!OutputKeys.empty())
+    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);
+            });
+          }
+        });
       });
     }
   });

Diff for: 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(

Diff for: 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
-    // inputs.
+    // Store CachedDiagnostics in the CAS/Cache.
+    // 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,

Diff for: lib/Frontend/CachingUtils.cpp

@@ -209,19 +209,15 @@ bool replayCachedCompilerOutputs(
               Outputs.try_emplace(ID, File);
             });
 
-    // Nothing to replay.
-    if (Outputs.empty())
-      return;
+    // Add cached diagnostic entry for lookup. Output path doesn't matter here.
+    Outputs.try_emplace(file_types::ID::TY_CachedDiagnostics,
+                        "<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;
 

Diff for: 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;
 }

Diff for: test/CAS/Inputs/ExtractOutputKey.py

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

Diff for: 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-NEXT: "OutputKind": "object"
-// CHECK-NEXT: "Input"
+// CHECK: "Input": "{{.*}}{{/|\\}}cache_key_compute.swift"
 // CHECK-NEXT: "CacheKey"
 // PLUGIN-SAME: myfirst-llvmcas://
 
-// CHECK: test.swiftmodule
-// CHECK-NEXT: "OutputKind": "swiftmodule"
-// CHECK-NEXT: "Input"
-// CHECK-NEXT: "CacheKey"
-// PLUGIN-SAME: myfirst-llvmcas://
-
-// CHECK: test.d
-// CHECK-NEXT: "OutputKind": "dependencies"
-// CHECK-NEXT: "Input"
-// CHECK-NEXT: "CacheKey"
-// PLUGIN-SAME: myfirst-llvmcas://
+// CHECK-NEXT: "Outputs": [
+// CHECK-NEXT: {
+// CHECK-NEXT: "Kind": "object",
+// CHECK-NEXT: "Path":
+// 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.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://