refactor: Move nullifier tree batch insertion logic into bberg

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

@@ -13,13 +13,9 @@ struct nullifier_leaf {
     index_t nextIndex;
     fr nextValue;
 
-    bool operator==(nullifier_leaf const&) const = default;
+    static nullifier_leaf empty() { return nullifier_leaf{ fr::zero(), 0, fr::zero() }; }
 
-    std::ostream& operator<<(std::ostream& os)
-    {
-        os << "value = " << value << "\nnextIdx = " << nextIndex << "\nnextVal = " << nextValue;
-        return os;
-    }
+    bool operator==(nullifier_leaf const&) const = default;
 
     void read(uint8_t const*& it)
     {
@@ -38,8 +34,17 @@ struct nullifier_leaf {
     }
 
     barretenberg::fr hash() const { return stdlib::merkle_tree::hash_multiple_native({ value, nextIndex, nextValue }); }
+
+    bool is_empty() const { return value == 0 && nextIndex == 0 && nextValue == 0; }
 };
 
+inline std::ostream& operator<<(std::ostream& os, nullifier_leaf const& leaf)
+{
+    return os << "value = " << leaf.value << "\n"
+              << "nextIdx = " << leaf.nextIndex << "\n"
+              << "nextVal = " << leaf.nextValue;
+}
+
 /**
  * @brief Wrapper for the Nullifier leaf class that allows for 0 values
  *
@@ -78,7 +83,14 @@ class WrappedNullifierLeaf {
      *
      * @param value
      */
-    void set(nullifier_leaf value) { data.emplace(value); }
+    void set(nullifier_leaf value)
+    {
+        if (value.is_empty()) {
+            data = std::nullopt;
+        } else {
+            data.emplace(value);
+        }
+    }
 
     /**
      * @brief Return the hash of the wrapped object, other return the zero hash of 0
@@ -99,6 +111,17 @@ class WrappedNullifierLeaf {
     std::optional<nullifier_leaf> data;
 };
 
+inline std::ostream& operator<<(std::ostream& os, WrappedNullifierLeaf const& leaf)
+{
+    if (!leaf.has_value()) {
+        return os << "value = 0\n"
+                  << "nextIdx = 0\n"
+                  << "nextVal = 0";
+    } else {
+        return os << leaf.unwrap();
+    }
+}
+
 inline std::pair<size_t, bool> find_closest_leaf(std::vector<WrappedNullifierLeaf> const& leaves_, fr const& new_value)
 {
     std::vector<uint256_t> diff;

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

@@ -72,6 +72,180 @@ fr NullifierMemoryTree::update_element(fr const& value)
     return root;
 }
 
+// Check for a larger value in an array
+bool check_has_less_than(std::vector<fr> const& values, fr const& value)
+{
+    // Must perform comparisons on integers
+    auto const value_as_uint = uint256_t(value);
+    for (auto const& v : values) {
+        if (uint256_t(v) < value_as_uint) {
+            return true;
+        }
+    }
+    return false;
+}
+
+/**
+ * @brief Insert a batch of values into the tree, returning the low nullifiers membership information (leaf, sibling
+ * path, index)
+ *
+ * Special Considerations
+ *
+ * Short algorithm description:
+ * When batch inserting values into the tree, we first update their "low_nullifiers" (the node that will insert to the
+ * inserted value).
+ *  - For each low nullifier that we update, we need to perform a membership check against the current root of the tree.
+ *  - Once membership is confirmed, we can update the leaf, then use the same sibling path to update the root.
+ *  - The next membership check will be against the new root.
+ *
+ * If the low nullifier for a value exists within the batch being inserted, then we cannot perform a membership check,
+ * as it has not yet been inserted! In this case we provide an all 0 sibling path and all 0 low nullifier, all
+ * corresponding aztec circuits account for this case.
+ *
+ * A description of the algorithm can be found here:
+ * https://colab.research.google.com/drive/1A0gizduSi4FIiIJZ8OylwIpO9-OTqV-R
+ *
+ * @param values - An array of values to insert into the tree.
+ * @return std::vector<LowLeafWitnessData>
+ */
+LowLeafWitnessData NullifierMemoryTree::batch_insert(std::vector<fr> const& values)
+{
+    // Start insertion index
+    size_t const start_insertion_index = this->size();
+
+    // Low nullifiers
+    auto values_size = values.size();
+    std::vector<nullifier_leaf> low_nullifiers(values_size);
+    std::vector<nullifier_leaf> pending_insertion_tree(values_size);
+
+    // Low nullifier sibling paths
+    std::vector<std::vector<fr>> sibling_paths(values_size);
+
+    // Low nullifier indexes
+    std::vector<uint32_t> low_nullifier_indexes(values_size);
+
+    // Keep track of the currently touched nodes while updating
+    std::map<size_t, std::vector<fr>> touched_nodes;
+
+    // Keep track of 0 values
+    std::vector<fr> const empty_sp(depth_, 0);
+    auto const empty_leaf = nullifier_leaf::empty();
+    uint32_t const empty_index = 0;
+
+    // Find the leaf with the value closest and less than `value` for each value
+    for (size_t i = 0; i < values.size(); ++i) {
+
+        auto new_value = values[i];
+        auto insertion_index = start_insertion_index + i;
+
+        size_t current = 0;
+        bool is_already_present = false;
+        std::tie(current, is_already_present) = find_closest_leaf(leaves_, new_value);
+
+        // If the inserted value is 0, then we ignore and provide a dummy low nullifier
+        if (new_value == 0) {
+            sibling_paths[i] = empty_sp;
+            low_nullifier_indexes[i] = empty_index;
+            low_nullifiers[i] = empty_leaf;
+            pending_insertion_tree[i] = empty_leaf;
+            continue;
+        }
+
+        // If the low_nullifier node has been touched this sub tree insertion, we provide a dummy sibling path
+        // It will be up to the circuit to check if the included node is valid vs the other nodes that have been
+        // inserted before it If it has not been touched, we provide a sibling path then update the nodes pointers
+        auto prev_nodes = touched_nodes.find(current);
+
+        bool has_less_than = false;
+        if (prev_nodes != touched_nodes.end()) {
+            has_less_than = check_has_less_than(prev_nodes->second, new_value);
+        }
+        // If there is a lower value in the tree, we need to check the current low nullifiers for one that can be used
+        if (has_less_than) {
+            for (size_t j = 0; j < pending_insertion_tree.size(); ++j) {
+
+                nullifier_leaf& pending = pending_insertion_tree[j];
+                // Skip checking empty values
+                if (pending.is_empty()) {
+                    continue;
+                }
+
+                if (uint256_t(pending.value) < uint256_t(new_value) &&
+                    (uint256_t(pending.nextValue) > uint256_t(new_value) || pending.nextValue == fr::zero())) {
+
+                    // Add a new pending low nullifier for this value
+                    nullifier_leaf const current_low_leaf = { .value = new_value,
+                                                              .nextIndex = pending_insertion_tree[j].nextIndex,
+                                                              .nextValue = pending_insertion_tree[j].nextValue };
+
+                    pending_insertion_tree[i] = current_low_leaf;
+
+                    // Update the pending low nullifier to point at the new value
+                    pending.nextValue = new_value;
+                    pending.nextIndex = insertion_index;
+
+                    break;
+                }
+            }
+
+            // add empty low nullifier
+            sibling_paths[i] = empty_sp;
+            low_nullifier_indexes[i] = empty_index;
+            low_nullifiers[i] = empty_leaf;
+        } else {
+            // Update the touched mapping
+            if (prev_nodes == touched_nodes.end()) {
+                std::vector<fr> const new_touched_values = { new_value };
+                touched_nodes[current] = new_touched_values;
+            } else {
+                prev_nodes->second.push_back(new_value);
+            }
+
+            nullifier_leaf const low_nullifier = leaves_[current].unwrap();
+            std::vector<fr> const sibling_path = this->get_sibling_path(current);
+
+            sibling_paths[i] = sibling_path;
+            low_nullifier_indexes[i] = static_cast<uint32_t>(current);
+            low_nullifiers[i] = low_nullifier;
+
+            // TODO(SEAN): rename this and new leaf
+            nullifier_leaf insertion_leaf = { .value = new_value,
+                                              .nextIndex = low_nullifier.nextIndex,
+                                              .nextValue = low_nullifier.nextValue };
+            pending_insertion_tree[i] = insertion_leaf;
+
+            // Update the current low nullifier
+            nullifier_leaf const new_leaf = { .value = low_nullifier.value,
+                                              .nextIndex = insertion_index,
+                                              .nextValue = new_value };
+
+            // Update the old leaf in the tree
+            // update old value in tree
+            update_element_in_place(current, new_leaf);
+        }
+    }
+
+    // resize leaves array
+    this->leaves_.resize(this->leaves_.size() + pending_insertion_tree.size());
+    for (size_t i = 0; i < pending_insertion_tree.size(); ++i) {
+        nullifier_leaf const pending = pending_insertion_tree[i];
+
+        // Update the old leaf in the tree
+        // update old value in tree
+        update_element_in_place(start_insertion_index + i, pending);
+    }
+
+    // Return tuple of low nullifiers and sibling paths
+    return std::make_tuple(low_nullifiers, sibling_paths, low_nullifier_indexes);
+}
+
+// Update the value of a leaf in place
+fr NullifierMemoryTree::update_element_in_place(size_t index, const nullifier_leaf& leaf)
+{
+    this->leaves_[index].set(leaf);
+    return update_element(index, leaf.hash());
+}
+
 } // namespace merkle_tree
 } // namespace stdlib
 } // namespace proof_system::plonk

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

