Added initial changes to support multiple Merkle trees in the storage.

kvbc/cmf/categorized_kvbc_msgs.cmf

@@ -199,6 +199,7 @@ Msg NibblePath 5001 {
 
 Msg BatchedInternalNodeKey 5002 {
     # The version of this key is the *tree* version, not block version
+    string custom_prefix
     uint64 version
     NibblePath path
 }

kvbc/include/categorization/block_merkle_category.h

@@ -202,7 +202,7 @@ class BlockMerkleCategory {
     Reader(const storage::rocksdb::NativeClient& db) : db_{db} {}
 
     // Return the latest root node in the system.
-    sparse_merkle::BatchedInternalNode get_latest_root() const override;
+    sparse_merkle::BatchedInternalNode get_latest_root(std::string custom_prefix = "") const override;
 
     // Retrieve a BatchedInternalNode given an InternalNodeKey.
     //

kvbc/include/merkle_tree_db_adapter.h

@@ -219,7 +219,7 @@ class DBAdapter : public IDbAdapter {
     Reader(const DBAdapter &adapter) : adapter_{adapter} {}
 
     // Return the latest root node in the system.
-    sparse_merkle::BatchedInternalNode get_latest_root() const override;
+    sparse_merkle::BatchedInternalNode get_latest_root(std::string custom_prefix = "") const override;
 
     // Retrieve a BatchedInternalNode given an InternalNodeKey.
     //

kvbc/include/merkle_tree_serialization.h

@@ -88,6 +88,24 @@ inline void serializeImp(const sparse_merkle::NibblePath &path, std::string &out
   out.append(std::cbegin(path.data()), std::cend(path.data()));
 }
 
+inline void serializeImp(const sparse_merkle::CustomPrefix &custom_prefix, std::string &out) {
+  ConcordAssert(custom_prefix.size() < std::numeric_limits<uint8_t>::max());
+  static const auto SIZE_PREFIX_BYTES = sparse_merkle::CustomPrefix::SIZE_PREFIX_BYTES;
+  using SIZE_PREFIX_TYPE = sparse_merkle::CustomPrefix::SIZE_PREFIX_TYPE;
+  out.append(SIZE_PREFIX_BYTES, static_cast<SIZE_PREFIX_TYPE>(custom_prefix.size()));
+  out.append(std::cbegin(custom_prefix.value()), std::cend(custom_prefix.value()));
+}
+
+inline void serializeImp(const sparse_merkle::PaddedCustomPrefix &custom_prefix, std::string &out) {
+  ConcordAssert(custom_prefix.size() < std::numeric_limits<uint8_t>::max());
+  static const auto SIZE_PREFIX_BYTES = sparse_merkle::PaddedCustomPrefix::SIZE_PREFIX_BYTES;
+  using SIZE_PREFIX_TYPE = sparse_merkle::PaddedCustomPrefix::SIZE_PREFIX_TYPE;
+  out.append(SIZE_PREFIX_BYTES, static_cast<SIZE_PREFIX_TYPE>(custom_prefix.size()));
+  for (auto i = 0u; i < sparse_merkle::PaddedCustomPrefix::SIZE_IN_BYTES; i++) {
+    out.append(1, static_cast<uint8_t>(custom_prefix.get(i)));
+  }
+}
+
 inline std::string serializeImp(const sparse_merkle::NibblePath &path) {
   auto out = std::string{};
   serializeImp(path, out);
@@ -105,6 +123,7 @@ inline std::string serializeImp(const sparse_merkle::Hash &hash) {
 }
 
 inline void serializeImp(const sparse_merkle::InternalNodeKey &key, std::string &out) {
+  serializeImp(key.customPrefix(), out);
   serializeImp(key.version().value(), out);
   serializeImp(key.path(), out);
 }
@@ -118,6 +137,7 @@ inline std::string serializeImp(const sparse_merkle::InternalNodeKey &key) {
 inline void serializeImp(const sparse_merkle::LeafKey &key, std::string &out) {
   serializeImp(key.hash(), out);
   serializeImp(key.version().value(), out);
+  serializeImp(key.customPrefix(), out);
 }
 
 inline std::string serializeImp(const sparse_merkle::LeafKey &key) {
@@ -310,28 +330,65 @@ inline std::vector<std::uint8_t> deserialize<std::vector<std::uint8_t>>(const co
 
 template <>
 inline sparse_merkle::NibblePath deserialize<sparse_merkle::NibblePath>(const concordUtils::Sliver &buf) {
-  ConcordAssert(buf.length() >= sizeof(std::uint8_t));
+  static const int nibblePathSizePrefixInBytes = sizeof(std::uint8_t);
+  ConcordAssert(buf.length() >= nibblePathSizePrefixInBytes);
   // NOLINTNEXTLINE(bugprone-signed-char-misuse)
-  const auto nibbleCount = static_cast<std::size_t>(buf[0]);
-  const auto path = deserialize<std::vector<std::uint8_t>>(
-      concordUtils::Sliver{buf, sizeof(std::uint8_t), buf.length() - sizeof(std::uint8_t)});
+  size_t nibbleCount = static_cast<std::uint8_t>(buf[0]);
+  size_t pathBytesCount = sparse_merkle::NibblePath::pathLengthInBytes(nibbleCount);
+  std::vector<std::uint8_t> path;
+  if (pathBytesCount > 0) {
+    path =
+        deserialize<std::vector<std::uint8_t>>(concordUtils::Sliver{buf, nibblePathSizePrefixInBytes, pathBytesCount});
+  }
   return sparse_merkle::NibblePath{nibbleCount, path};
 }
 
+template <>
+inline sparse_merkle::CustomPrefix deserialize<sparse_merkle::CustomPrefix>(const concordUtils::Sliver &buf) {
+  static const auto SIZE_PREFIX_BYTES = sparse_merkle::CustomPrefix::SIZE_PREFIX_BYTES;
+  using SIZE_PREFIX_TYPE = sparse_merkle::CustomPrefix::SIZE_PREFIX_TYPE;
+  ConcordAssert(buf.length() >= SIZE_PREFIX_BYTES);
+  // NOLINTNEXTLINE(bugprone-signed-char-misuse)
+  const auto prefixSize = static_cast<SIZE_PREFIX_TYPE>(buf[0]);
+  ConcordAssert(buf.length() > prefixSize);
+  sparse_merkle::CustomPrefix prefix(prefixSize);
+  for (auto i = SIZE_PREFIX_BYTES; i <= prefixSize; ++i) {
+    prefix.set(i - SIZE_PREFIX_BYTES, buf[i]);
+  }
+  return prefix;
+}
+
+template <>
+inline sparse_merkle::PaddedCustomPrefix deserialize<sparse_merkle::PaddedCustomPrefix>(
+    const concordUtils::Sliver &buf) {
+  return sparse_merkle::PaddedCustomPrefix(deserialize<sparse_merkle::CustomPrefix>(buf));
+}
+
 template <>
 inline sparse_merkle::LeafKey deserialize<sparse_merkle::LeafKey>(const concordUtils::Sliver &buf) {
   ConcordAssert(buf.length() >= sparse_merkle::LeafKey::SIZE_IN_BYTES);
-  return sparse_merkle::LeafKey{deserialize<sparse_merkle::Hash>(buf),
-                                deserialize<sparse_merkle::Version>(concordUtils::Sliver{
-                                    buf, sparse_merkle::Hash::SIZE_IN_BYTES, sparse_merkle::Version::SIZE_IN_BYTES})};
+  return sparse_merkle::LeafKey{
+      deserialize<sparse_merkle::Hash>(buf),
+      deserialize<sparse_merkle::Version>(
+          concordUtils::Sliver{buf, sparse_merkle::Hash::SIZE_IN_BYTES, sparse_merkle::Version::SIZE_IN_BYTES}),
+      deserialize<sparse_merkle::PaddedCustomPrefix>(concordUtils::Sliver{
+          buf,
+          sparse_merkle::Hash::SIZE_IN_BYTES + sparse_merkle::Version::SIZE_IN_BYTES,
+          buf.length() - sparse_merkle::Hash::SIZE_IN_BYTES - sparse_merkle::Version::SIZE_IN_BYTES})};
 }
 
 template <>
 inline sparse_merkle::InternalNodeKey deserialize<sparse_merkle::InternalNodeKey>(const concordUtils::Sliver &buf) {
-  return sparse_merkle::InternalNodeKey{
-      deserialize<sparse_merkle::Version>(buf),
-      deserialize<sparse_merkle::NibblePath>(concordUtils::Sliver{
-          buf, sparse_merkle::Version::SIZE_IN_BYTES, buf.length() - sparse_merkle::Version::SIZE_IN_BYTES})};
+  static const auto SIZE_PREFIX_BYTES = sparse_merkle::CustomPrefix::SIZE_PREFIX_BYTES;  // 1 byte for the address size
+  auto prefix = deserialize<sparse_merkle::CustomPrefix>(buf);
+  auto version = deserialize<sparse_merkle::Version>(
+      concordUtils::Sliver{buf, SIZE_PREFIX_BYTES + prefix.size(), buf.length() - SIZE_PREFIX_BYTES - prefix.size()});
+  auto nibblePath = deserialize<sparse_merkle::NibblePath>(
+      concordUtils::Sliver{buf,
+                           SIZE_PREFIX_BYTES + prefix.size() + sparse_merkle::Version::SIZE_IN_BYTES,
+                           buf.length() - SIZE_PREFIX_BYTES - prefix.size() - sparse_merkle::Version::SIZE_IN_BYTES});
+
+  return sparse_merkle::InternalNodeKey{std::move(prefix), std::move(version), std::move(nibblePath)};
 }
 
 template <>

kvbc/include/sparse_merkle/base_types.h

@@ -53,6 +53,57 @@ class Version {
   Type value_ = 0;
 };
 
+// A type safe custom prefix wrapper
+class CustomPrefix {
+ public:
+  using Type = std::string;
+  using SIZE_PREFIX_TYPE = std::uint8_t;
+  static constexpr auto SIZE_IN_BYTES = 20 * sizeof(std::uint8_t);
+  static constexpr auto SIZE_PREFIX_BYTES = sizeof(std::uint8_t);
+  static constexpr auto TOTAL_SIZE = SIZE_PREFIX_BYTES + SIZE_IN_BYTES;
+
+ public:
+  CustomPrefix() = default;
+  CustomPrefix(size_t size) {
+    ConcordAssert(size <= SIZE_IN_BYTES);
+    prefix_.assign('\0', size);
+  }
+  CustomPrefix(const Type& val) : prefix_(val) { ConcordAssert(val.size() <= SIZE_IN_BYTES); }
+  CustomPrefix(Type&& val) : prefix_(std::move(val)) { ConcordAssert(val.size() <= SIZE_IN_BYTES); }
+
+  void set(std::uint8_t pos, char c) {
+    ConcordAssert(pos < prefix_.size());
+    prefix_[pos] = c;
+  }
+  char get(std::uint8_t pos) const {
+    ConcordAssert(pos < prefix_.size());
+    return prefix_[pos];
+  }
+  bool operator==(const CustomPrefix& other) const { return prefix_ == other.prefix_; }
+  bool operator!=(const CustomPrefix& other) const { return prefix_ != other.prefix_; }
+  bool operator<(const CustomPrefix& other) const { return prefix_ < other.prefix_; }
+
+  const Type& value() const { return prefix_; }
+  size_t size() const { return prefix_.size(); }
+
+ protected:
+  Type prefix_;
+};
+
+class PaddedCustomPrefix : public CustomPrefix {
+ public:
+  template <class T>
+  PaddedCustomPrefix(T&& arg) : CustomPrefix(std::forward<T>(arg)) {}
+  char get(std::uint8_t pos) const {
+    ConcordAssert(pos < SIZE_IN_BYTES);
+    if (pos < prefix_.size()) {
+      return prefix_[pos];
+    } else {
+      return '\0';
+    }
+  }
+};
+
 // A nibble is 4 bits, stored in the lower bits of a byte.
 class Nibble {
  public:
@@ -243,7 +294,7 @@ class NibblePath {
   NibblePath() : num_nibbles_(0) {}
   NibblePath(size_t num_nibbles, const std::vector<uint8_t>& path) : num_nibbles_(num_nibbles), path_(path) {
     ConcordAssert(num_nibbles < Hash::MAX_NIBBLES);
-    ConcordAssert(path.size() == (num_nibbles % 2 ? (num_nibbles + 1) / 2 : num_nibbles / 2));
+    ConcordAssert(path.size() == pathLengthInBytes(num_nibbles));
   }
 
   bool operator==(const NibblePath& other) const { return num_nibbles_ == other.num_nibbles_ && path_ == other.path_; }
@@ -263,6 +314,11 @@ class NibblePath {
     return false;
   }
 
+  // Return the count of elements (bytes) in path_
+  static size_t pathLengthInBytes(size_t num_nibbles) {
+    return (num_nibbles % 2 ? (num_nibbles + 1) / 2 : num_nibbles / 2);
+  }
+
   // Return the length of the path_ in nibbles
   size_t length() const { return num_nibbles_; }
 

kvbc/include/sparse_merkle/db_reader.h

@@ -12,6 +12,8 @@
 
 #pragma once
 
+#include <string>
+
 namespace concord {
 namespace kvbc {
 namespace sparse_merkle {
@@ -28,7 +30,7 @@ class IDBReader {
   virtual ~IDBReader() = default;
 
   // Return the latest root node in the system
-  virtual BatchedInternalNode get_latest_root() const = 0;
+  virtual BatchedInternalNode get_latest_root(std::string custom_prefix = "") const = 0;
 
   // Retrieve a BatchedInternalNode given an InternalNodeKey.
   //

kvbc/include/sparse_merkle/internal_node.h

@@ -387,8 +387,8 @@ class BatchedInternalNode {
                                      Version version,
                                      Nibble child_key,
                                      size_t prefix_bits_in_common,
-                                     LeafChild child1,
-                                     LeafChild child2);
+                                     const LeafChild& child1,
+                                     const LeafChild& child2);
 
   // Remove the LeafChild at the given index. Updates of children should be set
   // to new_version.

kvbc/include/sparse_merkle/keys.h

@@ -24,34 +24,47 @@ namespace sparse_merkle {
 // are part of the batch by pointing to their containing BatchedInternalNode.
 class InternalNodeKey {
  public:
-  InternalNodeKey(Version version, const NibblePath& path) : version_(version), path_(path) {}
+  template <class T>
+  InternalNodeKey(T&& custom_prefix, Version version, const NibblePath& path)
+      : custom_prefix_(std::forward<T>(custom_prefix)), version_(version), path_(path) {}
 
   // Return the root of a sparse merkle tree at a given version.
-  static InternalNodeKey root(Version version) { return InternalNodeKey(version, NibblePath()); }
+  template <class T>
+  static InternalNodeKey root(T&& custom_prefix, Version version) {
+    return InternalNodeKey(std::forward<T>(custom_prefix), version, NibblePath());
+  }
 
-  bool operator==(const InternalNodeKey& other) const { return version_ == other.version_ && path_ == other.path_; }
+  bool operator==(const InternalNodeKey& other) const {
+    return custom_prefix_ == other.custom_prefix_ && version_ == other.version_ && path_ == other.path_;
+  }
 
-  // Compare by version then by path
+  // Compare by address, version and then by path
   bool operator<(const InternalNodeKey& other) const {
-    if (version_ == other.version_) {
-      return path_ < other.path_;
+    if (custom_prefix_ != other.custom_prefix_) {
+      return custom_prefix_ < other.custom_prefix_;
     }
-    return version_ < other.version_;
+    if (version_ != other.version_) {
+      return version_ < other.version_;
+    }
+    return path_ < other.path_;
   }
 
   std::string toString() const {
     if (path_.empty()) {
-      return std::string("<ROOT>") + "-" + version_.toString();
+      return custom_prefix_.value() + std::string("-<ROOT>") + "-" + version_.toString();
     }
-    return path_.toString() + "-" + version_.toString();
+    return custom_prefix_.value() + "-" + path_.toString() + "-" + version_.toString();
   }
 
   Version version() const { return version_; }
 
+  const CustomPrefix& customPrefix() const { return custom_prefix_; }
+
   const NibblePath& path() const { return path_; }
   NibblePath& path() { return path_; }
 
  private:
+  CustomPrefix custom_prefix_;
   Version version_;
   NibblePath path_;
 };
@@ -63,36 +76,46 @@ class InternalNodeKey {
 // merkle tree.
 class LeafKey {
  public:
-  static constexpr auto SIZE_IN_BYTES = Hash::SIZE_IN_BYTES + Version::SIZE_IN_BYTES;
+  static constexpr auto SIZE_IN_BYTES = PaddedCustomPrefix::TOTAL_SIZE + Hash::SIZE_IN_BYTES + Version::SIZE_IN_BYTES;
 
  public:
-  LeafKey(Hash key, Version version) : key_(key), version_(version) {}
+  template <class T>
+  LeafKey(Hash key, Version version, T&& address) : key_(key), version_(version), prefix_(std::forward<T>(address)) {}
+  LeafKey(Hash key, Version version) : key_(key), version_(version), prefix_("") {}
 
-  bool operator==(const LeafKey& other) const { return key_ == other.key_ && version_ == other.version_; }
+  bool operator==(const LeafKey& other) const {
+    return prefix_ == other.prefix_ && key_ == other.key_ && version_ == other.version_;
+  }
   bool operator!=(const LeafKey& other) const { return !(*this == other); }
 
-  // Compare by key_ then version_
+  // Compare by prefix_, key_ and then version_
   bool operator<(const LeafKey& other) const {
-    if (key_ < other.key_) {
-      return true;
+    if (prefix_ != other.prefix_) {
+      return prefix_ < other.prefix_;
+    }
+    if (key_ != other.key_) {
+      return key_ < other.key_;
     }
-    return key_ == other.key_ && version_ < other.version_;
+    return version_ < other.version_;
   }
 
-  std::string toString() const { return key_.toString() + "-" + version_.toString(); }
+  std::string toString() const { return key_.toString() + "-" + version_.toString() + "-" + prefix_.value(); }
 
   Nibble getNibble(const size_t n) const {
     ConcordAssert(n < Hash::MAX_NIBBLES);
     return key_.getNibble(n);
   }
 
+  const PaddedCustomPrefix& customPrefix() const { return prefix_; }
+
   Version version() const { return version_; }
 
   const Hash& hash() const { return key_; }
 
  private:
   Hash key_;
   Version version_;
+  PaddedCustomPrefix prefix_;
 };
 
 std::ostream& operator<<(std::ostream& os, const LeafKey&);

kvbc/include/sparse_merkle/tree.h

@@ -43,11 +43,15 @@ namespace sparse_merkle {
 class Tree {
  public:
   Tree() = default;
-  explicit Tree(std::shared_ptr<IDBReader> db_reader) : db_reader_(db_reader) { reset(); }
+  explicit Tree(std::shared_ptr<IDBReader> db_reader, std::string custom_prefix = "")
+      : db_reader_(db_reader), custom_prefix_(custom_prefix) {
+    reset();
+  }
 
   const Hash& get_root_hash() const { return root_.hash(); }
   Version get_version() const { return root_.version(); }
   bool empty() const { return root_.numChildren() == 0; }
+  std::string toString() { return root_.toString(); }
 
   // Add or update key-value pairs given in `updates`, and remove keys in
   // `deleted_keys`.
@@ -74,14 +78,15 @@ class Tree {
   //
   // This is necessary to do before updates, as we only allow updating the
   // latest tree.
-  void reset() { root_ = db_reader_->get_latest_root(); }
+  void reset() { root_ = db_reader_->get_latest_root(custom_prefix_); }
 
   UpdateBatch update_impl(const concord::kvbc::SetOfKeyValuePairs& updates,
                           const concord::kvbc::KeysVector& deleted_keys,
                           detail::UpdateCache& cache);
 
   std::shared_ptr<IDBReader> db_reader_;
   BatchedInternalNode root_;
+  std::string custom_prefix_;
 };
 
 }  // namespace sparse_merkle