chore: Refactor indexed tree to use traits (#10361)

noir-projects/noir-protocol-circuits/crates/rollup-lib/src/base/components/nullifier_tree.nr

@@ -8,7 +8,6 @@ use dep::types::{
         NULLIFIER_TREE_HEIGHT,
     },
     merkle_tree::{indexed_tree, MembershipWitness},
-    utils::field::full_field_less_than,
 };
 
 pub(crate) fn nullifier_tree_batch_insert(
@@ -20,50 +19,13 @@ pub(crate) fn nullifier_tree_batch_insert(
     nullifier_predecessor_preimages: [NullifierLeafPreimage; MAX_NULLIFIERS_PER_TX],
     nullifier_predecessor_membership_witnesses: [MembershipWitness<NULLIFIER_TREE_HEIGHT>; MAX_NULLIFIERS_PER_TX],
 ) -> AppendOnlyTreeSnapshot {
-    indexed_tree::batch_insert(
+    indexed_tree::batch_insert::<_, _, _, _, NULLIFIER_SUBTREE_HEIGHT, NULLIFIER_TREE_HEIGHT>(
         start_snapshot,
         nullifiers,
         sorted_nullifiers,
         sorted_nullifiers_indexes,
         nullifier_subtree_sibling_path,
         nullifier_predecessor_preimages,
-        nullifier_predecessor_membership_witnesses.map(
-            |witness: MembershipWitness<NULLIFIER_TREE_HEIGHT>| {
-                MembershipWitness {
-                    leaf_index: witness.leaf_index,
-                    sibling_path: witness.sibling_path,
-                }
-            },
-        ),
-        |low_leaf: NullifierLeafPreimage, nullifier: Field| {
-            // Is valid low leaf
-            let is_less_than_nullifier = full_field_less_than(low_leaf.nullifier, nullifier);
-            let is_next_greater_than = full_field_less_than(nullifier, low_leaf.next_nullifier);
-
-            (!low_leaf.is_empty())
-                & is_less_than_nullifier
-                & (
-                    is_next_greater_than
-                        | ((low_leaf.next_index == 0) & (low_leaf.next_nullifier == 0))
-                )
-        },
-        |low_leaf: NullifierLeafPreimage, nullifier: Field, nullifier_index: u32| {
-            // Update low leaf
-            NullifierLeafPreimage {
-                nullifier: low_leaf.nullifier,
-                next_nullifier: nullifier,
-                next_index: nullifier_index,
-            }
-        },
-        |nullifier: Field, low_leaf: NullifierLeafPreimage| {
-            // Build insertion leaf
-            NullifierLeafPreimage {
-                nullifier: nullifier,
-                next_nullifier: low_leaf.next_nullifier,
-                next_index: low_leaf.next_index,
-            }
-        },
-        [0; NULLIFIER_SUBTREE_HEIGHT],
-        [0; NULLIFIER_TREE_HEIGHT],
+        nullifier_predecessor_membership_witnesses,
     )
 }

noir-projects/noir-protocol-circuits/crates/rollup-lib/src/base/components/public_data_tree.nr

@@ -4,7 +4,6 @@ use dep::types::{
     data::{PublicDataTreeLeaf, PublicDataTreeLeafPreimage},
     merkle_tree::{indexed_tree, MembershipWitness},
     traits::is_empty,
-    utils::field::full_field_less_than,
 };
 
 pub(crate) fn public_data_tree_insert(
@@ -21,49 +20,6 @@ pub(crate) fn public_data_tree_insert(
             low_leaf_preimage,
             low_leaf_membership_witness,
             sibling_path,
-            |low_preimage: PublicDataTreeLeafPreimage, write: PublicDataTreeLeaf| {
-                // Is valid low preimage
-                let is_update = low_preimage.slot == write.slot;
-                let is_low_empty = low_preimage.is_empty();
-
-                let is_less_than_slot = full_field_less_than(low_preimage.slot, write.slot);
-                let is_next_greater_than = full_field_less_than(write.slot, low_preimage.next_slot);
-                let is_in_range = is_less_than_slot
-                    & (
-                        is_next_greater_than
-                            | ((low_preimage.next_index == 0) & (low_preimage.next_slot == 0))
-                    );
-
-                (!is_low_empty) & (is_update | is_in_range)
-            },
-            |low_preimage: PublicDataTreeLeafPreimage, write: PublicDataTreeLeaf, write_index: u32| {
-                // Update low leaf
-                let is_update = low_preimage.slot == write.slot;
-                if is_update {
-                    PublicDataTreeLeafPreimage {
-                        slot: low_preimage.slot,
-                        value: write.value,
-                        next_slot: low_preimage.next_slot,
-                        next_index: low_preimage.next_index,
-                    }
-                } else {
-                    PublicDataTreeLeafPreimage {
-                        slot: low_preimage.slot,
-                        value: low_preimage.value,
-                        next_slot: write.slot,
-                        next_index: write_index,
-                    }
-                }
-            },
-            |write: PublicDataTreeLeaf, low_preimage: PublicDataTreeLeafPreimage| {
-                // Build insertion leaf
-                PublicDataTreeLeafPreimage {
-                    slot: write.slot,
-                    value: write.value,
-                    next_slot: low_preimage.next_slot,
-                    next_index: low_preimage.next_index,
-                }
-            },
         )
     } else {
         start_snapshot

noir-projects/noir-protocol-circuits/crates/types/src/abis/nullifier_leaf_preimage.nr

@@ -39,7 +39,7 @@ impl LeafPreimage for NullifierLeafPreimage {
     }
 }
 
-impl IndexedTreeLeafPreimage for NullifierLeafPreimage {
+impl IndexedTreeLeafPreimage<Field> for NullifierLeafPreimage {
     fn get_key(self) -> Field {
         self.nullifier
     }
@@ -48,9 +48,26 @@ impl IndexedTreeLeafPreimage for NullifierLeafPreimage {
         self.next_nullifier
     }
 
+    fn points_to_infinity(self) -> bool {
+        (self.next_nullifier == 0) & (self.next_index == 0)
+    }
+
     fn as_leaf(self) -> Field {
         self.hash()
     }
+
+    fn update_pointers(self, next_key: Field, next_index: u32) -> Self {
+        Self { nullifier: self.nullifier, next_nullifier: next_key, next_index }
+    }
+
+    fn update_value(self, _nullifier: Field) -> Self {
+        assert(false, "Tried to update a nullifier");
+        Self::empty()
+    }
+
+    fn build_insertion_leaf(nullifier: Field, low_leaf: Self) -> Self {
+        Self { nullifier, next_nullifier: low_leaf.next_nullifier, next_index: low_leaf.next_index }
+    }
 }
 
 impl Readable<ScopedReadRequest> for NullifierLeafPreimage {

noir-projects/noir-protocol-circuits/crates/types/src/data/public_data_tree_leaf_preimage.nr

@@ -1,4 +1,8 @@
-use crate::{merkle_tree::leaf_preimage::IndexedTreeLeafPreimage, traits::{Empty, Hash}};
+use crate::{
+    data::public_data_tree_leaf::PublicDataTreeLeaf,
+    merkle_tree::leaf_preimage::IndexedTreeLeafPreimage,
+    traits::{Empty, Hash},
+};
 
 pub struct PublicDataTreeLeafPreimage {
     pub slot: Field,
@@ -13,6 +17,15 @@ impl Empty for PublicDataTreeLeafPreimage {
     }
 }
 
+impl Eq for PublicDataTreeLeafPreimage {
+    fn eq(self, other: Self) -> bool {
+        (self.slot == other.slot)
+            & (self.value == other.value)
+            & (self.next_slot == other.next_slot)
+            & (self.next_index == other.next_index)
+    }
+}
+
 impl Hash for PublicDataTreeLeafPreimage {
     fn hash(self) -> Field {
         if self.is_empty() {
@@ -28,7 +41,7 @@ impl Hash for PublicDataTreeLeafPreimage {
     }
 }
 
-impl IndexedTreeLeafPreimage for PublicDataTreeLeafPreimage {
+impl IndexedTreeLeafPreimage<PublicDataTreeLeaf> for PublicDataTreeLeafPreimage {
     fn get_key(self) -> Field {
         self.slot
     }
@@ -37,9 +50,35 @@ impl IndexedTreeLeafPreimage for PublicDataTreeLeafPreimage {
         self.next_slot
     }
 
+    fn points_to_infinity(self) -> bool {
+        (self.next_slot == 0) & (self.next_index == 0)
+    }
+
     fn as_leaf(self) -> Field {
         self.hash()
     }
+
+    fn update_pointers(self, next_slot: Field, next_index: u32) -> Self {
+        Self { slot: self.slot, value: self.value, next_slot, next_index }
+    }
+
+    fn update_value(self, write: PublicDataTreeLeaf) -> Self {
+        Self {
+            slot: self.slot,
+            value: write.value,
+            next_slot: self.next_slot,
+            next_index: self.next_index,
+        }
+    }
+
+    fn build_insertion_leaf(write: PublicDataTreeLeaf, low_leaf: Self) -> Self {
+        Self {
+            slot: write.slot,
+            value: write.value,
+            next_slot: low_leaf.next_slot,
+            next_index: low_leaf.next_index,
+        }
+    }
 }
 
 impl PublicDataTreeLeafPreimage {

noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree.nr

@@ -7,7 +7,7 @@ use crate::{
         membership::{assert_check_membership, MembershipWitness},
         root::{calculate_empty_tree_root, calculate_subtree_root, root_from_sibling_path},
     },
-    traits::{Empty, Hash, is_empty},
+    traits::{Empty, is_empty},
     utils::arrays::check_permutation,
 };
 
@@ -19,15 +19,10 @@ pub fn batch_insert<Value, Leaf, let SubtreeWidth: u32, let SiblingPathLength: u
     new_subtree_sibling_path: [Field; SiblingPathLength],
     low_leaf_preimages: [Leaf; SubtreeWidth],
     low_leaf_membership_witnesses: [MembershipWitness<TreeHeight>; SubtreeWidth],
-    is_valid_low_leaf: fn(Leaf, Value) -> bool,
-    update_low_leaf: fn(Leaf, Value, u32) -> Leaf,
-    build_insertion_leaf: fn(Value, Leaf) -> Leaf,
-    _subtree_height: [Field; SubtreeHeight],
-    _tree_height: [Field; TreeHeight],
 ) -> AppendOnlyTreeSnapshot
 where
-    Value: Eq + Empty,
-    Leaf: Hash + Empty,
+    Value: IndexedTreeLeafValue,
+    Leaf: IndexedTreeLeafPreimage<Value>,
 {
     // A permutation to the values is provided to make the insertion use only one insertion strategy
     // However, for the actual insertion in the tree the original order is respected, the sorting is only used for validation of the links
@@ -36,7 +31,7 @@ where
 
     // Now, update the existing leaves with the new leaves
     let mut current_tree_root = start_snapshot.root;
-    let mut insertion_subtree = [Leaf::empty(); SubtreeWidth];
+    let mut insertion_subtree = [Empty::empty(); SubtreeWidth];
     let start_insertion_index = start_snapshot.next_available_leaf_index;
 
     for i in 0..sorted_values.len() {
@@ -45,11 +40,23 @@ where
             let low_leaf_preimage = low_leaf_preimages[i];
             let witness = low_leaf_membership_witnesses[i];
 
-            assert(is_valid_low_leaf(low_leaf_preimage, value), "Invalid low leaf");
+            // validate the low leaf
+            assert(!is_empty(low_leaf_preimage), "Empty low leaf");
+            let value_key = value.get_key();
+            let low_leaf_key = low_leaf_preimage.get_key();
+            let low_leaf_next_key = low_leaf_preimage.get_next_key();
+            let is_update = value_key == low_leaf_key;
+
+            let is_less_than_slot = low_leaf_key.lt(value_key);
+            let is_next_greater_than = value_key.lt(low_leaf_next_key);
+            let is_in_range =
+                is_less_than_slot & (is_next_greater_than | low_leaf_preimage.points_to_infinity());
+
+            assert(is_update | is_in_range, "Invalid low leaf");
 
             // perform membership check for the low leaf against the original root
             assert_check_membership(
-                low_leaf_preimage.hash(),
+                low_leaf_preimage.as_leaf(),
                 witness.leaf_index,
                 witness.sibling_path,
                 current_tree_root,
@@ -58,19 +65,26 @@ where
             let value_index = sorted_values_indexes[i];
 
             // Calculate the new value of the low_leaf
-            let updated_low_leaf = update_low_leaf(
-                low_leaf_preimage,
-                value,
-                start_insertion_index as u32 + value_index,
-            );
+            let updated_low_leaf = if is_update {
+                low_leaf_preimage.update_value(value)
+            } else {
+                low_leaf_preimage.update_pointers(
+                    value_key,
+                    start_insertion_index as u32 + value_index,
+                )
+            };
 
             current_tree_root = root_from_sibling_path(
-                updated_low_leaf.hash(),
+                updated_low_leaf.as_leaf(),
                 witness.leaf_index,
                 witness.sibling_path,
             );
 
-            insertion_subtree[value_index] = build_insertion_leaf(value, low_leaf_preimage);
+            insertion_subtree[value_index] = if is_update {
+                Empty::empty()
+            } else {
+                Leaf::build_insertion_leaf(value, low_leaf_preimage)
+            };
         }
     }
 
@@ -85,7 +99,7 @@ where
     );
 
     // Create new subtree to insert into the whole indexed tree
-    let subtree_root = calculate_subtree_root(insertion_subtree.map(|leaf: Leaf| leaf.hash()));
+    let subtree_root = calculate_subtree_root(insertion_subtree.map(|leaf: Leaf| leaf.as_leaf()));
 
     // Calculate the new root
     // We are inserting a subtree rather than a full tree here
@@ -108,15 +122,24 @@ pub fn insert<Value, Leaf, let TreeHeight: u32>(
     low_leaf_preimage: Leaf,
     low_leaf_membership_witness: MembershipWitness<TreeHeight>,
     insertion_sibling_path: [Field; TreeHeight],
-    is_valid_low_leaf: fn(Leaf, Value) -> bool,
-    update_low_leaf: fn(Leaf, Value, u32) -> Leaf,
-    build_insertion_leaf: fn(Value, Leaf) -> Leaf,
 ) -> AppendOnlyTreeSnapshot
 where
     Value: IndexedTreeLeafValue,
-    Leaf: IndexedTreeLeafPreimage,
+    Leaf: IndexedTreeLeafPreimage<Value>,
 {
-    assert(is_valid_low_leaf(low_leaf_preimage, value), "Invalid low leaf");
+    // validate the low leaf
+    assert(!is_empty(low_leaf_preimage), "Empty low leaf");
+    let value_key = value.get_key();
+    let low_leaf_key = low_leaf_preimage.get_key();
+    let low_leaf_next_key = low_leaf_preimage.get_next_key();
+    let is_update = value_key == low_leaf_key;
+
+    let is_less_than_slot = low_leaf_key.lt(value_key);
+    let is_next_greater_than = value_key.lt(low_leaf_next_key);
+    let is_in_range =
+        is_less_than_slot & (is_next_greater_than | low_leaf_preimage.points_to_infinity());
+
+    assert(is_update | is_in_range, "Invalid low leaf");
 
     // perform membership check for the low leaf against the original root
     assert_check_membership(
@@ -127,8 +150,11 @@ where
     );
 
     // Calculate the new value of the low_leaf
-    let updated_low_leaf =
-        update_low_leaf(low_leaf_preimage, value, snapshot.next_available_leaf_index);
+    let updated_low_leaf = if is_update {
+        low_leaf_preimage.update_value(value)
+    } else {
+        low_leaf_preimage.update_pointers(value_key, snapshot.next_available_leaf_index)
+    };
 
     // Update low leaf
     snapshot.root = root_from_sibling_path(
@@ -137,11 +163,11 @@ where
         low_leaf_membership_witness.sibling_path,
     );
 
-    if low_leaf_preimage.get_key() == value.get_key() {
+    if is_update {
         // If it's an update, we don't need to insert the new leaf and advance the tree
         snapshot
     } else {
-        let insertion_leaf = build_insertion_leaf(value, low_leaf_preimage);
+        let insertion_leaf = Leaf::build_insertion_leaf(value, low_leaf_preimage);
         assert_check_membership(
             0,
             snapshot.next_available_leaf_index as Field,