fix: small nullifier tree issue

cpp/src/aztec3/circuits/rollup/base/.test.cpp

@@ -1,8 +1,3 @@
-// #include <barretenberg/common/serialize.hpp>
-// #include <barretenberg/stdlib/types/types.hpp>
-// #include <aztec3/oracle/oracle.hpp>
-// #include <aztec3/circuits/apps/oracle_wrapper.hpp>
-// #include <barretenberg/numeric/random/engine.hpp>
 #include "aztec3/circuits/abis/append_only_tree_snapshot.hpp"
 #include "aztec3/circuits/abis/membership_witness.hpp"
 #include "aztec3/circuits/abis/private_kernel/new_contract_data.hpp"
@@ -36,12 +31,9 @@
 #include <aztec3/circuits/abis/private_kernel/constant_data.hpp>
 #include <aztec3/circuits/abis/private_kernel/old_tree_roots.hpp>
 #include <aztec3/circuits/abis/private_kernel/globals.hpp>
-// #include <aztec3/circuits/abis/private_kernel/private_inputs.hpp>
-// #include <aztec3/circuits/abis/private_kernel/private_inputs.hpp>
 
 #include <aztec3/circuits/apps/function_execution_context.hpp>
 
-// #include <aztec3/circuits/mock/mock_circuit.hpp>
 #include <aztec3/circuits/mock/mock_kernel_circuit.hpp>
 
 #include <barretenberg/common/map.hpp>
@@ -418,6 +410,12 @@ TEST_F(base_rollup_tests, new_nullifier_tree_all_larger)
     AppendOnlyTreeSnapshot<NT> nullifier_tree_start_snapshot = std::get<1>(inputs_and_snapshots);
     AppendOnlyTreeSnapshot<NT> nullifier_tree_end_snapshot = std::get<2>(inputs_and_snapshots);
 
+    // info("testing inputs");
+    // info(testing_inputs.kernel_data[0].public_inputs.end.new_nullifiers);
+    // info(testing_inputs.kernel_data[1].public_inputs.end.new_nullifiers);
+    // info("should only be one non zero");
+    // info(testing_inputs.low_nullifier_membership_witness);
+
     /**
      * RUN
      */
@@ -469,26 +467,69 @@ TEST_F(base_rollup_tests, new_nullifier_tree_sparse)
     ASSERT_EQ(outputs.end_nullifier_tree_snapshot, nullifier_tree_end_snapshot);
 }
 
