feat(memory_tree|a3): add sibling path calculations (#301)

* feat(memory_tree): frontier paths * fix: sibling path rename, nullifier tree -> protected vars * clean * test: sibling path test --------- Co-authored-by: cheethas <[email protected]>
AztecProtocol · Apr 5, 2023 · 7cbc8e0 · 7cbc8e0
1 parent 0c20f56
commit 7cbc8e0
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 3 deletions.
diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/hash_path.hpp b/cpp/src/barretenberg/stdlib/merkle_tree/hash_path.hpp
@@ -11,6 +11,7 @@ namespace merkle_tree {
 using namespace barretenberg;
 
 typedef std::vector<std::pair<fr, fr>> fr_hash_path;
+typedef std::vector<fr> fr_sibling_path;
 template <typename Ctx> using hash_path = std::vector<std::pair<field_t<Ctx>, field_t<Ctx>>>;
 
 inline fr_hash_path get_new_hash_path(fr_hash_path const& old_path, uint128_t index, fr const& value)

diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/memory_tree.cpp b/cpp/src/barretenberg/stdlib/merkle_tree/memory_tree.cpp
@@ -40,6 +40,24 @@ fr_hash_path MemoryTree::get_hash_path(size_t index)
     return path;
 }
 
+fr_sibling_path MemoryTree::get_sibling_path(size_t index)
+{
+    fr_sibling_path path(depth_);
+    size_t offset = 0;
+    size_t layer_size = total_size_;
+    for (size_t i = 0; i < depth_; i++) {
+        if (index % 2 == 0) {
+            path[i] = hashes_[offset + index + 1];
+        } else {
+            path[i] = hashes_[offset + index - 1];
+        }
+        offset += layer_size;
+        layer_size >>= 1;
+        index >>= 1;
+    }
+    return path;
+}
+
 fr MemoryTree::update_element(size_t index, fr const& value)
 {
     size_t offset = 0;

diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/memory_tree.hpp b/cpp/src/barretenberg/stdlib/merkle_tree/memory_tree.hpp
@@ -30,6 +30,8 @@ class MemoryTree {
 
     fr_hash_path get_hash_path(size_t index);
 
+    fr_sibling_path get_sibling_path(size_t index);
+
     fr update_element(size_t index, fr const& value);
 
     fr root() const { return root_; }

diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/memory_tree.test.cpp b/cpp/src/barretenberg/stdlib/merkle_tree/memory_tree.test.cpp
@@ -1,6 +1,5 @@
 #include "memory_tree.hpp"
 #include <gtest/gtest.h>
-#include "barretenberg/stdlib/types/types.hpp"
 
 using namespace barretenberg;
 using namespace plonk::stdlib::merkle_tree;
@@ -43,3 +42,39 @@ TEST(stdlib_merkle_tree, test_memory_store)
     EXPECT_EQ(db.get_hash_path(3), expected);
     EXPECT_EQ(db.root(), root);
 }
+
+TEST(stdlib_merkle_tree, test_memory_store_sibling_path)
+{
+    fr e00 = 0;
+    fr e01 = VALUES[1];
+    fr e02 = VALUES[2];
+    fr e03 = VALUES[3];
+    fr e10 = hash_pair_native(e00, e01);
+    fr e11 = hash_pair_native(e02, e03);
+    fr root = hash_pair_native(e10, e11);
+
+    MemoryTree db(2);
+    for (size_t i = 0; i < 4; ++i) {
+        db.update_element(i, VALUES[i]);
+    }
+
+    // Check correct paths are generated for each layer 0 element
+    fr_sibling_path expected00 = {
+        e01,
+        e11,
+    };
+    fr_sibling_path expected01 = { e00, e11 };
+    fr_sibling_path expected02 = {
+        e03,
+        e10,
+    };
+    fr_sibling_path expected03 = {
+        e02,
+        e10,
+    };
+    EXPECT_EQ(db.get_sibling_path(0), expected00);
+    EXPECT_EQ(db.get_sibling_path(1), expected01);
+    EXPECT_EQ(db.get_sibling_path(2), expected02);
+    EXPECT_EQ(db.get_sibling_path(3), expected03);
+    EXPECT_EQ(db.root(), root);
+}
diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_leaf.hpp b/cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_leaf.hpp
@@ -44,9 +44,10 @@ inline std::pair<size_t, bool> find_closest_leaf(std::vector<nullifier_leaf> con
 {
     std::vector<uint256_t> diff;
     bool repeated = false;
+    auto new_value_ = uint256_t(new_value);
+
     for (size_t i = 0; i < leaves_.size(); i++) {
         auto leaf_value_ = uint256_t(leaves_[i].value);
-        auto new_value_ = uint256_t(new_value);
         if (leaf_value_ > new_value_) {
             diff.push_back(leaf_value_);
         } else if (leaf_value_ == new_value_) {

diff --git a/cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_memory_tree.hpp b/cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_memory_tree.hpp
@@ -79,8 +79,9 @@ class NullifierMemoryTree : public MemoryTree {
 
     const std::vector<barretenberg::fr>& get_hashes() { return hashes_; }
     const std::vector<nullifier_leaf>& get_leaves() { return leaves_; }
+    const nullifier_leaf& get_leaf(size_t index) { return leaves_[index]; }
 
-  private:
+  protected:
     using MemoryTree::depth_;
     using MemoryTree::hashes_;
     using MemoryTree::root_;