[Caching] Change swift cache key computation

Update swift cache key computation mechanism from one cache key per
output, to one cache key per primary input file (for all outputs that
associated with that input).

The new schema allows fewer cache lookups while still preserving most of
the flexibility for batch mode and incremental mode.

Author: cachemeifyoucan

include/swift-c/DependencyScan/DependencyScan.h

@@ -437,17 +437,6 @@ typedef struct swiftscan_cas_options_s *swiftscan_cas_options_t;
 /// ActionCache.
 typedef struct swiftscan_cas_s *swiftscan_cas_t;
 
-/// Enum types for output types for cache key computation.
-/// TODO: complete the list.
-typedef enum {
-  SWIFTSCAN_OUTPUT_TYPE_OBJECT = 0,
-  SWIFTSCAN_OUTPUT_TYPE_SWIFTMODULE = 1,
-  SWIFTSCAN_OUTPUT_TYPE_SWIFTINTERFACE = 2,
-  SWIFTSCAN_OUTPUT_TYPE_SWIFTPRIVATEINTERFACE = 3,
-  SWIFTSCAN_OUTPUT_TYPE_CLANG_MODULE = 4,
-  SWIFTSCAN_OUTPUT_TYPE_CLANG_PCH = 5
-} swiftscan_output_kind_t;
-
 /// Create a \c CASOptions for creating CAS inside scanner specified.
 SWIFTSCAN_PUBLIC swiftscan_cas_options_t swiftscan_cas_options_create(void);
 
@@ -489,13 +478,15 @@ SWIFTSCAN_PUBLIC swiftscan_string_ref_t
 swiftscan_cas_store(swiftscan_cas_t cas, uint8_t *data, unsigned size,
                     swiftscan_string_ref_t *error);
 
-/// Compute \c CacheKey for output of \c kind from the compiler invocation \c
-/// argc and \c argv with \c input. Return \c CacheKey as string.
+/// Compute \c CacheKey for the outputs of a primary input file from a compiler
+/// invocation with command-line \c argc and \c argv. 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_compute_cache_key(
-    swiftscan_cas_t cas, int argc, const char **argv, const char *input,
-    swiftscan_output_kind_t kind, swiftscan_string_ref_t *error);
+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);
 
 //===----------------------------------------------------------------------===//
 

include/swift/AST/DiagnosticsFrontend.def

@@ -494,7 +494,7 @@ ERROR(error_caching_no_cas_fs, none,
 ERROR(error_prefix_mapping, none, "cannot create scanner prefix mapping: '%0'", (StringRef))
 
 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))
+REMARK(output_cache_miss, none, "cache miss for input file '%0': key '%1'", (StringRef, StringRef))
 
 // CAS related diagnostics
 ERROR(error_invalid_cas_id, none, "invalid CASID '%0' (%1)", (StringRef, StringRef))

include/swift/Basic/FileTypes.h

@@ -58,6 +58,9 @@ bool isAfterLLVM(ID Id);
 /// These need to be passed to the Swift frontend
 bool isPartOfSwiftCompilation(ID Id);
 