-TEST_F(base_rollup_tests, new_nullifier_tree_sparse_attack)
+TEST_F(base_rollup_tests, nullifier_tree_regression)
 {
-    // @todo THIS SHOULD NOT BE PASSING. The circuit should fail with an assert as we are trying to double-spend.
-    /**
-     * DESCRIPTION
-     */
+    // Regression test caught when testing the typescript nullifier tree implementation
     DummyComposer composer = DummyComposer();
     BaseRollupInputs empty_inputs = dummy_base_rollup_inputs_with_vk_proof();
 
-    std::array<fr, KERNEL_NEW_NULLIFIERS_LENGTH* 2> new_nullifiers = { 11, 0, 11, 0, 0, 0, 0, 0 };
+    // This test runs after some data has already been inserted into the tree
+    // TODO: explain data
+    std::vector<fr> initial_values(23, 0);
+    for (size_t i = 0; i < 7; i++) {
+        initial_values[i] = i + 1;
+    }
+    // Note these are hex representations
+    initial_values[7] = uint256_t("2bb9aa4a22a6ae7204f2c67abaab59cead6558cde4ee25ce3464704cb2e38136");
+    initial_values[8] = uint256_t("16a732095298ccca828c4d747813f8bd46e188079ed17904e2c9de50760833c8");
+
+    std::array<fr, KERNEL_NEW_NULLIFIERS_LENGTH* 2> new_nullifiers = { 0 };
+    new_nullifiers[0] = uint256_t("16da4f27fb78de7e0db4c5a04b569bc46382c5f471da2f7d670beff1614e0118"),
+    new_nullifiers[1] = uint256_t("26ab07ce103a55e29f11478eaa36cebd10c4834b143a7debcc7ef53bfdb547dd");
+
     std::tuple<BaseRollupInputs, AppendOnlyTreeSnapshot<NT>, AppendOnlyTreeSnapshot<NT>> inputs_and_snapshots =
-        utils::generate_nullifier_tree_testing_values(empty_inputs, new_nullifiers, 1);
+        utils::generate_nullifier_tree_testing_values(empty_inputs, new_nullifiers, initial_values);
     BaseRollupInputs testing_inputs = std::get<0>(inputs_and_snapshots);
+    AppendOnlyTreeSnapshot<NT> nullifier_tree_start_snapshot = std::get<1>(inputs_and_snapshots);
+    AppendOnlyTreeSnapshot<NT> nullifier_tree_end_snapshot = std::get<2>(inputs_and_snapshots);
+
+    // We now have the correct tree - can debug in circuit
+    /**
+     * RUN
+     */
 
-    // Run the circuit (SHOULD FAIL WITH AN ASSERT INSTEAD OF THIS!)
+    // Run the circuit
     BaseOrMergeRollupPublicInputs outputs =
         aztec3::circuits::rollup::native_base_rollup::base_rollup_circuit(composer, testing_inputs);
+
+    /**
+     * ASSERT
+     */
+    // Start state
+    ASSERT_EQ(outputs.start_nullifier_tree_snapshot, nullifier_tree_start_snapshot);
+
+    // End state
+    ASSERT_EQ(outputs.end_nullifier_tree_snapshot, nullifier_tree_end_snapshot);
 }
 
