[feat] Batch inclusion proof verification

Scarb.lock

@@ -23,3 +23,6 @@ version = "0.1.0"
 [[package]]
 name = "utreexo"
 version = "0.1.0"
+dependencies = [
+ "utils",
+]

packages/utils/src/lib.cairo

@@ -5,6 +5,7 @@ pub mod hash;
 pub mod merkle_tree;
 pub mod numeric;
 pub mod sha256;
+pub mod sort;
 
 #[cfg(target: 'test')]
 pub mod hex;

packages/utils/src/sort.cairo

@@ -0,0 +1,39 @@
+/// Bubble sort from
+/// https://github.com/keep-starknet-strange/alexandria/blob/main/packages/sorting/src/bubble_sort.cairo
+pub fn bubble_sort<T, +Copy<T>, +Drop<T>, +PartialOrd<T>>(mut array: Span<T>) -> Array<T> {
+    if array.len() == 0 {
+        return array![];
+    }
+    if array.len() == 1 {
+        return array![*array[0]];
+    }
+    let mut idx1 = 0;
+    let mut idx2 = 1;
+    let mut sorted_iteration = true;
+    let mut sorted_array = array![];
+
+    loop {
+        if idx2 == array.len() {
+            sorted_array.append(*array[idx1]);
+            if sorted_iteration {
+                break;
+            }
+            array = sorted_array.span();
+            sorted_array = array![];
+            idx1 = 0;
+            idx2 = 1;
+            sorted_iteration = true;
+        } else {
+            if *array[idx1] <= *array[idx2] {
+                sorted_array.append(*array[idx1]);
+                idx1 = idx2;
+                idx2 += 1;
+            } else {
+                sorted_array.append(*array[idx2]);
+                idx2 += 1;
+                sorted_iteration = false;
+            }
+        };
+    };
+    sorted_array
+}

packages/utreexo/Scarb.toml

@@ -3,6 +3,9 @@ name = "utreexo"
 version = "0.1.0"
 edition = "2024_07"
 
+[dependencies]
+utils = { path = "../utils" }
+
 [dev-dependencies]
 cairo_test.workspace = true
 

packages/utreexo/src/stump/accumulator.cairo

@@ -1,5 +1,5 @@
 use crate::stump::state::UtreexoStumpState;