+/// Returns true of the type of the output is produced from a diagnostic engine.
+bool isProducedFromDiagnostics(ID Id);
+
 static inline void forAllTypes(llvm::function_ref<void(file_types::ID)> fn) {
   for (uint8_t i = 0; i < static_cast<uint8_t>(TY_INVALID); ++i)
     fn(static_cast<ID>(i));

include/swift/Frontend/CASOutputBackends.h

@@ -0,0 +1,60 @@
+//===------CASOutputBackends.h-- ---------------------------------*-C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_FRONTEND_CASOUTPUTBACKENDS_H
+#define SWIFT_FRONTEND_CASOUTPUTBACKENDS_H
+
+#include "swift/Frontend/FrontendInputsAndOutputs.h"
+#include "swift/Frontend/FrontendOptions.h"
+#include "llvm/CAS/ActionCache.h"
+#include "llvm/CAS/ObjectStore.h"
+#include "llvm/Support/VirtualOutputBackends.h"
+#include "llvm/Support/VirtualOutputFile.h"
+
+namespace swift::cas {
+
+class SwiftCASOutputBackend : public llvm::vfs::OutputBackend {
+  void anchor() override;
+
+protected:
+  llvm::IntrusiveRefCntPtr<OutputBackend> cloneImpl() const override;
+
+  llvm::Expected<std::unique_ptr<llvm::vfs::OutputFileImpl>>
+  createFileImpl(llvm::StringRef ResolvedPath,
+                 llvm::Optional<llvm::vfs::OutputConfig> Config) override;
+
+  virtual llvm::Error storeImpl(llvm::StringRef Path, llvm::StringRef Bytes,
+                                llvm::StringRef CorrespondingInput,
+                                file_types::ID OutputKind);
+
+private:
+  file_types::ID getOutputFileType(llvm::StringRef Path) const;
+
+public:
+  SwiftCASOutputBackend(llvm::cas::ObjectStore &CAS,
+                        llvm::cas::ActionCache &Cache,
+                        llvm::cas::ObjectRef BaseKey,
+                        const FrontendInputsAndOutputs &InputsAndOutputs,
+                        FrontendOptions::ActionType Action);
+  ~SwiftCASOutputBackend();
+
+  llvm::Error storeCachedDiagnostics(llvm::StringRef InputFile,
+                                     llvm::StringRef Bytes);
+
+private:
+  class Implementation;
+  Implementation &Impl;
+};
+
+} // end namespace swift::cas
+
+#endif

include/swift/Frontend/CachingUtils.h

@@ -13,8 +13,10 @@
 #ifndef SWIFT_FRONTEND_CACHINGUTILS_H
 #define SWIFT_FRONTEND_CACHINGUTILS_H
 
+#include "swift/Frontend/CASOutputBackends.h"
 #include "swift/Frontend/CachedDiagnostics.h"
 #include "swift/Frontend/FrontendInputsAndOutputs.h"
+#include "swift/Frontend/FrontendOptions.h"
 #include "clang/CAS/CASOptions.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/CAS/ActionCache.h"
@@ -29,11 +31,12 @@ namespace swift {
 
 /// Create a swift caching output backend that stores the output from
 /// compiler into a CAS.
-llvm::IntrusiveRefCntPtr<llvm::vfs::OutputBackend>
+llvm::IntrusiveRefCntPtr<cas::SwiftCASOutputBackend>
 createSwiftCachingOutputBackend(
     llvm::cas::ObjectStore &CAS, llvm::cas::ActionCache &Cache,
     llvm::cas::ObjectRef BaseKey,
-    const FrontendInputsAndOutputs &InputsAndOutputs);
+    const FrontendInputsAndOutputs &InputsAndOutputs,
+    FrontendOptions::ActionType Action);
 
 /// Replay the output of the compilation from cache.
 /// Return true if outputs are replayed, false otherwise.
@@ -46,24 +49,35 @@ bool replayCachedCompilerOutputs(
 /// Load the cached compile result from cache.
 std::unique_ptr<llvm::MemoryBuffer> loadCachedCompileResultFromCacheKey(
     llvm::cas::ObjectStore &CAS, llvm::cas::ActionCache &Cache,
-    DiagnosticEngine &Diag, llvm::StringRef CacheKey,
+    DiagnosticEngine &Diag, llvm::StringRef CacheKey, file_types::ID Type,
     llvm::StringRef Filename = "");
 
-/// Store compiler output.
-llvm::Error storeCachedCompilerOutput(llvm::cas::ObjectStore &CAS,
-                                      llvm::cas::ActionCache &Cache,
-                                      StringRef Path, StringRef Bytes,
-                                      llvm::cas::ObjectRef BaseKey,
-                                      StringRef CorrespondingInput,
-                                      file_types::ID OutputKind);
-
 llvm::Expected<llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem>>
 createCASFileSystem(llvm::cas::ObjectStore &CAS, ArrayRef<std::string> FSRoots,
                     ArrayRef<std::string> IncludeTreeRoots);
 
 std::vector<std::string> remapPathsFromCommandLine(
     ArrayRef<std::string> Args,
     llvm::function_ref<std::string(StringRef)> RemapCallback);
-} // namespace swift
+
+namespace cas {
+class CachedResultLoader {
+public:
+  CachedResultLoader(llvm::cas::ObjectStore &CAS, llvm::cas::ActionCache &Cache,
+                     llvm::cas::ObjectRef CacheKey)
+      : CAS(CAS), Cache(Cache), CacheKey(CacheKey) {}
+
+  using CallbackTy =
+      llvm::function_ref<llvm::Error(file_types::ID, llvm::cas::ObjectRef)>;
+  // Replay the cached result, return false if a cache miss happened.
+  llvm::Expected<bool> replay(CallbackTy Callback);
+
+private:
+  llvm::cas::ObjectStore &CAS;
+  llvm::cas::ActionCache &Cache;
+  llvm::cas::ObjectRef CacheKey;
+};
+} // end namespace cas
+} // end namespace swift
 
 #endif

include/swift/Frontend/CompileJobCacheKey.h

@@ -41,8 +41,7 @@ createCompileJobBaseCacheKey(llvm::cas::ObjectStore &CAS,
 llvm::Expected<llvm::cas::ObjectRef>
 createCompileJobCacheKeyForOutput(llvm::cas::ObjectStore &CAS,
                                   llvm::cas::ObjectRef BaseKey,
-                                  StringRef ProducingInput,
-                                  file_types::ID OutputType);
+                                  StringRef ProducingInput);
 } // namespace swift
 
 #endif

include/swift/Frontend/CompileJobCacheResult.h