@@ -7,6 +7,9 @@ namespace proof_system::plonk {
 namespace stdlib {
 namespace merkle_tree {
 
+// tuple(nullifier_leaf, sibling_path, index) // utility alias
+using LowLeafWitnessData = std::tuple<std::vector<nullifier_leaf>, std::vector<std::vector<fr>>, std::vector<uint32_t>>;
+
 using namespace barretenberg;
 
 /**
@@ -76,15 +79,22 @@ class NullifierMemoryTree : public MemoryTree {
     using MemoryTree::root;
     using MemoryTree::update_element;
 
-    fr update_element(fr const& value);
-
+    // Inspectors
+    size_t size() { return leaves_.size(); }
+    fr total_size() const { return total_size_; }
+    fr depth() const { return depth_; }
     const std::vector<barretenberg::fr>& get_hashes() { return hashes_; }
     const WrappedNullifierLeaf get_leaf(size_t index)
     {
         return (index < leaves_.size()) ? leaves_[index] : WrappedNullifierLeaf::zero();
     }
     const std::vector<WrappedNullifierLeaf>& get_leaves() { return leaves_; }
 
+    // Mutators
+    fr update_element(fr const& value);
+    fr update_element_in_place(size_t index, const nullifier_leaf& value);
+    LowLeafWitnessData batch_insert(std::vector<fr> const& values);
+
   protected:
     using MemoryTree::depth_;
     using MemoryTree::hashes_;

cpp/src/barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_memory_tree.test.cpp

@@ -327,6 +327,110 @@ TEST(crypto_nullifier_tree, test_nullifier_memory_appending_zero)
     EXPECT_EQ(tree.get_hash_path(6), expected);
     EXPECT_EQ(tree.get_hash_path(7), expected);
 }
+
+TEST(crypto_nullifier_tree, test_nullifier_memory_tree_batch_insert)
+{
+    // Create a depth-3 indexed merkle tree
+    constexpr size_t depth = 3;
+    NullifierMemoryTree tree(depth);
+
+    /**
+     * Initial State:
+     *
+     *  index     0       1       2       3        4       5       6       7
+     *  ---------------------------------------------------------------------
+     *  val       0       30      10      20       0       0       0       0
+     *  nextIdx   2       0       3       1        0       0       0       0
+     *  nextVal   10      0       20      30       0       0       0       0
+     */
+    tree.update_element(30);
+    tree.update_element(10);
+    tree.update_element(20);
+
+    EXPECT_EQ(tree.get_leaves().size(), 4);
+    EXPECT_EQ(tree.get_leaves()[0].hash(), WrappedNullifierLeaf({ 0, 2, 10 }).hash());
+    EXPECT_EQ(tree.get_leaves()[1].hash(), WrappedNullifierLeaf({ 30, 0, 0 }).hash());
+    EXPECT_EQ(tree.get_leaves()[2].hash(), WrappedNullifierLeaf({ 10, 3, 20 }).hash());
+    EXPECT_EQ(tree.get_leaves()[3].hash(), WrappedNullifierLeaf({ 20, 1, 30 }).hash());
+
+    // Perform batch insertion
+    // What should we expect?
+    // 50's low nullifier should be the value 30, index 1
+    // 25's low nullifier should be the value 20, index 3, sibling path will be such that the low nullifier for 30 has
+    // been updated! 80's low nullifier should be the value 30 - this is already used, so it will return 0 75's low
+    // nullifier should be the value 30 - this is already used, so it will return 0
+    std::vector<fr> batch_values = { 50, 25, 80, 75 };
+    LowLeafWitnessData low_leaf_witnesses = tree.batch_insert(batch_values);
+
+    // Destructure return type
+    std::vector<nullifier_leaf> low_leaves = std::get<0>(low_leaf_witnesses);
+    std::vector<std::vector<fr>> low_leaf_sibling_paths = std::get<1>(low_leaf_witnesses);
+    std::vector<uint32_t> low_leaf_witness_indexes = std::get<2>(low_leaf_witnesses);
+
+    /**
+     * State after batch insertion:
+     *
+     *  index     0       1       2       3        4       5       6       7
+     *  ---------------------------------------------------------------------
+     *  val       0       30      10      20       50      25      80      75
+     *  nextIdx   2       4       3       5        7       1       0       6
+     *  nextVal   10      50      20      25       75      30      0       80
+     */
+    // Check that insertions have been performed correctly
+    auto leaves = tree.get_leaves();
+
+    EXPECT_EQ(leaves.size(), 8);
+    EXPECT_EQ(leaves[0].hash(), WrappedNullifierLeaf({ 0, 2, 10 }).hash());
+    EXPECT_EQ(leaves[1].hash(), WrappedNullifierLeaf({ 30, 4, 50 }).hash());
+    EXPECT_EQ(leaves[2].hash(), WrappedNullifierLeaf({ 10, 3, 20 }).hash());
+    EXPECT_EQ(leaves[3].hash(), WrappedNullifierLeaf({ 20, 5, 25 }).hash());
+    EXPECT_EQ(leaves[4].hash(), WrappedNullifierLeaf({ 50, 7, 75 }).hash());
+    EXPECT_EQ(leaves[5].hash(), WrappedNullifierLeaf({ 25, 1, 30 }).hash());
+    EXPECT_EQ(leaves[6].hash(), WrappedNullifierLeaf({ 80, 0, 0 }).hash());
+    EXPECT_EQ(leaves[7].hash(), WrappedNullifierLeaf({ 75, 6, 80 }).hash());
+
+    // Check the correct low leaf witness indexes have been returned.
+    // 50's low index is 30, 25's low index is 20, 80's low index is 30 (as it is being revisited it is 0),
+    // 75's low index is 30 (as it being revisited it is 0 )
+    auto expected_low_leaf_witness_indexes = std::vector<uint32_t>{ 1, 3, 0, 0 };
+    EXPECT_EQ(low_leaf_witness_indexes, expected_low_leaf_witness_indexes);
+
+    // Check the low nullifier for the first insertion
+    EXPECT_EQ(low_leaves[0], WrappedNullifierLeaf({ 30, 0, 0 }));
+
+    // Below, nodes are indexed by their layer, then their index e.g. e12 -> layer 1, index 2
+    auto zero_leaf_hash = WrappedNullifierLeaf().hash();
+    auto e00 = WrappedNullifierLeaf({ 0, 2, 10 }).hash();
+    auto e02 = WrappedNullifierLeaf({ 10, 3, 20 }).hash();
+    auto e03 = WrappedNullifierLeaf({ 20, 1, 30 }).hash();
+
+    auto e11 = hash_pair_native(e02, e03);
+    auto e12 = hash_pair_native(zero_leaf_hash, zero_leaf_hash);
+    auto e13 = hash_pair_native(zero_leaf_hash, zero_leaf_hash);
+
+    auto e21 = hash_pair_native(e12, e13);
+
+    auto expected_low_leaf_sibling_path = std::vector<fr>{ e00, e11, e21 };
+    EXPECT_EQ(low_leaf_sibling_paths[0], expected_low_leaf_sibling_path);
+
+    // Check the low nullifier for the second insertion
+    EXPECT_EQ(low_leaves[1], nullifier_leaf({ 20, 1, 30 }));
+
+    // Update sibling paths after first low nullifier update below.
+    auto updated_e01 = WrappedNullifierLeaf({ 30, 4, 50 }).hash();
+    auto updated_e10 = hash_pair_native(e00, updated_e01);
+
+    auto expected_second_low_leaf_sibling_path = std::vector<fr>{ e02, updated_e10, e21 };
+    EXPECT_EQ(low_leaf_sibling_paths[1], expected_second_low_leaf_sibling_path);
+
+    // Check the remaining insertions are zero paths
+    auto zero_path = std::vector<fr>{ fr::zero(), fr::zero(), fr::zero() };
+    EXPECT_EQ(low_leaves[2], nullifier_leaf());
+    EXPECT_EQ(low_leaf_sibling_paths[2], zero_path);
+    EXPECT_EQ(low_leaves[3], nullifier_leaf());
+    EXPECT_EQ(low_leaf_sibling_paths[3], zero_path);
+}
+
 TEST(crypto_nullifier_tree, test_nullifier_tree)
 {
     // Create a depth-8 indexed merkle tree