diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.cpp b/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.cpp index 74f0ba80e3..d0319960a7 100644 --- a/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.cpp +++ b/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.cpp @@ -106,8 +106,8 @@ template fr_hash_path MerkleTree::get_hash_path(index_t current = hash_pair_native(path[j].first, path[j].second); } } else { - // Requesting path to a different, indepenent element. - // We know that this element exits in an empty subtree, of height determined by the common bits in the + // Requesting path to a different, independent element. + // We know that this element exists in an empty subtree, of height determined by the common bits in the // stumps index and the requested index. size_t common_bits = numeric::count_leading_zeros(diff); size_t common_height = sizeof(index_t) * 8 - common_bits - 1; @@ -133,6 +133,65 @@ template fr_hash_path MerkleTree::get_hash_path(index_t return path; } +template fr_sibling_path MerkleTree::get_sibling_path(index_t index) +{ + fr_sibling_path path(depth_); + + std::vector data; + bool status = store_.get(root().to_buffer(), data); + + for (size_t i = depth_ - 1; i < depth_; --i) { + if (!status) { + // This is an empty subtree. Fill in zero value. + path[i] = zero_hashes_[i]; + continue; + } + + if (data.size() == 64) { + // This is a regular node with left and right trees. Descend according to index path. + bool is_right = bit_set(index, i); + path[i] = from_buffer(data, is_right ? 0 : 32); + + auto it = data.data() + (is_right ? 32 : 0); + status = store_.get(std::vector(it, it + 32), data); + } else { + // This is a stump. The sibling path can be fully restored from this node. + // In case of a stump, we store: [key : (value, local_index, true)], i.e. 65-byte data. + ASSERT(data.size() == 65); + fr current = from_buffer(data, 0); + index_t element_index = from_buffer(data, 32); + index_t subtree_index = numeric::keep_n_lsb(index, i + 1); + index_t diff = element_index ^ subtree_index; + + // Populate the sibling path with zero hashes. + for (size_t j = 0; j <= i; ++j) { + path[j] = zero_hashes_[j]; + } + + // If diff == 0, we are requesting the sibling path of the only non-zero element in the stump which is + // populated only with zero hashes. + if (diff == 1) { + // Requesting path of the sibling of the non-zero leaf in the stump. + // Set the bottom element of the path to the non-zero leaf (the rest is already populated correctly + // with zero hashes). + path[0] = current; + } else if (diff > 1) { + // Requesting path to a different, independent element. + // We know that this element exists in an empty subtree, of height determined by the common bits in the + // stumps index and the requested index. + size_t common_bits = numeric::count_leading_zeros(diff); + size_t common_height = sizeof(index_t) * 8 - common_bits - 1; + + // Insert the only non-zero sibling at the common height. + path[common_height] = compute_zero_path_hash(common_height, element_index, current); + } + break; + } + } + + return path; +} + template fr MerkleTree::update_element(index_t index, fr const& value) { auto leaf = value; diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.hpp b/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.hpp index 161941207e..de28fb9bc5 100644 --- a/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.hpp +++ b/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.hpp @@ -21,6 +21,8 @@ template class MerkleTree { fr_hash_path get_hash_path(index_t index); + fr_sibling_path get_sibling_path(index_t index); + fr update_element(index_t index, fr const& value); fr root() const; diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.test.cpp b/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.test.cpp index db210fc8fe..7c9b5b7471 100644 --- a/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.test.cpp +++ b/cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.test.cpp @@ -105,6 +105,28 @@ TEST(stdlib_merkle_tree, test_get_hash_path) EXPECT_EQ(db.get_hash_path(512), memdb.get_hash_path(512)); } +TEST(stdlib_merkle_tree, test_get_sibling_path) +{ + MemoryTree memdb(10); + + MemoryStore store; + auto db = MerkleTree(store, 10); + + EXPECT_EQ(memdb.get_sibling_path(512), db.get_sibling_path(512)); + + memdb.update_element(512, VALUES[512]); + db.update_element(512, VALUES[512]); + + EXPECT_EQ(db.get_sibling_path(512), memdb.get_sibling_path(512)); + + for (size_t i = 0; i < 1024; ++i) { + memdb.update_element(i, VALUES[i]); + db.update_element(i, VALUES[i]); + } + + EXPECT_EQ(db.get_sibling_path(512), memdb.get_sibling_path(512)); +} + TEST(stdlib_merkle_tree, test_get_hash_path_layers) { { @@ -133,4 +155,36 @@ TEST(stdlib_merkle_tree, test_get_hash_path_layers) EXPECT_NE(before[2], after[2]); } } + +TEST(stdlib_merkle_tree, test_get_sibling_path_layers) +{ + { + MemoryStore store; + auto db = MerkleTree(store, 3); + + auto before = db.get_sibling_path(1); + db.update_element(0, VALUES[1]); + auto after = db.get_sibling_path(1); + + EXPECT_NE(before[0], after[0]); + EXPECT_EQ(before[1], after[1]); + EXPECT_EQ(before[2], after[2]); + } + + { + MemoryStore store; + auto db = MerkleTree(store, 3); + + auto before = db.get_sibling_path(7); + db.update_element(0x0, VALUES[1]); + auto after = db.get_sibling_path(7); + + info("before: ", before[0], " ", before[1], " ", before[2]); + info("after: ", after[0], " ", after[1], " ", after[2]); + + EXPECT_EQ(before[0], after[0]); + EXPECT_EQ(before[1], after[1]); + EXPECT_NE(before[2], after[2]); + } +} } // namespace proof_system::test_stdlib_merkle_tree \ No newline at end of file diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_tree.hpp b/cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_tree.hpp index 4e8155937a..d04c904046 100644 --- a/cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_tree.hpp +++ b/cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_tree.hpp @@ -25,12 +25,19 @@ template class NullifierTree : public MerkleTree { fr update_element(fr const& value); - private: + protected: using MerkleTree::update_element; using MerkleTree::get_element; using MerkleTree::compute_zero_path_hash; - private: + // const std::vector& get_hashes() { return hashes_; } + const WrappedNullifierLeaf get_leaf(size_t index) + { + // return (index < leaves_.size()) ? leaves_[index] : WrappedNullifierLeaf::zero(); + return (index < leaves.size()) ? leaves[index] : WrappedNullifierLeaf::zero(); + } + const std::vector& get_leaves() { return leaves; } + using MerkleTree::store_; using MerkleTree::zero_hashes_; using MerkleTree::depth_;