@@ -0,0 +1,114 @@
+//===------CompileJobCacheResult.h-- -----------------------------*-C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_FRONTEND_COMPILEJOBCACHERESULT_H
+#define SWIFT_FRONTEND_COMPILEJOBCACHERESULT_H
+
+#include "swift/Basic/FileTypes.h"
+#include "llvm/CAS/CASNodeSchema.h"
+#include "llvm/CAS/ObjectStore.h"
+
+namespace swift::cas {
+
+class CompileJobResultSchema;
+class CompileJobCacheResult : public llvm::cas::ObjectProxy {
+public:
+  /// A single output file or stream.
+  struct Output {
+    /// The CAS object for this output.
+    llvm::cas::ObjectRef Object;
+    /// The output kind.
+    file_types::ID Kind;
+
+    bool operator==(const Output &Other) const {
+      return Object == Other.Object && Kind == Other.Kind;
+    }
+  };
+
+  /// Retrieves each \c Output from this result.
+  llvm::Error
+  forEachOutput(llvm::function_ref<llvm::Error(Output)> Callback) const;
+
+  /// Loads all outputs concurrently and passes the resulting \c ObjectProxy
+  /// objects to \p Callback. If there was an error during loading then the
+  /// callback will not be invoked.
+  llvm::Error forEachLoadedOutput(
+      llvm::function_ref<llvm::Error(Output, std::optional<ObjectProxy>)>
+          Callback);
+
+  size_t getNumOutputs() const;
+
+  Output getOutput(size_t I) const;
+
+  /// Retrieves a specific output specified by \p Kind, if it exists.
+  llvm::Optional<Output> getOutput(file_types::ID Kind) const;
+
+  /// Print this result to \p OS.
+  llvm::Error print(llvm::raw_ostream &OS);
+
+  /// Helper to build a \c CompileJobCacheResult from individual outputs.
+  class Builder {
+  public:
+    Builder();
+    ~Builder();
+    /// Treat outputs added for \p Path as having the given \p Kind. Otherwise
+    /// they will have kind \c Unknown.
+    void addKindMap(file_types::ID Kind, StringRef Path);
+    /// Add an output with an explicit \p Kind.
+    void addOutput(file_types::ID Kind, llvm::cas::ObjectRef Object);
+    /// Add an output for the given \p Path. There must be a a kind map for it.
+    llvm::Error addOutput(StringRef Path, llvm::cas::ObjectRef Object);
+    /// Build a single \c ObjectRef representing the provided outputs. The
+    /// result can be used with \c CompileJobResultSchema to retrieve the
+    /// original outputs.
+    llvm::Expected<llvm::cas::ObjectRef> build(llvm::cas::ObjectStore &CAS);
+
+  private:
+    struct PrivateImpl;
+    PrivateImpl &Impl;
+  };
+
+private:
+  llvm::cas::ObjectRef getOutputObject(size_t I) const;
+  llvm::cas::ObjectRef getPathsListRef() const;
+  file_types::ID getOutputKind(size_t I) const;
+  llvm::Expected<llvm::cas::ObjectRef> getOutputPath(size_t I) const;
+
+private:
+  friend class CompileJobResultSchema;
+  CompileJobCacheResult(const ObjectProxy &);
+};
+
+class CompileJobResultSchema
+    : public llvm::RTTIExtends<CompileJobResultSchema, llvm::cas::NodeSchema> {
+public:
+  static char ID;
+
+  CompileJobResultSchema(llvm::cas::ObjectStore &CAS);
+
+  /// Attempt to load \p Ref as a \c CompileJobCacheResult if it matches the
+  /// schema.
+  llvm::Expected<CompileJobCacheResult> load(llvm::cas::ObjectRef Ref) const;
+
+  bool isRootNode(const CompileJobCacheResult::ObjectProxy &Node) const final;
+  bool isNode(const CompileJobCacheResult::ObjectProxy &Node) const final;
+
+  /// Get this schema's marker node.
+  llvm::cas::ObjectRef getKindRef() const { return KindRef; }
+
+private:
+  llvm::cas::ObjectRef KindRef;
+};
+
+} // namespace swift::cas
+
+#endif

include/swift/Frontend/Frontend.h

@@ -30,6 +30,7 @@
 #include "swift/Basic/LangOptions.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/ClangImporter/ClangImporter.h"
+#include "swift/Frontend/CASOutputBackends.h"
 #include "swift/Frontend/CachedDiagnostics.h"
 #include "swift/Frontend/DiagnosticVerifier.h"
 #include "swift/Frontend/FrontendOptions.h"
@@ -483,6 +484,10 @@ class CompilerInstance {
   /// Virtual OutputBackend.
   llvm::IntrusiveRefCntPtr<llvm::vfs::OutputBackend> OutputBackend = nullptr;
 
+  /// CAS OutputBackend.
+  llvm::IntrusiveRefCntPtr<swift::cas::SwiftCASOutputBackend> CASOutputBackend =
+      nullptr;
+
   /// The verification output backend.
   using HashBackendTy = llvm::vfs::HashingOutputBackend<llvm::BLAKE3>;
   llvm::IntrusiveRefCntPtr<HashBackendTy> HashBackend;
@@ -537,6 +542,10 @@ class CompilerInstance {
   llvm::vfs::OutputBackend &getOutputBackend() const {
     return *OutputBackend;
   }
+  swift::cas::SwiftCASOutputBackend &getCASOutputBackend() const {
+    return *CASOutputBackend;
+  }
+
   void
   setOutputBackend(llvm::IntrusiveRefCntPtr<llvm::vfs::OutputBackend> Backend) {
     OutputBackend = std::move(Backend);