[CAS] Add ActionCache to LLVMCAS Library (#114097)

ActionCache is used to store a mapping from CASID to CASID. The current
implementation of the ActionCache can only be used to associate the
key/value from the same hash context.

ActionCache has two operations: `put` to store the key/value and `get`
to
lookup the key/value mapping. ActionCache uses the same TrieRawHashMap
data structure to store the mapping, where is CASID of the key is the
hash to index the map.

While CASIDs for key/value are often associcate with actual CAS
ObjectStore, it doesn't provide the guarantee of the existence of such
object in any ObjectStore.

Author: cachemeifyoucan

llvm/docs/ContentAddressableStorage.md

@@ -83,6 +83,19 @@ Expected<ObjectProxy> Loaded = ObjectStore.getProxy(C);
 It also provides APIs to convert between `ObjectRef`, `ObjectProxy` and
 `CASID`.
 
+### ActionCache
+
+`ActionCache` is a key value storage can be used to associate two CASIDs. It is
+usually used with an `ObjectStore` to map an input CASObject to an output CASObject
+with their CASIDs.
+
+`ActionCache` has APIs like following:
+
+```
+CASID A, B;
+Error E = ActionCache.put(A, B);
+Expected<std::optional<CASID>> Result = ActionCache.get(A);
+```
 
 
 ## CAS Library Implementation Guide

llvm/include/llvm/CAS/ActionCache.h

@@ -0,0 +1,102 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file contains the declaration of the ActionCache class, which is the
+/// base class for ActionCache implementations.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CAS_ACTIONCACHE_H
+#define LLVM_CAS_ACTIONCACHE_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/CAS/CASID.h"
+#include "llvm/CAS/CASReference.h"
+#include "llvm/Support/Error.h"
+
+namespace llvm::cas {
+
+class ObjectStore;
+class CASID;
+class ObjectProxy;
+
+/// A key for caching an operation.
+/// It is implemented as a bag of bytes and provides a convenient constructor
+/// for CAS types.
+class CacheKey {
+public:
+  StringRef getKey() const { return Key; }
+
+  CacheKey(const CASID &ID);
+  CacheKey(const ObjectProxy &Proxy);
+  CacheKey(const ObjectStore &CAS, const ObjectRef &Ref);
+
+private:
+  std::string Key;
+};
+
+/// A cache from a key (that describes an action) to the result of performing
+/// that action.
+///
+/// Actions are expected to be pure. Storing mappings from one action to
+/// multiple results will result in error (cache poisoning).
+class ActionCache {
+  virtual void anchor();
+
+public:
+  /// Get a previously computed result for \p ActionKey.
+  ///
+  /// \param CanBeDistributed is a hint to the underlying implementation that if
+  /// it is true, the lookup is profitable to be done on a distributed caching
+  /// level, not just locally. The implementation is free to ignore this flag.
+  Expected<std::optional<CASID>> get(const CacheKey &ActionKey,
+                                     bool CanBeDistributed = false) const {
+    return getImpl(arrayRefFromStringRef(ActionKey.getKey()), CanBeDistributed);
+  }
+
+  /// Cache \p Result for the \p ActionKey computation.
+  ///
+  /// \param CanBeDistributed is a hint to the underlying implementation that if
+  /// it is true, the association is profitable to be done on a distributed
+  /// caching level, not just locally. The implementation is free to ignore this
+  /// flag.
+  Error put(const CacheKey &ActionKey, const CASID &Result,
+            bool CanBeDistributed = false) {
+    assert(Result.getContext().getHashSchemaIdentifier() ==
+               getContext().getHashSchemaIdentifier() &&
+           "Hash schema mismatch");
+    return putImpl(arrayRefFromStringRef(ActionKey.getKey()), Result,
+                   CanBeDistributed);
+  }
+
+  virtual ~ActionCache() = default;
+
+protected:
+  // Implementation detail for \p get method.
+  virtual Expected<std::optional<CASID>>
+  getImpl(ArrayRef<uint8_t> ResolvedKey, bool CanBeDistributed) const = 0;
+
+  // Implementation details for \p put method.
+  virtual Error putImpl(ArrayRef<uint8_t> ResolvedKey, const CASID &Result,
+                        bool CanBeDistributed) = 0;
+
+  ActionCache(const CASContext &Context) : Context(Context) {}
+
+  const CASContext &getContext() const { return Context; }
+
+private:
+  const CASContext &Context;
+};
+
+/// Create an action cache in memory.
+std::unique_ptr<ActionCache> createInMemoryActionCache();
+
+} // end namespace llvm::cas
+
+#endif // LLVM_CAS_ACTIONCACHE_H

llvm/lib/CAS/ActionCache.cpp

@@ -0,0 +1,22 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CAS/ActionCache.h"
+#include "llvm/CAS/CASID.h"
+#include "llvm/CAS/ObjectStore.h"
+
+using namespace llvm;
+using namespace llvm::cas;
+
+void ActionCache::anchor() {}
+
+CacheKey::CacheKey(const CASID &ID) : Key(toStringRef(ID.getHash()).str()) {}
+CacheKey::CacheKey(const ObjectProxy &Proxy)
+    : CacheKey(Proxy.getCAS(), Proxy.getRef()) {}
+CacheKey::CacheKey(const ObjectStore &CAS, const ObjectRef &Ref)
+    : Key(toStringRef(CAS.getID(Ref).getHash())) {}

llvm/lib/CAS/ActionCaches.cpp

@@ -0,0 +1,101 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file This file implements the underlying ActionCache implementations.
+///
+//===----------------------------------------------------------------------===//
+
+#include "BuiltinCAS.h"
+#include "llvm/ADT/TrieRawHashMap.h"
+#include "llvm/CAS/ActionCache.h"
+#include "llvm/Support/BLAKE3.h"
+
+#define DEBUG_TYPE "cas-action-caches"
+
+using namespace llvm;
+using namespace llvm::cas;
+
+namespace {
+
+using HasherT = BLAKE3;
+using HashType = decltype(HasherT::hash(std::declval<ArrayRef<uint8_t> &>()));
+
+template <size_t Size> class CacheEntry {
+public:
+  CacheEntry() = default;
+  CacheEntry(ArrayRef<uint8_t> Hash) { llvm::copy(Hash, Value.data()); }
+  CacheEntry(const CacheEntry &Entry) { llvm::copy(Entry.Value, Value.data()); }
+  ArrayRef<uint8_t> getValue() const { return Value; }
+
+private:
+  std::array<uint8_t, Size> Value;
+};
+
+/// Builtin InMemory ActionCache that stores the mapping in memory.
+class InMemoryActionCache final : public ActionCache {
+public:
+  InMemoryActionCache()
+      : ActionCache(builtin::BuiltinCASContext::getDefaultContext()) {}
+
+  Error putImpl(ArrayRef<uint8_t> ActionKey, const CASID &Result,
+                bool CanBeDistributed) final;
+  Expected<std::optional<CASID>> getImpl(ArrayRef<uint8_t> ActionKey,
+                                         bool CanBeDistributed) const final;
+
+private:
+  using DataT = CacheEntry<sizeof(HashType)>;
+  using InMemoryCacheT = ThreadSafeTrieRawHashMap<DataT, sizeof(HashType)>;
+
+  InMemoryCacheT Cache;
+};
+} // end namespace
+
+static Error createResultCachePoisonedError(ArrayRef<uint8_t> KeyHash,
+                                            const CASContext &Context,
+                                            CASID Output,
+                                            ArrayRef<uint8_t> ExistingOutput) {
+  std::string Existing =
+      CASID::create(&Context, toStringRef(ExistingOutput)).toString();
+  SmallString<64> Key;
+  toHex(KeyHash, /*LowerCase=*/true, Key);
+  return createStringError(std::make_error_code(std::errc::invalid_argument),
+                           "cache poisoned for '" + Key + "' (new='" +
+                               Output.toString() + "' vs. existing '" +
+                               Existing + "')");
+}
+
+Expected<std::optional<CASID>>
+InMemoryActionCache::getImpl(ArrayRef<uint8_t> Key,
+                             bool /*CanBeDistributed*/) const {
+  auto Result = Cache.find(Key);
+  if (!Result)
+    return std::nullopt;
+  return CASID::create(&getContext(), toStringRef(Result->Data.getValue()));
+}
+
+Error InMemoryActionCache::putImpl(ArrayRef<uint8_t> Key, const CASID &Result,
+                                   bool /*CanBeDistributed*/) {
+  DataT Expected(Result.getHash());
+  const InMemoryCacheT::value_type &Cached = *Cache.insertLazy(
+      Key, [&](auto ValueConstructor) { ValueConstructor.emplace(Expected); });
+
+  const DataT &Observed = Cached.Data;
+  if (Expected.getValue() == Observed.getValue())
+    return Error::success();
+
+  return createResultCachePoisonedError(Key, getContext(), Result,
+                                        Observed.getValue());
+}
+
+namespace llvm::cas {
+
+std::unique_ptr<ActionCache> createInMemoryActionCache() {
+  return std::make_unique<InMemoryActionCache>();
+}
+
+} // namespace llvm::cas

llvm/lib/CAS/CMakeLists.txt

@@ -1,4 +1,6 @@
 add_llvm_component_library(LLVMCAS
+  ActionCache.cpp
+  ActionCaches.cpp
   BuiltinCAS.cpp
   InMemoryCAS.cpp
   ObjectStore.cpp

llvm/unittests/CAS/ActionCacheTest.cpp

@@ -0,0 +1,73 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file implements the tests for ActionCaches.
+///
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CAS/ActionCache.h"
+#include "CASTestConfig.h"
+#include "llvm/CAS/ObjectStore.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::cas;
+
+TEST_P(CASTest, ActionCacheHit) {
+  std::shared_ptr<ObjectStore> CAS = createObjectStore();
+  std::unique_ptr<ActionCache> Cache = createActionCache();
+
+  std::optional<ObjectProxy> ID;
+  ASSERT_THAT_ERROR(CAS->createProxy({}, "1").moveInto(ID), Succeeded());
+  std::optional<CASID> ResultID;
+  ASSERT_THAT_ERROR(Cache->put(*ID, *ID), Succeeded());
+  ASSERT_THAT_ERROR(Cache->get(*ID).moveInto(ResultID), Succeeded());
+  ASSERT_TRUE(ResultID);
+  std::optional<ObjectRef> Result = CAS->getReference(*ResultID);
+  ASSERT_TRUE(Result);
+  ASSERT_EQ(*ID, *Result);
+}
+
+TEST_P(CASTest, ActionCacheMiss) {
+  std::shared_ptr<ObjectStore> CAS = createObjectStore();
+  std::unique_ptr<ActionCache> Cache = createActionCache();
+
+  std::optional<ObjectProxy> ID1, ID2;
+  ASSERT_THAT_ERROR(CAS->createProxy({}, "1").moveInto(ID1), Succeeded());
+  ASSERT_THAT_ERROR(CAS->createProxy({}, "2").moveInto(ID2), Succeeded());
+  ASSERT_THAT_ERROR(Cache->put(*ID1, *ID2), Succeeded());
+  // This is a cache miss for looking up a key doesn't exist.
+  std::optional<CASID> Result1;
+  ASSERT_THAT_ERROR(Cache->get(*ID2).moveInto(Result1), Succeeded());
+  ASSERT_FALSE(Result1);
+
+  ASSERT_THAT_ERROR(Cache->put(*ID2, *ID1), Succeeded());
+  // Cache hit after adding the value.
+  std::optional<CASID> Result2;
+  ASSERT_THAT_ERROR(Cache->get(*ID2).moveInto(Result2), Succeeded());
+  ASSERT_TRUE(Result2);
+  std::optional<ObjectRef> Ref = CAS->getReference(*Result2);
+  ASSERT_TRUE(Ref);
+  ASSERT_EQ(*ID1, *Ref);
+}
+
+TEST_P(CASTest, ActionCacheRewrite) {
+  std::shared_ptr<ObjectStore> CAS = createObjectStore();
+  std::unique_ptr<ActionCache> Cache = createActionCache();
+
+  std::optional<ObjectProxy> ID1, ID2;
+  ASSERT_THAT_ERROR(CAS->createProxy({}, "1").moveInto(ID1), Succeeded());
+  ASSERT_THAT_ERROR(CAS->createProxy({}, "2").moveInto(ID2), Succeeded());
+  ASSERT_THAT_ERROR(Cache->put(*ID1, *ID1), Succeeded());
+  // Writing to the same key with different value is error.
+  ASSERT_THAT_ERROR(Cache->put(*ID1, *ID2), Failed());
+  // Writing the same value multiple times to the same key is fine.
+  ASSERT_THAT_ERROR(Cache->put(*ID1, *ID1), Succeeded());
+}

llvm/unittests/CAS/CASTestConfig.cpp

@@ -13,9 +13,8 @@
 using namespace llvm;
 using namespace llvm::cas;
 
-CASTestingEnv createInMemory(int I) {
-  std::unique_ptr<ObjectStore> CAS = createInMemoryCAS();
-  return CASTestingEnv{std::move(CAS)};
+static CASTestingEnv createInMemory(int I) {
+  return CASTestingEnv{createInMemoryCAS(), createInMemoryActionCache()};
 }
 
 INSTANTIATE_TEST_SUITE_P(InMemoryCAS, CASTest,

llvm/unittests/CAS/CASTestConfig.h

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CAS/ActionCache.h"
 #include "llvm/CAS/ObjectStore.h"
 #include "gtest/gtest.h"
 
@@ -14,6 +15,7 @@
 
 struct CASTestingEnv {
   std::unique_ptr<llvm::cas::ObjectStore> CAS;
+  std::unique_ptr<llvm::cas::ActionCache> Cache;
 };
 
 class CASTest
@@ -25,6 +27,10 @@ class CASTest
     auto TD = GetParam()(++(*NextCASIndex));
     return std::move(TD.CAS);
   }
+  std::unique_ptr<llvm::cas::ActionCache> createActionCache() {
+    auto TD = GetParam()(++(*NextCASIndex));
+    return std::move(TD.Cache);
+  }
   void SetUp() { NextCASIndex = 0; }
   void TearDown() { NextCASIndex = std::nullopt; }
 };

llvm/unittests/CAS/CMakeLists.txt

@@ -5,6 +5,7 @@ set(LLVM_LINK_COMPONENTS
   )
 
 add_llvm_unittest(CASTests
+  ActionCacheTest.cpp
   CASTestConfig.cpp
   ObjectStoreTest.cpp
   )