-use crate::stump::proof::UtreexoBatchProof;
+use crate::stump::proof::{UtreexoBatchProof, UtreexoBatchProofTrait};
 
 #[generate_trait]
 pub impl StumpUtreexoAccumulatorImpl of StumpUtreexoAccumulator {
@@ -13,7 +13,36 @@ pub impl StumpUtreexoAccumulatorImpl of StumpUtreexoAccumulator {
     fn verify(
         self: @UtreexoStumpState, proof: @UtreexoBatchProof, del_hashes: Span<felt252>
     ) -> Result<(), ByteArray> {
-        // TODO
+        if (*proof.targets).is_empty() {
+            return Result::Ok(());
+        };
+
+        let computed_roots: Span<felt252> = proof.compute_roots(del_hashes, *self.num_leaves)?;
+
+        let mut number_matched_roots: u32 = 0;
+
+        for i in 0
+            ..computed_roots
+                .len() {
+                    for root in *self
+                        .roots {
+                            match root {
+                                Option::Some(root) => {
+                                    if (computed_roots[i] == root) {
+                                        number_matched_roots += 1;
+                                    };
+                                },
+                                Option::None => {},
+                            };
+                        };
+                };
+
+        let computed_roots_len = computed_roots.len();
+
+        if (computed_roots_len != number_matched_roots && computed_roots_len != 0) {
+            return Result::Err("Proof verification failed");
+        }
+
         Result::Ok(())
     }
 
@@ -22,3 +51,18 @@ pub impl StumpUtreexoAccumulatorImpl of StumpUtreexoAccumulator {
         *self
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{UtreexoStumpState, StumpUtreexoAccumulator, UtreexoBatchProof};
+
+    #[test]
+    fn test_verification_1() {
+        let state = UtreexoStumpState { roots: array![Option::Some(0x371cb6995ea5e7effcd2e174de264b5b407027a75a231a70c2c8d196107f0e7)].span(), num_leaves: 2 };
+        let batch_proof = UtreexoBatchProof { targets: array![0].span(), proof: array![2].span() };
+        let del_hashes = array![1];
+
+        let result = state.verify(@batch_proof, del_hashes.span(),);
+        assert_eq!(result, Result::Ok(()));
+    }
+}

packages/utreexo/src/stump/proof.cairo

@@ -1,10 +1,13 @@
 use core::fmt::{Display, Formatter, Error};
+use core::num::traits::Bounded;
+use crate::parent_hash;
+use utils::{bit_shifts::shr, sort::bubble_sort};
 
 /// Utreexo inclusion proof for multiple outputs.
 /// Compatible with https://github.com/utreexo/utreexo
 #[derive(Drop, Copy)]
 pub struct UtreexoBatchProof {
-    /// List of sibling nodes required to calculate the root.
+    /// List of sibling nodes required to calculate the roots.
     pub proof: Span<felt252>,
     /// Indices of leaves to be deleted (ordered starting from 0, left to right).
     pub targets: Span<u64>,
@@ -27,3 +30,170 @@ impl UtreexoBatchProofDisplay of Display<UtreexoBatchProof> {
         Result::Ok(())
     }
 }
+
+#[generate_trait]
+pub impl UtreexoBatchProofImpl of UtreexoBatchProofTrait {
+    /// Computes a set of roots given a proof and leaves hashes.
+    fn compute_roots(
+        self: @UtreexoBatchProof, mut del_hashes: Span<felt252>, num_leaves: u64,
+    ) -> Result<Span<felt252>, ByteArray> {
+        // Where all the parent hashes we've calculated in a given row will go to.
+        let mut calculated_root_hashes: Array<felt252> = array![];
+        // Target leaves
+        let mut leaf_nodes: Array<(u64, felt252)> = array![];
+
+        // Append targets with their hashes.
+        let mut positions = *self.targets;
+        while let Option::Some(rhs) = del_hashes.pop_front() {
+            if let Option::Some(lhs) = positions.pop_front() {
+                leaf_nodes.append((*lhs, *rhs));
+            }
+        };
+
+        let mut leaf_nodes: Array<(u64, felt252)> = bubble_sort(leaf_nodes.span());
+        let mut inner_result = Result::Ok((array![].span()));
+
+        // Proof nodes.
+        let mut sibling_nodes: Array<felt252> = (*self.proof).into();
+        // Queue of computed intermediate nodes.
+        let mut computed_nodes: Array<(u64, felt252)> = array![];
+        // Actual length of the current row
+        let mut actual_row_len: u64 = num_leaves;
+        // Length of the "padded" row which is always power of two.
+        let mut row_len: u64 = next_power_of_two(num_leaves);
+        // Total padded length of processed rows (excluding the current one)
+        let mut row_len_acc: u64 = 0;
+        // Next position of the target leaf and the leaf itself.
+        let (mut next_leaf_pos, mut next_leaf) = leaf_nodes.pop_front().unwrap();
+        // Next computed node
+        let mut next_computed: felt252 = 0;
+        // Position of the next computed node
+        let mut next_computed_pos: u64 = Bounded::<u64>::MAX;
+
+        while row_len != 0 {
+            let (pos, node) = if next_leaf_pos < next_computed_pos {
+                let res = (next_leaf_pos, next_leaf);
+                let (a, b) = leaf_nodes.pop_front().unwrap_or((Bounded::<u64>::MAX, 0));
+                next_leaf_pos = a;
+                next_leaf = b;
+                res
+            } else if next_computed_pos != Bounded::<u64>::MAX {
+                let res = (next_computed_pos, next_computed);
+                let (a, b) = computed_nodes.pop_front().unwrap_or((Bounded::<u64>::MAX, 0));
+                next_computed_pos = a;
+                next_computed = b;
+                res
+            } else {
+                // Out of nodes, terminating here.
+                break;
+            };
+
+            // If we are beyond current row, level up.
+            while pos >= row_len_acc + row_len {
+                row_len_acc += row_len;
+                row_len /= 2;
+                actual_row_len /= 2;
+
+                if row_len == 0 {
+                    inner_result =
+                        Result::Err(
+                            format!("Position {pos} is out of the forest range {row_len_acc}")
+                        );
+                    break;
+                }
+            };
+
+            // If row length is odd and we are at the edge this is a root.
+            if pos == row_len_acc + actual_row_len - 1 && actual_row_len % 2 == 1 {
+                calculated_root_hashes.append(node);
+                row_len_acc += row_len;
+                row_len /= 2;
+                actual_row_len /= 2;
+                continue;
+            };
+
+            let parent_node = if (pos - row_len_acc) % 2 == 0 {
+                // Right sibling can be both leaf/computed or proof.
+                let right_sibling = if next_leaf_pos == pos + 1 {
+                    let res = next_leaf;
+                    let (a, b) = leaf_nodes.pop_front().unwrap_or((Bounded::<u64>::MAX, 0));
+                    next_leaf_pos = a;
+                    next_leaf = b;
+                    res
+                } else if next_computed_pos == pos + 1 {
+                    let res = next_computed;
+                    let (a, b) = computed_nodes.pop_front().unwrap_or((Bounded::<u64>::MAX, 0));
+                    next_computed_pos = a;
+                    next_computed = b;
+                    res
+                } else {
+                    if sibling_nodes.is_empty() {
+                        inner_result = Result::Err("Proof is empty");
+                        break;
+                    };
+                    sibling_nodes.pop_front().unwrap()
+                };
+                parent_hash(node, right_sibling)
+            } else {
+                // Left sibling always from proof.
+                if let Option::Some(left_sibling) = sibling_nodes.pop_front() {
+                    parent_hash(left_sibling, node)
+                } else {
+                    inner_result = Result::Err("Proof is empty");
+                    break;
+                }
+            };
+
+            let parent_pos = row_len_acc + row_len + (pos - row_len_acc) / 2;
+
+            if next_computed_pos == Bounded::<u64>::MAX {
+                next_computed_pos = parent_pos;
+                next_computed = parent_node;
+            } else {
+                computed_nodes.append((parent_pos, parent_node));
+            }
+        };
+
+        if !sibling_nodes.is_empty() {
+            return Result::Err("Proof should be empty");
+        }
+
+        if inner_result != Result::Ok((array![].span())) {
+            inner_result
+        } else {
+            Result::Ok((calculated_root_hashes.span()))
+        }
+    }
+}
+
+/// Computes the next power of two of a u64 variable.
+fn next_power_of_two(mut n: u64) -> u64 {
+    if n == 0 {
+        return 1;
+    }
+
+    n -= 1;
+    n = n | shr(n, 1_u64);
+    n = n | shr(n, 2_u64);
+    n = n | shr(n, 4_u64);
+    n = n | shr(n, 8_u64);
+    n = n | shr(n, 16_u64);
+    n = n | shr(n, 32_u64);
+
+    n + 1
+}
+
+/// PartialOrd implementation for tuple (u32, felt252).
+impl PartialOrdTupleU64Felt252 of PartialOrd<(u64, felt252)> {
+    fn lt(lhs: (u64, felt252), rhs: (u64, felt252)) -> bool {
+        let (a, _) = lhs;
+        let (b, _) = rhs;
+
+        if a < b {
+            true
+        } else {
+            false
+        }
+    }
+}
+

packages/utreexo/src/vanilla/proof.cairo

@@ -27,7 +27,7 @@ impl UtreexoProofDisplay of Display<UtreexoProof> {
 
 #[generate_trait]
 pub impl UtreexoProofImpl of UtreexoProofTrait {
-    /// Computes the root given a a proof and leaf hash.
+    /// Computes the root given a proof and leaf hash.
     ///
     /// Traverses the tree from leaf to root, hashing paired nodes.
     /// Proof order is bottom-up. Returns the computed root.