Skip to content

Commit

Permalink
feat: Add get_sibling_path method in MerkleTree (#584)
Browse files Browse the repository at this point in the history
  • Loading branch information
@benesjan
benesjan authored Jul 11, 2023
1 parent 0ac4bc3 commit b3db9f8
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 6 deletions.
75 changes: 69 additions & 6 deletions cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@ namespace merkle_tree {

using namespace barretenberg;

// Size of merkle tree nodes in bytes.
constexpr size_t REGULAR_NODE_SIZE = 64;
constexpr size_t STUMP_NODE_SIZE = 65;

template <typename T> inline bool bit_set(T const& index, size_t i)
{
return bool((index >> i) & 0x1);
Expand Down Expand Up @@ -76,7 +80,7 @@ template <typename Store> fr_hash_path MerkleTree<Store>::get_hash_path(index_t
continue;
}

if (data.size() == 64) {
if (data.size() == REGULAR_NODE_SIZE) {
// This is a regular node with left and right trees. Descend according to index path.
auto left = from_buffer<fr>(data, 0);
auto right = from_buffer<fr>(data, 32);
Expand All @@ -87,7 +91,7 @@ template <typename Store> fr_hash_path MerkleTree<Store>::get_hash_path(index_t
} else {
// This is a stump. The hash 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);
ASSERT(data.size() == STUMP_NODE_SIZE);
fr current = from_buffer<fr>(data, 0);
index_t element_index = from_buffer<index_t>(data, 32);
index_t subtree_index = numeric::keep_n_lsb(index, i + 1);
Expand All @@ -106,8 +110,8 @@ template <typename Store> fr_hash_path MerkleTree<Store>::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;
Expand All @@ -133,6 +137,65 @@ template <typename Store> fr_hash_path MerkleTree<Store>::get_hash_path(index_t
return path;
}

template <typename Store> fr_sibling_path MerkleTree<Store>::get_sibling_path(index_t index)
{
fr_sibling_path path(depth_);

std::vector<uint8_t> 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() == REGULAR_NODE_SIZE) {
// 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<fr>(data, is_right ? 0 : 32);

auto it = data.data() + (is_right ? 32 : 0);
status = store_.get(std::vector<uint8_t>(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() == STUMP_NODE_SIZE);
fr current = from_buffer<fr>(data, 0);
index_t element_index = from_buffer<index_t>(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 <typename Store> fr MerkleTree<Store>::update_element(index_t index, fr const& value)
{
auto leaf = value;
Expand Down Expand Up @@ -206,7 +269,7 @@ fr MerkleTree<Store>::update_element(fr const& root, fr const& value, index_t in
return key;
}

if (data.size() == 65) {
if (data.size() == STUMP_NODE_SIZE) {
// We've come across a stump.
index_t existing_index = from_buffer<index_t>(data, 32);

Expand All @@ -224,7 +287,7 @@ fr MerkleTree<Store>::update_element(fr const& root, fr const& value, index_t in
return fork_stump(existing_value, existing_index, value, index, height, common_height);
} else {
// If its not a stump, the data size must be 64 bytes.
ASSERT(data.size() == 64);
ASSERT(data.size() == REGULAR_NODE_SIZE);
bool is_right = bit_set(index, height - 1);
fr subtree_root = from_buffer<fr>(data, is_right ? 32 : 0);
fr subtree_root_copy = subtree_root;
Expand Down
2 changes: 2 additions & 0 deletions cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ template <typename Store> 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;
Expand Down
51 changes: 51 additions & 0 deletions cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.test.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
{
Expand Down Expand Up @@ -133,4 +155,33 @@ 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);

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

0 comments on commit b3db9f8

Please sign in to comment.