[CAS] Teach swift-frontend to replay result from cache

Author: cachemeifyoucan

include/swift/AST/DiagnosticsFrontend.def

@@ -489,6 +489,9 @@ REMARK(matching_output_produced,none,
 ERROR(error_caching_no_cas_fs, none,
     "caching is enabled without -cas-fs option, input is not immutable", ())
 
+REMARK(replay_output, none, "replay output file '%0': key '%1'", (StringRef, StringRef))
+REMARK(output_cache_miss, none, "cache miss output file '%0': key '%1'", (StringRef, StringRef))
+
 // CAS related diagnostics
 ERROR(error_create_cas, none, "failed to create CAS '%0' (%1)", (StringRef, StringRef))
 ERROR(error_invalid_cas_id, none, "invalid CASID '%0' (%1)", (StringRef, StringRef))

include/swift/Frontend/CachingUtils.h

@@ -13,6 +13,7 @@
 #ifndef SWIFT_FRONTEND_CACHINGUTILS_H
 #define SWIFT_FRONTEND_CACHINGUTILS_H
 
+#include "swift/Frontend/CachedDiagnostics.h"
 #include "swift/Frontend/FrontendInputsAndOutputs.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/CAS/ActionCache.h"
@@ -31,6 +32,14 @@ createSwiftCachingOutputBackend(
     llvm::cas::ObjectRef BaseKey,
     const FrontendInputsAndOutputs &InputsAndOutputs);
 
+/// Replay the output of the compilation from cache.
+/// Return true if outputs are replayed, false otherwise.
+bool replayCachedCompilerOutputs(
+    llvm::cas::ObjectStore &CAS, llvm::cas::ActionCache &Cache,
+    llvm::cas::ObjectRef BaseKey, DiagnosticEngine &Diag,
+    const FrontendInputsAndOutputs &InputsAndOutputs,
+    CachingDiagnosticsProcessor &CDP);
+
 /// Load the cached compile result from cache.
 std::unique_ptr<llvm::MemoryBuffer> loadCachedCompileResultFromCacheKey(
     llvm::cas::ObjectStore &CAS, llvm::cas::ActionCache &Cache,

lib/Frontend/CachingUtils.cpp

@@ -133,6 +133,134 @@ createSwiftCachingOutputBackend(
                                                     InputsAndOutputs);
 }
 
+bool replayCachedCompilerOutputs(
+    ObjectStore &CAS, ActionCache &Cache, ObjectRef BaseKey,
+    DiagnosticEngine &Diag, const FrontendInputsAndOutputs &InputsAndOutputs,
+    CachingDiagnosticsProcessor &CDP) {
+  clang::cas::CompileJobResultSchema Schema(CAS);
+  bool CanReplayAllOutput = true;
+  struct OutputEntry {
+    std::string Path;
+    std::string Key;
+    llvm::cas::ObjectProxy Proxy;
+  };
+  SmallVector<OutputEntry> OutputProxies;
+
+  auto replayOutputFile = [&](StringRef InputName, file_types::ID OutputKind,
+                              StringRef OutputPath) -> Optional<OutputEntry> {
+    LLVM_DEBUG(llvm::dbgs()
+                   << "DEBUG: lookup output \'" << OutputPath << "\' type \'"
+                   << file_types::getTypeName(OutputKind) << "\' input \'"
+                   << InputName << "\n";);
+
+    auto OutputKey =
+        createCompileJobCacheKeyForOutput(CAS, BaseKey, InputName, OutputKind);
+    if (!OutputKey) {
+      Diag.diagnose(SourceLoc(), diag::error_cas,
+                    toString(OutputKey.takeError()));
+      return None;
+    }
+    auto OutputKeyID = CAS.getID(*OutputKey);
+    auto Lookup = Cache.get(OutputKeyID);
+    if (!Lookup) {
+      Diag.diagnose(SourceLoc(), diag::error_cas, toString(Lookup.takeError()));
+      return None;
+    }
+    if (!*Lookup) {
+      Diag.diagnose(SourceLoc(), diag::output_cache_miss, OutputPath,
+                    OutputKeyID.toString());
+      return None;
+    }
+    auto OutputRef = CAS.getReference(**Lookup);
+    if (!OutputRef) {
+      return None;
+    }
+    auto Result = Schema.load(*OutputRef);
+    if (!Result) {
+      Diag.diagnose(SourceLoc(), diag::error_cas, toString(Result.takeError()));
+      return None;
+    }
+    auto MainOutput = Result->getOutput(
+        clang::cas::CompileJobCacheResult::OutputKind::MainOutput);
+    if (!MainOutput) {
+      return None;
+    }
+    auto LoadedResult = CAS.getProxy(MainOutput->Object);
+    if (!LoadedResult) {
+      Diag.diagnose(SourceLoc(), diag::error_cas,
+                    toString(LoadedResult.takeError()));
+      return None;
+    }
+
+    return OutputEntry{OutputPath.str(), OutputKeyID.toString(), *LoadedResult};
+  };
+
+  auto replayOutputFromInput = [&](const InputFile &Input) {
+    auto InputPath = Input.getFileName();
+    if (!Input.outputFilename().empty()) {
+      if (auto Result = replayOutputFile(
+              InputPath, InputsAndOutputs.getPrincipalOutputType(),
+              Input.outputFilename()))
+        OutputProxies.emplace_back(*Result);
+      else
+        CanReplayAllOutput = false;
+    }
+
+    Input.getPrimarySpecificPaths()
+        .SupplementaryOutputs.forEachSetOutputAndType(
+            [&](const std::string &File, file_types::ID ID) {
+              if (ID == file_types::ID::TY_SerializedDiagnostics)
+                return;
+
+              if (auto Result = replayOutputFile(InputPath, ID, File))
+                OutputProxies.emplace_back(*Result);
+              else
+                CanReplayAllOutput = false;
+            });
+  };
+
+  llvm::for_each(InputsAndOutputs.getAllInputs(), replayOutputFromInput);
+
+  auto DiagnosticsOutput = replayOutputFile(
+      "<cached-diagnostics>", file_types::ID::TY_CachedDiagnostics,
+      "<cached-diagnostics>");
+  if (!DiagnosticsOutput)
+    CanReplayAllOutput = false;
+
+  if (!CanReplayAllOutput)
+    return false;
+
+  // Replay Diagnostics first so the output failures comes after.
+  // Also if the diagnostics replay failed, proceed to re-compile.
+  if (auto E = CDP.replayCachedDiagnostics(
+          DiagnosticsOutput->Proxy.getData())) {
+    Diag.diagnose(SourceLoc(), diag::error_replay_cached_diag,
+                  toString(std::move(E)));
+    return false;
+  }
+
+  // Replay the result only when everything is resolved.
+  // Use on disk output backend directly here to write to disk.
+  llvm::vfs::OnDiskOutputBackend Backend;
+  for (auto &Output : OutputProxies) {
+    auto File = Backend.createFile(Output.Path);
+    if (!File) {
+      Diag.diagnose(SourceLoc(), diag::error_opening_output, Output.Path,
+                    toString(File.takeError()));
+      continue;
+    }
+    *File << Output.Proxy.getData();
+    if (auto E = File->keep()) {
+      Diag.diagnose(SourceLoc(), diag::error_closing_output, Output.Path,
+                    toString(std::move(E)));
+      continue;
+    }
+    Diag.diagnose(SourceLoc(), diag::replay_output, Output.Path, Output.Key);
+  }
+
+  return true;
+}
+
 static Expected<std::unique_ptr<llvm::MemoryBuffer>>
 loadCachedCompileResultFromCacheKeyImpl(ObjectStore &CAS, ActionCache &Cache,
                                         StringRef CacheKey,

lib/Frontend/CompileJobCacheKey.cpp

@@ -48,6 +48,13 @@ llvm::Expected<llvm::cas::ObjectRef> swift::createCompileJobBaseCacheKey(
       SkipNext = true;
       continue;
     }
+    // FIXME: Use a heuristic to remove all the flags that affect output paths.
+    // Those should not affect compile cache key.
+    if (Arg.startswith("-emit-")) {
+      if (Arg.endswith("-path"))
+        SkipNext = true;
+      continue;
+    }
     CommandLine.append(Arg);
     CommandLine.push_back(0);
   }

lib/FrontendTool/FrontendTool.cpp

@@ -51,6 +51,9 @@
 #include "swift/ConstExtract/ConstExtract.h"
 #include "swift/DependencyScan/ScanDependencies.h"
 #include "swift/Frontend/AccumulatingDiagnosticConsumer.h"
+#include "swift/Frontend/CachedDiagnostics.h"
+#include "swift/Frontend/CachingUtils.h"
+#include "swift/Frontend/CompileJobCacheKey.h"
 #include "swift/Frontend/Frontend.h"
 #include "swift/Frontend/ModuleInterfaceLoader.h"
 #include "swift/Frontend/ModuleInterfaceSupport.h"
@@ -1380,6 +1383,35 @@ static bool performAction(CompilerInstance &Instance,
   return Instance.getASTContext().hadError();
 }
 
+/// Try replay the compiler result from cache.
+///
+/// Return true if all the outputs are fetched from cache. Otherwise, return
+/// false and will not replay any output.
+static bool tryReplayCompilerResults(CompilerInstance &Instance) {
+  if (!Instance.supportCaching())
+    return false;
+
+  assert(Instance.getCompilerBaseKey() &&
+         "Instance is not setup correctly for replay");
+
+  auto *CDP = Instance.getCachingDiagnosticsProcessor();
+  assert(CDP && "CachingDiagnosticsProcessor needs to be setup for replay");
+
+  // Don't capture diagnostics from replay.
+  CDP->endDiagnosticCapture();
+
+  bool replayed = replayCachedCompilerOutputs(
+      Instance.getObjectStore(), Instance.getActionCache(),
+      *Instance.getCompilerBaseKey(), Instance.getDiags(),
+      Instance.getInvocation().getFrontendOptions().InputsAndOutputs, *CDP);
+
+  // If we didn't replay successfully, re-start capture.
+  if (!replayed)
+    CDP->startDiagnosticCapture();
+
+  return replayed;
+}
+
 /// Performs the compile requested by the user.
 /// \param Instance Will be reset after performIRGeneration when the verifier
 ///                 mode is NoVerify and there were no errors.
@@ -1391,6 +1423,9 @@ static bool performCompile(CompilerInstance &Instance,
   const auto &opts = Invocation.getFrontendOptions();
   const FrontendOptions::ActionType Action = opts.RequestedAction;
 
+  if (tryReplayCompilerResults(Instance))
+    return false;
+
   // To compile LLVM IR, just pass it off unmodified.
   if (opts.InputsAndOutputs.shouldTreatAsLLVM())
     return compileLLVMIR(Instance);

test/CAS/cache_replay.swift

@@ -0,0 +1,19 @@
+// RUN: %empty-directory(%t)
+
+/// Run the command first time, expect cache miss.
+/// FIXME: This command doesn't use `-cas-fs` so it is not a good cache entry. It is currently allowed so it is easier to write tests.
+// RUN: %target-swift-frontend -enable-cas %s -emit-module -emit-module-path %t/Test.swiftmodule -c -emit-dependencies \
+// RUN:  -module-name Test -o %t/test.o -cas-path %t/cas -allow-unstable-cache-key-for-testing 2>&1 | %FileCheck --check-prefix=CACHE-MISS %s
+
+/// Expect cache hit for second time.
+// RUN: %target-swift-frontend -enable-cas %s -emit-module -emit-module-path %t/Test.swiftmodule -c -emit-dependencies \
+// RUN:  -module-name Test -o %t/test.o -cas-path %t/cas -allow-unstable-cache-key-for-testing 2>&1 | %FileCheck --check-prefix=CACHE-HIT %s
+
+/// Expect cache hit a subset of outputs.
+// RUN: %target-swift-frontend -enable-cas %s -emit-module -emit-module-path %t/Test.swiftmodule -c \
+// RUN:  -module-name Test -o %t/test.o -cas-path %t/cas -allow-unstable-cache-key-for-testing 2>&1 | %FileCheck --check-prefix=CACHE-HIT %s
+
+// CACHE-MISS: remark: cache miss output file
+// CACHE-HIT: remark: replay output file
+
+func testFunc() {}