diff --git a/kvbc/cmf/categorized_kvbc_msgs.cmf b/kvbc/cmf/categorized_kvbc_msgs.cmf index 46be39088e..278ebfbaa9 100644 --- a/kvbc/cmf/categorized_kvbc_msgs.cmf +++ b/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 } diff --git a/kvbc/include/categorization/block_merkle_category.h b/kvbc/include/categorization/block_merkle_category.h index b079abc732..36aaaddc3b 100644 --- a/kvbc/include/categorization/block_merkle_category.h +++ b/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. // diff --git a/kvbc/include/merkle_tree_db_adapter.h b/kvbc/include/merkle_tree_db_adapter.h index 78d3b8524f..56472d40fb 100644 --- a/kvbc/include/merkle_tree_db_adapter.h +++ b/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. // diff --git a/kvbc/include/merkle_tree_serialization.h b/kvbc/include/merkle_tree_serialization.h index 5062a8f46d..565f2615d6 100644 --- a/kvbc/include/merkle_tree_serialization.h +++ b/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::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(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::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(custom_prefix.size())); + for (auto i = 0u; i < sparse_merkle::PaddedCustomPrefix::SIZE_IN_BYTES; i++) { + out.append(1, static_cast(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 deserialize>(const co template <> inline sparse_merkle::NibblePath deserialize(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(buf[0]); - const auto path = deserialize>( - concordUtils::Sliver{buf, sizeof(std::uint8_t), buf.length() - sizeof(std::uint8_t)}); + size_t nibbleCount = static_cast(buf[0]); + size_t pathBytesCount = sparse_merkle::NibblePath::pathLengthInBytes(nibbleCount); + std::vector path; + if (pathBytesCount > 0) { + path = + deserialize>(concordUtils::Sliver{buf, nibblePathSizePrefixInBytes, pathBytesCount}); + } return sparse_merkle::NibblePath{nibbleCount, path}; } +template <> +inline sparse_merkle::CustomPrefix deserialize(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(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( + const concordUtils::Sliver &buf) { + return sparse_merkle::PaddedCustomPrefix(deserialize(buf)); +} + template <> inline sparse_merkle::LeafKey deserialize(const concordUtils::Sliver &buf) { ConcordAssert(buf.length() >= sparse_merkle::LeafKey::SIZE_IN_BYTES); - return sparse_merkle::LeafKey{deserialize(buf), - deserialize(concordUtils::Sliver{ - buf, sparse_merkle::Hash::SIZE_IN_BYTES, sparse_merkle::Version::SIZE_IN_BYTES})}; + return sparse_merkle::LeafKey{ + deserialize(buf), + deserialize( + concordUtils::Sliver{buf, sparse_merkle::Hash::SIZE_IN_BYTES, sparse_merkle::Version::SIZE_IN_BYTES}), + deserialize(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(const concordUtils::Sliver &buf) { - return sparse_merkle::InternalNodeKey{ - deserialize(buf), - deserialize(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(buf); + auto version = deserialize( + concordUtils::Sliver{buf, SIZE_PREFIX_BYTES + prefix.size(), buf.length() - SIZE_PREFIX_BYTES - prefix.size()}); + auto nibblePath = deserialize( + 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 <> diff --git a/kvbc/include/sparse_merkle/base_types.h b/kvbc/include/sparse_merkle/base_types.h index b03412f08e..47d564b1fb 100644 --- a/kvbc/include/sparse_merkle/base_types.h +++ b/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 + PaddedCustomPrefix(T&& arg) : CustomPrefix(std::forward(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& 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_; } diff --git a/kvbc/include/sparse_merkle/db_reader.h b/kvbc/include/sparse_merkle/db_reader.h index a23850d773..9909190385 100644 --- a/kvbc/include/sparse_merkle/db_reader.h +++ b/kvbc/include/sparse_merkle/db_reader.h @@ -12,6 +12,8 @@ #pragma once +#include + 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. // diff --git a/kvbc/include/sparse_merkle/internal_node.h b/kvbc/include/sparse_merkle/internal_node.h index a587154571..52e9e7ca3a 100644 --- a/kvbc/include/sparse_merkle/internal_node.h +++ b/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. diff --git a/kvbc/include/sparse_merkle/keys.h b/kvbc/include/sparse_merkle/keys.h index 52575ee1f4..4161ccc5f4 100644 --- a/kvbc/include/sparse_merkle/keys.h +++ b/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 + InternalNodeKey(T&& custom_prefix, Version version, const NibblePath& path) + : custom_prefix_(std::forward(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 + static InternalNodeKey root(T&& custom_prefix, Version version) { + return InternalNodeKey(std::forward(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("") + "-" + version_.toString(); + return custom_prefix_.value() + std::string("-") + "-" + 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,29 +76,38 @@ 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 + LeafKey(Hash key, Version version, T&& address) : key_(key), version_(version), prefix_(std::forward(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_; } @@ -93,6 +115,7 @@ class LeafKey { private: Hash key_; Version version_; + PaddedCustomPrefix prefix_; }; std::ostream& operator<<(std::ostream& os, const LeafKey&); diff --git a/kvbc/include/sparse_merkle/tree.h b/kvbc/include/sparse_merkle/tree.h index 59911a36d6..0497a9af9a 100644 --- a/kvbc/include/sparse_merkle/tree.h +++ b/kvbc/include/sparse_merkle/tree.h @@ -43,11 +43,15 @@ namespace sparse_merkle { class Tree { public: Tree() = default; - explicit Tree(std::shared_ptr db_reader) : db_reader_(db_reader) { reset(); } + explicit Tree(std::shared_ptr 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,7 +78,7 @@ 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, @@ -82,6 +86,7 @@ class Tree { std::shared_ptr db_reader_; BatchedInternalNode root_; + std::string custom_prefix_; }; } // namespace sparse_merkle diff --git a/kvbc/include/sparse_merkle/update_cache.h b/kvbc/include/sparse_merkle/update_cache.h index bea85595e6..668bcdc0a4 100644 --- a/kvbc/include/sparse_merkle/update_cache.h +++ b/kvbc/include/sparse_merkle/update_cache.h @@ -34,8 +34,8 @@ namespace detail { // A new instance of this structure is created during every update call. class UpdateCache { public: - UpdateCache(const BatchedInternalNode& root, const std::shared_ptr& db_reader) - : version_(root.version() + 1), db_reader_(db_reader), original_root_(root) {} + UpdateCache(const BatchedInternalNode& root, const std::shared_ptr& db_reader, std::string custom_prefix) + : version_(root.version() + 1), db_reader_(db_reader), original_root_(root), custom_prefix_(custom_prefix) {} const StaleNodeIndexes& stale() const { return stale_; } const auto& internalNodes() const { return internal_nodes_; } @@ -55,8 +55,8 @@ class UpdateCache { void putStale(const std::optional& key); void putStale(const InternalNodeKey& key); - void put(const NibblePath& path, const BatchedInternalNode& node); - void remove(const NibblePath& path); + void put(const InternalNodeKey& key, const BatchedInternalNode& node); + void remove(const InternalNodeKey& key); private: // The version of the tree after this update is complete. @@ -73,7 +73,8 @@ class UpdateCache { // // These nodes are mutable and are all being updated to a single version. // Therefore we key them by their NibblePath alone, without a version. - std::map internal_nodes_; + std::map internal_nodes_; + std::string custom_prefix_; }; } // namespace detail diff --git a/kvbc/include/sparse_merkle/walker.h b/kvbc/include/sparse_merkle/walker.h index 854cb2cf82..02431a1091 100644 --- a/kvbc/include/sparse_merkle/walker.h +++ b/kvbc/include/sparse_merkle/walker.h @@ -28,7 +28,8 @@ typedef std::stack> NodeSt // A class for ascending and descending the tree. This is used for updates. class Walker { public: - Walker(UpdateCache& cache) : cache_(cache), current_node_(cache.getRoot()) { + Walker(UpdateCache& cache, std::string custom_prefix) + : cache_(cache), current_node_(cache.getRoot()), custom_prefix_(custom_prefix) { // Save the version of the root node, in case we only update the root node. stale_version_ = current_node_.version(); } @@ -56,6 +57,8 @@ class Walker { void insertEmptyRootAtCurrentVersion(); + const std::string& customPrefix() const { return custom_prefix_; } + private: void ascend(); std::pair pop(); @@ -74,6 +77,7 @@ class Walker { NodeStack stack_; BatchedInternalNode current_node_; NibblePath nibble_path_; + std::string custom_prefix_; }; } // namespace detail diff --git a/kvbc/src/categorization/block_merkle_category.cpp b/kvbc/src/categorization/block_merkle_category.cpp index abf13aa439..4d19fa4163 100644 --- a/kvbc/src/categorization/block_merkle_category.cpp +++ b/kvbc/src/categorization/block_merkle_category.cpp @@ -71,17 +71,17 @@ VersionedKey leafKeyToVersionedKey(const sparse_merkle::LeafKey& leaf_key) { // favor of the one immediately below it. BatchedInternalNodeKey toBatchedInternalNodeKey(const sparse_merkle::InternalNodeKey& key) { auto path = NibblePath{static_cast(key.path().length()), key.path().data()}; - return BatchedInternalNodeKey{key.version().value(), std::move(path)}; + return BatchedInternalNodeKey{key.customPrefix().value(), key.version().value(), std::move(path)}; } BatchedInternalNodeKey toBatchedInternalNodeKey(sparse_merkle::InternalNodeKey&& key) { auto path = NibblePath{static_cast(key.path().length()), key.path().move_data()}; - return BatchedInternalNodeKey{key.version().value(), std::move(path)}; + return BatchedInternalNodeKey{key.customPrefix().value(), key.version().value(), std::move(path)}; } -std::vector rootKey(uint64_t version) { +std::vector rootKey(uint64_t version, const std::string& custom_prefix = "") { auto v = sparse_merkle::Version(version); - return serialize(toBatchedInternalNodeKey(sparse_merkle::InternalNodeKey::root(v))); + return serialize(toBatchedInternalNodeKey(sparse_merkle::InternalNodeKey::root(custom_prefix, v))); } // A key used in tree_.update() @@ -775,8 +775,8 @@ uint64_t BlockMerkleCategory::getLatestTreeVersion() const { return 0; } -sparse_merkle::BatchedInternalNode BlockMerkleCategory::Reader::get_latest_root() const { - if (auto latest_root_key = db_.get(BLOCK_MERKLE_INTERNAL_NODES_CF, rootKey(0))) { +sparse_merkle::BatchedInternalNode BlockMerkleCategory::Reader::get_latest_root(std::string custom_prefix) const { + if (auto latest_root_key = db_.get(BLOCK_MERKLE_INTERNAL_NODES_CF, rootKey(0, custom_prefix))) { if (auto serialized = db_.get(BLOCK_MERKLE_INTERNAL_NODES_CF, *latest_root_key)) { return deserializeBatchedInternalNode(*serialized); } diff --git a/kvbc/src/merkle_tree_db_adapter.cpp b/kvbc/src/merkle_tree_db_adapter.cpp index c799c8b823..c5fc5eeb07 100644 --- a/kvbc/src/merkle_tree_db_adapter.cpp +++ b/kvbc/src/merkle_tree_db_adapter.cpp @@ -446,7 +446,7 @@ std::pair DBAdap return smTree_.update_with_cache(updates, KeysVector{std::cbegin(deletes), std::cend(deletes)}); } -BatchedInternalNode DBAdapter::Reader::get_latest_root() const { +BatchedInternalNode DBAdapter::Reader::get_latest_root(std::string custom_prefix) const { const auto lastBlock = adapter_.getLastReachableBlockId(); if (lastBlock == 0) { return BatchedInternalNode{}; @@ -462,7 +462,7 @@ BatchedInternalNode DBAdapter::Reader::get_latest_root() const { return BatchedInternalNode{}; } - return get_internal(InternalNodeKey::root(stateRootVersion)); + return get_internal(InternalNodeKey::root(custom_prefix, stateRootVersion)); } BatchedInternalNode DBAdapter::Reader::get_internal(const InternalNodeKey &key) const { @@ -708,7 +708,7 @@ KeysVector DBAdapter::internalProvableKeysForVersion(const Version &version) con // Rely on the fact that root internal keys always precede non-root ones - due to lexicographical ordering and root // internal keys having empty nibble paths. See InternalNodeKey serialization code. return keysForVersion(db_, - DBKeyManipulator::genInternalDbKey(InternalNodeKey::root(version)), + DBKeyManipulator::genInternalDbKey(InternalNodeKey::root("", version)), version, EKeySubtype::Internal, [](const Key &key) { return DBKeyManipulator::extractVersionFromInternalKey(key); }); diff --git a/kvbc/src/sparse_merkle/internal_node.cpp b/kvbc/src/sparse_merkle/internal_node.cpp index 193b46ecdb..1dc6d1115c 100644 --- a/kvbc/src/sparse_merkle/internal_node.cpp +++ b/kvbc/src/sparse_merkle/internal_node.cpp @@ -173,8 +173,12 @@ size_t BatchedInternalNode::insertInternalChildren(size_t index, return index; } -BatchedInternalNode::InsertResult BatchedInternalNode::insertTwoLeafChildren( - size_t index, Version version, Nibble child_key, size_t prefix_bits_in_common, LeafChild child1, LeafChild child2) { +BatchedInternalNode::InsertResult BatchedInternalNode::insertTwoLeafChildren(size_t index, + Version version, + Nibble child_key, + size_t prefix_bits_in_common, + const LeafChild& child1, + const LeafChild& child2) { size_t child1_index = 0; size_t child2_index = 0; // prefix_bits_in_common start from MSB. At most there can be 3 bits in common diff --git a/kvbc/src/sparse_merkle/keys.cpp b/kvbc/src/sparse_merkle/keys.cpp index fd9de2e134..221593b042 100644 --- a/kvbc/src/sparse_merkle/keys.cpp +++ b/kvbc/src/sparse_merkle/keys.cpp @@ -15,7 +15,7 @@ namespace concord::kvbc::sparse_merkle { std::ostream& operator<<(std::ostream& os, const LeafKey& key) { - os << key.hash() << "-" << key.version(); + os << key.hash() << "-" << key.version() << "-" << key.customPrefix().value(); return os; } diff --git a/kvbc/src/sparse_merkle/tree.cpp b/kvbc/src/sparse_merkle/tree.cpp index 238f7e241c..22c1853f5f 100644 --- a/kvbc/src/sparse_merkle/tree.cpp +++ b/kvbc/src/sparse_merkle/tree.cpp @@ -101,7 +101,7 @@ void remove(Walker& walker, const Hash& key_hash) { if (auto rv = std::get_if(&result)) { histograms.remove_depth->record(walker.depth()); - auto stale = LeafKey(key_hash, rv->version); + auto stale = LeafKey(key_hash, rv->version, walker.customPrefix()); return walker.ascendToRoot(stale); } @@ -111,7 +111,7 @@ void remove(Walker& walker, const Hash& key_hash) { } if (auto rv = std::get_if(&result)) { - walker.markStale(LeafKey(key_hash, rv->removed_version)); + walker.markStale(LeafKey(key_hash, rv->removed_version, walker.customPrefix())); return removeBatchedInternalNode(walker, rv->promoted); } @@ -133,14 +133,14 @@ UpdateBatch Tree::update(const concord::kvbc::SetOfKeyValuePairs& updates, histograms.num_deleted_keys->record(deleted_keys.size()); TimeRecorder scoped_timer(*histograms.update); reset(); - UpdateCache cache(root_, db_reader_); + UpdateCache cache(root_, db_reader_, custom_prefix_); return update_impl(updates, deleted_keys, cache); } std::pair Tree::update_with_cache(const concord::kvbc::SetOfKeyValuePairs& updates, const concord::kvbc::KeysVector& deleted_keys) { reset(); - UpdateCache cache(root_, db_reader_); + UpdateCache cache(root_, db_reader_, custom_prefix_); auto batch = update_impl(updates, deleted_keys, cache); return std::make_pair(batch, cache); } @@ -155,7 +155,7 @@ UpdateBatch Tree::update_impl(const concord::kvbc::SetOfKeyValuePairs& updates, // Deletes come before inserts because it makes more semantic sense. A user can delete a key and then write a new // version, but it makes no sense to add a new version and then delete a key. for (auto& key : deleted_keys) { - Walker walker(cache); + Walker walker(cache, custom_prefix_); auto key_hash = hasher.hash(key.data(), key.length()); sparse_merkle::remove(walker, key_hash); } @@ -169,9 +169,9 @@ UpdateBatch Tree::update_impl(const concord::kvbc::SetOfKeyValuePairs& updates, leaf_hash = hasher.hash(val.data(), val.length()); } LeafNode leaf_node{val}; - LeafKey leaf_key{hasher.hash(key.data(), key.length()), version}; + LeafKey leaf_key{hasher.hash(key.data(), key.length()), version, custom_prefix_}; LeafChild child{leaf_hash, leaf_key}; - Walker walker(cache); + Walker walker(cache, custom_prefix_); insert(walker, child); batch.leaf_nodes.emplace_back(leaf_key, leaf_node); } @@ -180,7 +180,7 @@ UpdateBatch Tree::update_impl(const concord::kvbc::SetOfKeyValuePairs& updates, batch.stale = cache.stale(); batch.stale.stale_since_version = version; for (auto& it : cache.internalNodes()) { - batch.internal_nodes.emplace_back(InternalNodeKey(version, it.first), it.second); + batch.internal_nodes.emplace_back(it.first, it.second); } updateBatchHistograms(batch); diff --git a/kvbc/src/sparse_merkle/update_cache.cpp b/kvbc/src/sparse_merkle/update_cache.cpp index 2d1e2a2248..be0996dfcb 100644 --- a/kvbc/src/sparse_merkle/update_cache.cpp +++ b/kvbc/src/sparse_merkle/update_cache.cpp @@ -15,7 +15,7 @@ namespace concord::kvbc::sparse_merkle::detail { BatchedInternalNode UpdateCache::getInternalNode(const InternalNodeKey& key) { - auto it = internal_nodes_.find(key.path()); + auto it = internal_nodes_.find(key); if (it != internal_nodes_.end()) { return it->second; } @@ -29,17 +29,17 @@ void UpdateCache::putStale(const std::optional& key) { } const BatchedInternalNode& UpdateCache::getRoot() { - auto it = internal_nodes_.find(NibblePath()); + auto it = internal_nodes_.find({custom_prefix_, version_, NibblePath()}); if (it != internal_nodes_.end()) { return it->second; } return original_root_; } -void UpdateCache::put(const NibblePath& path, const BatchedInternalNode& node) { internal_nodes_[path] = node; } +void UpdateCache::put(const InternalNodeKey& key, const BatchedInternalNode& node) { internal_nodes_[key] = node; } void UpdateCache::putStale(const InternalNodeKey& key) { stale_.internal_keys.insert(key); } -void UpdateCache::remove(const NibblePath& path) { internal_nodes_.erase(path); } +void UpdateCache::remove(const InternalNodeKey& key) { internal_nodes_.erase(key); } } // namespace concord::kvbc::sparse_merkle::detail diff --git a/kvbc/src/sparse_merkle/walker.cpp b/kvbc/src/sparse_merkle/walker.cpp index 5687f934b0..40b322c2a1 100644 --- a/kvbc/src/sparse_merkle/walker.cpp +++ b/kvbc/src/sparse_merkle/walker.cpp @@ -36,7 +36,7 @@ void Walker::descend(const Hash& key, Version next_version) { stack_.push(current_node_); Nibble next_nibble = key.getNibble(depth()); nibble_path_.append(next_nibble); - InternalNodeKey next_internal_key{next_version, nibble_path_}; + InternalNodeKey next_internal_key{custom_prefix_, next_version, nibble_path_}; current_node_ = cache_.getInternalNode(next_internal_key); stale_version_ = current_node_.version(); } @@ -68,7 +68,7 @@ void Walker::ascend() { std::optional Walker::removeCurrentNode() { markCurrentNodeStale(); - cache_.remove(nibble_path_); + cache_.remove({custom_prefix_, stale_version_, nibble_path_}); if (!atRoot()) { return pop().first; } @@ -78,7 +78,7 @@ std::optional Walker::removeCurrentNode() { void Walker::insertEmptyRootAtCurrentVersion() { auto children = BatchedInternalNode::Children{}; children[0] = InternalChild{PLACEHOLDER_HASH, cache_.version()}; - cache_.put(NibblePath{}, BatchedInternalNode{children}); + cache_.put({custom_prefix_, cache_.version(), NibblePath{}}, BatchedInternalNode{children}); } std::pair Walker::pop() { @@ -97,10 +97,10 @@ void Walker::markCurrentNodeStale() { // Don't mark the initial root stale, and don't mark an already cached node // stale. if (stale_version_ != Version(0) && stale_version_ != version()) { - cache_.putStale(InternalNodeKey(stale_version_, nibble_path_)); + cache_.putStale(InternalNodeKey(custom_prefix_, stale_version_, nibble_path_)); } } -void Walker::cacheCurrentNode() { cache_.put(nibble_path_, current_node_); } +void Walker::cacheCurrentNode() { cache_.put({custom_prefix_, cache_.version(), nibble_path_}, current_node_); } } // namespace concord::kvbc::sparse_merkle::detail diff --git a/kvbc/test/sparse_merkle/test_db.h b/kvbc/test/sparse_merkle/test_db.h index bac6126db5..69ed0946e5 100644 --- a/kvbc/test/sparse_merkle/test_db.h +++ b/kvbc/test/sparse_merkle/test_db.h @@ -15,32 +15,35 @@ #include "sparse_merkle/keys.h" #include "sparse_merkle/db_reader.h" #include "sparse_merkle/internal_node.h" +#include "kvbc/include/categorization/details.h" using namespace std; using namespace concord::kvbc::sparse_merkle; -class TestDB : public concord::kvbc::sparse_merkle::IDBReader { +class TestDB : public IDBReader { public: void put(const LeafKey& key, const LeafNode& val) { leaf_nodes_[key] = val; } void put(const InternalNodeKey& key, const BatchedInternalNode& val) { + ConcordAssert(internal_nodes_.find(key) == internal_nodes_.end()); + internal_nodes_[key] = val; - if (latest_version_ < key.version()) { - latest_version_ = key.version(); + if (latest_version_[key.customPrefix().value()] < key.version()) { + latest_version_[key.customPrefix().value()] = key.version(); } } - BatchedInternalNode get_latest_root() const override { - if (latest_version_ == 0) { + BatchedInternalNode get_latest_root(std::string custom_prefix = "") const override { + if (latest_version_.find(custom_prefix) == latest_version_.end()) { return BatchedInternalNode(); } - auto root_key = InternalNodeKey::root(latest_version_); + auto root_key = InternalNodeKey::root(custom_prefix, latest_version_.at(custom_prefix)); return internal_nodes_.at(root_key); } BatchedInternalNode get_internal(const InternalNodeKey& key) const override { return internal_nodes_.at(key); } private: - Version latest_version_ = 0; + map latest_version_; map leaf_nodes_; map internal_nodes_; }; diff --git a/kvbc/test/sparse_merkle/tree_test.cpp b/kvbc/test/sparse_merkle/tree_test.cpp index 1bc24e4e68..61b8cc3b56 100644 --- a/kvbc/test/sparse_merkle/tree_test.cpp +++ b/kvbc/test/sparse_merkle/tree_test.cpp @@ -274,7 +274,7 @@ TEST(tree_tests, insert_single_leaf_node) { // Only a single node was created. It's key is made up of an empty nibble // path, since there are no collisions requring going deeper into the tree. auto& root = batch.internal_nodes.at(0); - ASSERT_EQ(InternalNodeKey(1, NibblePath()), root.first); + ASSERT_EQ(InternalNodeKey("", 1, NibblePath()), root.first); // There should be a root and a single leaf ASSERT_EQ(2, root.second.numChildren()); diff --git a/kvbc/test/sparse_merkle_storage/db_adapter_unit_test.cpp b/kvbc/test/sparse_merkle_storage/db_adapter_unit_test.cpp index b0c86a7a93..148103cb9f 100644 --- a/kvbc/test/sparse_merkle_storage/db_adapter_unit_test.cpp +++ b/kvbc/test/sparse_merkle_storage/db_adapter_unit_test.cpp @@ -113,7 +113,7 @@ bool hasProvableStaleIndexKeysSince(const std::shared_ptr &db, const bool hasInternalKeysForVersion(const std::shared_ptr &db, const Version &version) { auto iter = db->getIteratorGuard(); - const auto key = iter->seekAtLeast(DBKeyManipulator::genInternalDbKey(InternalNodeKey::root(version))).first; + const auto key = iter->seekAtLeast(DBKeyManipulator::genInternalDbKey(InternalNodeKey::root("", version))).first; if (!key.empty() && DBKeyManipulator::getDBKeyType(key) == EDBKeyType::Key && DBKeyManipulator::getKeySubtype(key) == EKeySubtype::Internal && DBKeyManipulator::extractVersionFromInternalKey(key) == version) { diff --git a/kvbc/test/sparse_merkle_storage/serialization_unit_test.cpp b/kvbc/test/sparse_merkle_storage/serialization_unit_test.cpp index 0a6828afc0..711959b3ff 100644 --- a/kvbc/test/sparse_merkle_storage/serialization_unit_test.cpp +++ b/kvbc/test/sparse_merkle_storage/serialization_unit_test.cpp @@ -155,9 +155,10 @@ TEST(key_manipulator, extract_key_from_non_provable_stale_key) { TEST(key_manipulator, data_key_leaf) { const auto leafKey = LeafKey{defaultHash, defaultVersion}; const auto key = DBKeyManipulator::genDataDbKey(leafKey); + std::string emptyCustomPrefix(concord::kvbc::sparse_merkle::PaddedCustomPrefix::TOTAL_SIZE, '\0'); const auto expected = toSliver(serializeEnum(EDBKeyType::Key) + serializeEnum(EKeySubtype::Leaf) + defaultHashStrBuf + - serializeIntegral(leafKey.version().value())); - ASSERT_EQ(key.length(), 1 + 1 + 32 + 8); + serializeIntegral(leafKey.version().value()) + emptyCustomPrefix); + ASSERT_EQ(key.length(), 1 + 1 + 32 + 8 + 21); ASSERT_TRUE(key == expected); } @@ -166,9 +167,10 @@ TEST(key_manipulator, data_key_leaf) { TEST(key_manipulator, data_key_sliver) { const auto version = 42ul; const auto key = DBKeyManipulator::genDataDbKey(defaultSliver, version); + std::string emptyCustomPrefix(concord::kvbc::sparse_merkle::PaddedCustomPrefix::TOTAL_SIZE, '\0'); const auto expected = toSliver(serializeEnum(EDBKeyType::Key) + serializeEnum(EKeySubtype::Leaf) + defaultHashStrBuf + - serializeIntegral(version)); - ASSERT_EQ(key.length(), 1 + 1 + 32 + 8); + serializeIntegral(version) + emptyCustomPrefix); + ASSERT_EQ(key.length(), 1 + 1 + 32 + 8 + 21); ASSERT_TRUE(key == expected); } @@ -182,13 +184,14 @@ TEST(key_manipulator, internal_key_even) { // second byte of the nibble path is 0x34 (by appending 0x03 and 0x04) path.append(0x03); path.append(0x04); - const auto internalKey = InternalNodeKey{defaultVersion, path}; + const auto internalKey = InternalNodeKey{"", defaultVersion, path}; const auto key = DBKeyManipulator::genInternalDbKey(internalKey); // Expect that two bytes of nibbles have been written. const auto expected = toSliver(serializeEnum(EDBKeyType::Key) + serializeEnum(EKeySubtype::Internal) + + serializeIntegral(std::uint8_t{0x0}) + serializeIntegral(internalKey.version().value()) + serializeIntegral(std::uint8_t{4}) + serializeIntegral(std::uint8_t{0x12}) + serializeIntegral(std::uint8_t{0x34})); - ASSERT_EQ(key.length(), 1 + 1 + 8 + 1 + 2); + ASSERT_EQ(key.length(), 1 + 1 + 8 + 1 + 2 + 1); ASSERT_TRUE(key == expected); } @@ -201,13 +204,17 @@ TEST(key_manipulator, internal_key_odd) { path.append(0x02); // second byte of the nibble path is 0x30 (by appending 0x03) path.append(0x03); - const auto internalKey = InternalNodeKey{defaultVersion, path}; + std::string custom_prefix = "test"; + const auto internalKey = InternalNodeKey{custom_prefix, defaultVersion, path}; const auto key = DBKeyManipulator::genInternalDbKey(internalKey); // Expect that two bytes of nibbles have been written. const auto expected = toSliver(serializeEnum(EDBKeyType::Key) + serializeEnum(EKeySubtype::Internal) + + serializeIntegral(std::uint8_t{0x4}) + serializeIntegral(custom_prefix[0]) + + serializeIntegral(custom_prefix[1]) + serializeIntegral(custom_prefix[2]) + + serializeIntegral(custom_prefix[3]) + serializeIntegral(internalKey.version().value()) + serializeIntegral(std::uint8_t{3}) + serializeIntegral(std::uint8_t{0x12}) + serializeIntegral(std::uint8_t{0x30})); - ASSERT_EQ(key.length(), 1 + 1 + 8 + 1 + 2); + ASSERT_EQ(key.length(), 1 + 1 + 8 + 1 + 2 + 1 + custom_prefix.size()); ASSERT_TRUE(key == expected); } @@ -292,12 +299,13 @@ TEST(key_manipulator, stale_db_key_internal) { // second byte of the nibble path is 0x34 (by appending 0x03 and 0x04) path.append(0x03); path.append(0x04); - const auto internalKey = InternalNodeKey{defaultVersion, path}; + std::string custom_prefix = "abcdefg123"; + const auto internalKey = InternalNodeKey{custom_prefix, defaultVersion, path}; const auto key = DBKeyManipulator::genStaleDbKey(internalKey, defaultVersion); const auto expected = toSliver(serializeEnum(EDBKeyType::Key) + serializeEnum(EKeySubtype::ProvableStale) + serializeIntegral(defaultBlockId) + DBKeyManipulator::genInternalDbKey(internalKey).toString()); - ASSERT_EQ(key.length(), 1 + 1 + 8 + 1 + 1 + 8 + 1 + 2); + ASSERT_EQ(key.length(), 1 + 1 + 8 + 1 + 1 + 8 + 1 + 2 + 1 + custom_prefix.size()); ASSERT_TRUE(key == expected); } @@ -310,7 +318,7 @@ TEST(key_manipulator, stale_db_key_leaf) { const auto expected = toSliver(serializeEnum(EDBKeyType::Key) + serializeEnum(EKeySubtype::ProvableStale) + serializeIntegral(defaultBlockId) + DBKeyManipulator::genDataDbKey(leafKey).toString()); - ASSERT_EQ(key.length(), 1 + 1 + 8 + 1 + 1 + 32 + 8); + ASSERT_EQ(key.length(), 1 + 1 + 8 + 1 + 1 + 32 + 8 + 21); ASSERT_TRUE(key == expected); }