refactor: Exposing more of nullifier tree state

cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.cpp

@@ -106,8 +106,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;
@@ -133,6 +133,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() == 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<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() == 65);
+            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;

cpp/src/barretenberg/stdlib/merkle_tree/merkle_tree.hpp

@@ -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;

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

cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_tree.hpp

@@ -25,12 +25,19 @@ template <typename Store> class NullifierTree : public MerkleTree<Store> {
 
     fr update_element(fr const& value);
 
-  private:
+  protected:
     using MerkleTree<Store>::update_element;
     using MerkleTree<Store>::get_element;
     using MerkleTree<Store>::compute_zero_path_hash;
 
-  private:
+    // const std::vector<barretenberg::fr>& 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<WrappedNullifierLeaf>& get_leaves() { return leaves; }
+
     using MerkleTree<Store>::store_;
     using MerkleTree<Store>::zero_hashes_;
     using MerkleTree<Store>::depth_;