-// TEST_F(base_rollup_tests, new_commitments_tree) {}
+// TEST_F(base_rollup_tests, new_nullifier_tree_sparse_attack)
+// {
+//     // @todo THIS SHOULD NOT BE PASSING. The circuit should fail with an assert as we are trying to double-spend.
+//     /**
+//      * DESCRIPTION
+//      */
+
+//     BaseRollupInputs empty_inputs = dummy_base_rollup_inputs_with_vk_proof();
+
+//     std::array<fr, KERNEL_NEW_NULLIFIERS_LENGTH* 2> new_nullifiers = { 11, 0, 11, 0, 0, 0, 0, 0 };
+//     std::tuple<BaseRollupInputs, AppendOnlyTreeSnapshot<NT>, AppendOnlyTreeSnapshot<NT>> inputs_and_snapshots =
+//         utils::generate_nullifier_tree_testing_values(empty_inputs, new_nullifiers, 1);
+//     BaseRollupInputs testing_inputs = std::get<0>(inputs_and_snapshots);
+
+//     // Run the circuit (SHOULD FAIL WITH AN ASSERT INSTEAD OF THIS!)
+//     BaseOrMergeRollupPublicInputs outputs =
+//         aztec3::circuits::rollup::native_base_rollup::base_rollup_circuit(testing_inputs);
+// }
 
 TEST_F(base_rollup_tests, empty_block_calldata_hash)
 {

cpp/src/aztec3/circuits/rollup/base/native_base_rollup_circuit.cpp

@@ -273,7 +273,7 @@ AppendOnlySnapshot check_nullifier_tree_non_membership_and_insert_to_tree(DummyC
     // 2. If we receive the 0 nullifier leaf (where all values are 0, we skip insertion and leave a sparse subtree)
 
     // New nullifier subtree
-    std::array<NullifierLeaf, KERNEL_NEW_NULLIFIERS_LENGTH * 2> nullifier_leaves;
+    std::array<NullifierLeaf, KERNEL_NEW_NULLIFIERS_LENGTH * 2> nullifier_insertion_subtree;
 
     // This will update on each iteration
     auto current_nullifier_tree_root = baseRollupInputs.start_nullifier_tree_snapshot.root;
@@ -291,12 +291,22 @@ AppendOnlySnapshot check_nullifier_tree_non_membership_and_insert_to_tree(DummyC
 
             // Witness containing index and path
             auto nullifier_index = 4 * i + j;
+            info("nullifier index", nullifier_index);
+
+            // info(nullifier_insertion_subtree);
+
             auto witness = baseRollupInputs.low_nullifier_membership_witness[nullifier_index];
             // Preimage of the lo-index required for a non-membership proof
             auto low_nullifier_preimage = baseRollupInputs.low_nullifier_leaf_preimages[nullifier_index];
             // Newly created nullifier
             auto nullifier = new_nullifiers[j];
 
+            // preimage
+            info("preimage");
+            info("value: ", low_nullifier_preimage.leaf_value);
+            info("next index: ", low_nullifier_preimage.next_index);
+            info("next value: ", low_nullifier_preimage.next_value);
+
             // TODO: reason about this more strongly, can this cause issues?
             if (nullifier != 0) {
 
@@ -313,13 +323,19 @@ AppendOnlySnapshot check_nullifier_tree_non_membership_and_insert_to_tree(DummyC
                     // TODO: this is a hack, and insecure, we need to fix this
                     bool matched = false;
                     for (size_t k = 0; k < nullifier_index; k++) {
-                        if ((uint256_t(nullifier_leaves[k].nextValue) > uint256_t(nullifier) &&
-                             uint256_t(nullifier_leaves[k].value) < uint256_t(nullifier)) ||
-                            (nullifier_leaves[k].nextValue == 0 && nullifier_leaves[k].nextIndex == 0)) {
+                        if ((uint256_t(nullifier_insertion_subtree[k].nextValue) > uint256_t(nullifier) &&
+                             uint256_t(nullifier_insertion_subtree[k].value) < uint256_t(nullifier)) ||
+                            (nullifier_insertion_subtree[k].nextValue == 0 &&
+                             nullifier_insertion_subtree[k].nextIndex == 0)) {
 
                             matched = true;
-                            nullifier_leaves[k].nextIndex = new_index;
-                            nullifier_leaves[k].nextValue = nullifier;
+                            // Update pointers
+                            new_nullifier_leaf.nextIndex = nullifier_insertion_subtree[k].nextIndex;
+                            new_nullifier_leaf.nextValue = nullifier_insertion_subtree[k].nextValue;
+
+                            // Update child
+                            nullifier_insertion_subtree[k].nextIndex = new_index;
+                            nullifier_insertion_subtree[k].nextValue = nullifier;
                         }
                     }
                     // if not matched, our subtree will misformed - we must reject
@@ -359,15 +375,15 @@ AppendOnlySnapshot check_nullifier_tree_non_membership_and_insert_to_tree(DummyC
                         updated_low_nullifier.hash(), witness.leaf_index, witness.sibling_path);
                 }
 
-                nullifier_leaves[nullifier_index] = new_nullifier_leaf;
+                nullifier_insertion_subtree[nullifier_index] = new_nullifier_leaf;
             } else {
                 // 0 case
                 NullifierLeaf new_nullifier_leaf = {
                     .value = 0,
                     .nextIndex = 0,
                     .nextValue = 0,
                 };
-                nullifier_leaves[nullifier_index] = new_nullifier_leaf;
+                nullifier_insertion_subtree[nullifier_index] = new_nullifier_leaf;
             }
 
             // increment insertion index
@@ -377,7 +393,7 @@ AppendOnlySnapshot check_nullifier_tree_non_membership_and_insert_to_tree(DummyC
 
     // Create new nullifier subtree to insert into the whole nullifier tree
     auto nullifier_sibling_path = baseRollupInputs.new_nullifiers_subtree_sibling_path;
-    auto nullifier_subtree_root = create_nullifier_subtree(nullifier_leaves);
+    auto nullifier_subtree_root = create_nullifier_subtree(nullifier_insertion_subtree);
 
     // Calculate the new root
     // We are inserting a subtree rather than a full tree here

cpp/src/aztec3/circuits/rollup/base/nullifier_tree_testing_harness.cpp

@@ -1,6 +1,7 @@
 #include "nullifier_tree_testing_harness.hpp"
 #include <barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_memory_tree.hpp>
 #include <barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_leaf.hpp>
+#include <cstdint>
 #include <tuple>
 
 using NullifierMemoryTree = proof_system::plonk::stdlib::merkle_tree::NullifierMemoryTree;
@@ -10,13 +11,76 @@ NullifierMemoryTreeTestingHarness::NullifierMemoryTreeTestingHarness(size_t dept
     : NullifierMemoryTree(depth)
 {}
 
+// Check for a larger value in an array
+bool check_has_lesser_value(std::vector<fr> const& values, fr const& value)
+{
+    // Must perform comparisons on integers
+    uint256_t value_as_uint = uint256_t(value);
+    for (auto const& v : values) {
+        if (uint256_t(v) < value_as_uint) {
+            return true;
+        }
+    }
+    return false;
+}
+
+// TODO: test
+fr NullifierMemoryTreeTestingHarness::append_value(fr const& value)
+{
+    // If the value is 0, then we force insert the value to increase the size of the tree
+    // TODO: this is a hack
+    if (value == 0) {
+        nullifier_leaf new_leaf = { .value = 0, .nextIndex = 0, .nextValue = 0 };
+        auto new_leaf_hash = new_leaf.hash();
+        leaves_.push_back(new_leaf);
+        size_t new_leaf_index = leaves_.size() - 1;
+        auto root = update_element(new_leaf_index, new_leaf_hash);
+        return root;
+    }
+
+    // Find the leaf with the value closest and less than `value`
+    size_t current;
+    bool is_already_present;
+    std::tie(current, is_already_present) = find_closest_leaf(leaves_, value);
+
+    nullifier_leaf new_leaf = { .value = value,
+                                .nextIndex = leaves_[current].nextIndex,
+                                .nextValue = leaves_[current].nextValue };
+    if (!is_already_present) {
+        // Update the current leaf to point it to the new leaf
+        leaves_[current].nextIndex = leaves_.size();
+        leaves_[current].nextValue = value;
+
+        // Insert the new leaf with (nextIndex, nextValue) of the current leaf
+        leaves_.push_back(new_leaf);
+    }
+
+    // Update the old leaf in the tree
+    auto old_leaf_hash = leaves_[current].hash();
+    size_t old_leaf_index = current;
+    auto root = update_element(old_leaf_index, old_leaf_hash);
+
+    // Insert the new leaf in the tree
+    auto new_leaf_hash = new_leaf.hash();
+
+    size_t new_leaf_index = is_already_present ? old_leaf_index : leaves_.size() - 1;
+    root = update_element(new_leaf_index, new_leaf_hash);
+
+    return root;
+}
+
 // handle synthetic membership assertions
 std::tuple<std::vector<nullifier_leaf>, std::vector<std::vector<fr>>, std::vector<uint32_t>>
-NullifierMemoryTreeTestingHarness::circuit_prep_batch_insert(std::vector<fr> const& values,
-                                                             std::vector<fr> const& insertion_locations)
+NullifierMemoryTreeTestingHarness::circuit_prep_batch_insert(std::vector<fr> const& values)
 {
+    // Start insertion index
+    fr start_insertion_index = this->size();
+
+    info(values);
+
     // Low nullifiers
     std::vector<nullifier_leaf> low_nullifiers;
+    std::vector<nullifier_leaf> pending_insertion_tree;
 
     // Low nullifier sibling paths
     std::vector<std::vector<fr>> sibling_paths;
@@ -25,31 +89,80 @@ NullifierMemoryTreeTestingHarness::circuit_prep_batch_insert(std::vector<fr> con
     std::vector<uint32_t> low_nullifier_indexes;
 
     // Keep track of the currently touched nodes while updating
-    std::set<size_t> touched_nodes;
+    std::map<size_t, std::vector<fr>> touched_nodes;
+
+    // Keep track of 0 values
+    std::vector<fr> empty_sp(depth_, 0);
+    nullifier_leaf empty_leaf = { 0, 0, 0 };
+    uint32_t 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 value = values[i];
-        auto insertion_index = uint256_t(insertion_locations[i]);
+        auto new_value = values[i];
+        auto insertion_index = start_insertion_index + i;
 
         size_t current;
         bool is_already_present;
+        std::tie(current, is_already_present) = find_closest_leaf(leaves_, new_value);
 
-        std::tie(current, is_already_present) = find_closest_leaf(leaves_, value);
+        info("new value: ", new_value);
+        info("current");
+        info(current);
+
+        // If the inserted value is 0, then we ignore and provide a dummy low nullifier
+        if (new_value == 0) {
+            sibling_paths.push_back(empty_sp);
+            low_nullifier_indexes.push_back(empty_index);
+            low_nullifiers.push_back(empty_leaf);
+        }
 
         // 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
-        if (touched_nodes.contains(current)) {
-            std::vector<fr> sp(depth_, 0);
-            auto empty_leaf = nullifier_leaf{ 0, 0, 0 };
+        auto prev_nodes = touched_nodes.find(current);
+
+        bool has_lesser_value = false;
+        if (prev_nodes != touched_nodes.end()) {
+            has_lesser_value = check_has_lesser_value(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 (current == 0 || has_lesser_value) {
+        if (has_lesser_value) {
+
+            for (size_t j = 0; j < pending_insertion_tree.size(); ++j) {
+                // Skip checking empty values
+                if (pending_insertion_tree[j].value == 0) {
+                    continue;
+                }
+
+                if (pending_insertion_tree[j].value < new_value &&
+                    (pending_insertion_tree[j].nextValue > new_value || pending_insertion_tree[j].nextValue == 0)) {
+                    // Add a new pending low nullifier for this value
+                    nullifier_leaf new_leaf = { .value = new_value,
+                                                .nextIndex = pending_insertion_tree[j].nextIndex,
+                                                .nextValue = pending_insertion_tree[j].nextValue };
+                    pending_insertion_tree.push_back(new_leaf);
+
+                    // Update the pending low nullifier to point at the new value
+                    pending_insertion_tree[j].nextIndex = insertion_index;
+                    pending_insertion_tree[j].nextValue = new_value;
+
+                    break;
+                }
+            }
 
             // empty low nullifier
-            sibling_paths.push_back(sp);
-            low_nullifier_indexes.push_back(0);
+            sibling_paths.push_back(empty_sp);
+            low_nullifier_indexes.push_back(empty_index);
             low_nullifiers.push_back(empty_leaf);
         } else {
-            touched_nodes.insert(current);
+            // Update the touched mapping
+            if (prev_nodes == touched_nodes.end()) {
+                std::vector<fr> new_touched_values = { new_value };
+                touched_nodes[current] = new_touched_values;
+            } else {
+                prev_nodes->second.push_back(new_value);
+            }
 
             nullifier_leaf low_nullifier = leaves_[current];
             std::vector<fr> sibling_path = this->get_sibling_path(current);
@@ -61,7 +174,7 @@ NullifierMemoryTreeTestingHarness::circuit_prep_batch_insert(std::vector<fr> con
             // Update the current low nullifier
             nullifier_leaf new_leaf = { .value = low_nullifier.value,
                                         .nextIndex = insertion_index,
-                                        .nextValue = value };
+                                        .nextValue = new_value };
 
             // Update the old leaf in the tree
             update_element(current, new_leaf.hash());