Merge pull request #68946 from cachemeifyoucan/eng/PR-cache-key-rework

[Caching] Change swift cache key computation

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