Perf/fri merkle folding

crypto/crypto/src/merkle_tree/backends/field_element_vector.rs

@@ -56,6 +56,55 @@ where
     }
 }
 
+/// A backend for Merkle trees whose leaves are fixed-size groups of four field
+/// elements. This is the natural leaf for arity-4 FRI, where one fold orbit
+/// (four evaluations) is always opened together, so all four belong under a
+/// single leaf hash.
+#[derive(Clone)]
+pub struct FieldElementQuadBackend<F, D: Digest, const NUM_BYTES: usize> {
+    phantom1: PhantomData<F>,
+    phantom2: PhantomData<D>,
+}
+
+impl<F, D: Digest, const NUM_BYTES: usize> Default for FieldElementQuadBackend<F, D, NUM_BYTES> {
+    fn default() -> Self {
+        Self {
+            phantom1: PhantomData,
+            phantom2: PhantomData,
+        }
+    }
+}
+
+impl<F, D: Digest, const NUM_BYTES: usize> IsMerkleTreeBackend
+    for FieldElementQuadBackend<F, D, NUM_BYTES>
+where
+    F: IsField,
+    FieldElement<F>: AsBytes,
+    [u8; NUM_BYTES]: From<Output<D>>,
+{
+    type Node = [u8; NUM_BYTES];
+    type Data = [FieldElement<F>; 4];
+
+    fn hash_data(input: &[FieldElement<F>; 4]) -> [u8; NUM_BYTES] {
+        let mut hasher = D::new();
+        for element in input {
+            hasher.update(element.as_bytes());
+        }
+        let mut result_hash = [0_u8; NUM_BYTES];
+        result_hash.copy_from_slice(&hasher.finalize());
+        result_hash
+    }
+
+    fn hash_new_parent(left: &[u8; NUM_BYTES], right: &[u8; NUM_BYTES]) -> [u8; NUM_BYTES] {
+        let mut hasher = D::new();
+        hasher.update(left);
+        hasher.update(right);
+        let mut result_hash = [0_u8; NUM_BYTES];
+        result_hash.copy_from_slice(&hasher.finalize());
+        result_hash
+    }
+}
+
 #[derive(Clone)]
 pub struct FieldElementVectorBackend<F, D: Digest, const NUM_BYTES: usize> {
     phantom1: PhantomData<F>,

crypto/crypto/src/merkle_tree/backends/types.rs

@@ -2,7 +2,9 @@ use sha3::Keccak256;
 
 use super::{
     field_element::FieldElementBackend,
-    field_element_vector::{FieldElementPairBackend, FieldElementVectorBackend},
+    field_element_vector::{
+        FieldElementPairBackend, FieldElementQuadBackend, FieldElementVectorBackend,
+    },
 };
 
 // Field element backend definitions
@@ -13,3 +15,6 @@ pub type BatchKeccak256Backend<F> = FieldElementVectorBackend<F, Keccak256, 32>;
 
 // Fixed-size pair backends (more efficient for FRI layers)
 pub type PairKeccak256Backend<F> = FieldElementPairBackend<F, Keccak256, 32>;
+
+// Fixed-size quad backend: the leaf for arity-4 FRI fold orbits.
+pub type QuadKeccak256Backend<F> = FieldElementQuadBackend<F, Keccak256, 32>;

crypto/crypto/src/merkle_tree/merkle.rs

@@ -192,31 +192,53 @@ where
         &self.nodes
     }
 
-    /// Returns a Merkle proof for the element/s at position pos
-    /// For example, give me an inclusion proof for the 3rd element in the
-    /// Merkle tree
-    pub fn get_proof_by_pos(&self, pos: usize) -> Option<Proof<B::Node>> {
-        let pos = pos + self.node_count() / 2;
-        let Ok(merkle_path) = self.build_merkle_path(pos) else {
-            return None;
-        };
+    /// Total tree depth (root-to-leaf): `log2(number of leaves)`.
+    fn depth(&self) -> usize {
+        let leaves_len = self.node_count().div_ceil(2);
+        leaves_len.trailing_zeros() as usize
+    }
 
-        self.create_proof(merkle_path)
+    /// Returns the Merkle cap at `cap_height`: the `2^cap_height` nodes at tree
+    /// depth `cap_height`, ordered left to right. `cap_height` is clamped to the
+    /// tree depth; `cap_height == 0` yields a single-node cap equal to the root.
+    pub fn cap(&self, cap_height: usize) -> Vec<B::Node> {
+        let c = cap_height.min(self.depth());
+        let start = (1usize << c) - 1;
+        let end = (1usize << (c + 1)) - 1;
+        (start..end)
+            .map(|i| self.node_get(i).expect("cap node index in bounds").clone())
+            .collect()
     }
 
-    /// Creates a proof from a Merkle pasth
-    fn create_proof(&self, merkle_path: Vec<B::Node>) -> Option<Proof<B::Node>> {
+    /// Returns a Merkle proof for the element at position `pos` (full path to
+    /// the root).
+    pub fn get_proof_by_pos(&self, pos: usize) -> Option<Proof<B::Node>> {
+        self.get_proof_by_pos_capped(pos, 0)
+    }
+
+    /// Returns a Merkle proof for the element at position `pos` whose
+    /// authentication path stops at the cap level (`cap_height` levels below
+    /// the root). `cap_height == 0` is equivalent to [`get_proof_by_pos`].
+    pub fn get_proof_by_pos_capped(&self, pos: usize, cap_height: usize) -> Option<Proof<B::Node>> {
+        let pos = pos + self.node_count() / 2;
+        let merkle_path = self.build_merkle_path_capped(pos, cap_height).ok()?;
         Some(Proof { merkle_path })
     }
 
-    /// Returns the Merkle path for the element/s for the leaf at position pos
-    fn build_merkle_path(&self, pos: usize) -> Result<Vec<B::Node>, Error> {
-        // Pre-allocate based on tree depth (log2 of tree size)
-        let tree_depth = (self.node_count() + 1).ilog2() as usize;
-        let mut merkle_path = Vec::with_capacity(tree_depth);
+    /// Returns the authentication path from leaf node `pos` up to — but not
+    /// including — the cap level.
+    fn build_merkle_path_capped(
+        &self,
+        pos: usize,
+        cap_height: usize,
+    ) -> Result<Vec<B::Node>, Error> {
+        let c = cap_height.min(self.depth());
+        let cap_start = (1usize << c) - 1;
+        let cap_end = (1usize << (c + 1)) - 1;
+        let mut merkle_path = Vec::with_capacity(self.depth() - c);
         let mut pos = pos;
 
-        while pos != ROOT {
+        while !(cap_start..cap_end).contains(&pos) {
             let Some(node) = self.node_get(sibling_index(pos)) else {
                 // out of bounds, exit returning the current merkle_path
                 return Err(Error::OutOfBounds);
@@ -248,6 +270,16 @@ where
     /// - `Error::EmptyPositionList` if `pos_list` is empty
     /// - `Error::OutOfBounds` if any position in `pos_list` is >= number of leaves
     pub fn get_batch_proof(&self, pos_list: &[usize]) -> Result<BatchProof<B::Node>, Error> {
+        self.get_batch_proof_capped(pos_list, 0)
+    }
+
+    /// Batch proof variant whose authentication paths stop at the cap level.
+    /// `cap_height == 0` is equivalent to [`get_batch_proof`].
+    pub fn get_batch_proof_capped(
+        &self,
+        pos_list: &[usize],
+        cap_height: usize,
+    ) -> Result<BatchProof<B::Node>, Error> {
         if pos_list.is_empty() {
             return Err(Error::EmptyPositionList);
         }
@@ -268,7 +300,8 @@ where
             .map(|pos| pos + self.node_count() / 2)
             .collect::<Vec<usize>>();
         // We get the positions of the nodes for the batch proof.
-        let batch_auth_path_positions = self.get_batch_auth_path_positions(&leaf_positions);
+        let batch_auth_path_positions =
+            self.get_batch_auth_path_positions(&leaf_positions, cap_height);
 
         // We get the nodes for the batch proof.
         let batch_auth_path_nodes = batch_auth_path_positions
@@ -297,16 +330,20 @@ where
     ///
     /// This ordering is critical because the verifier consumes proof nodes level-by-level
     /// starting from leaves, so it needs leaf-level siblings first.
-    fn get_batch_auth_path_positions(&self, leaf_positions: &[usize]) -> Vec<usize> {
+    fn get_batch_auth_path_positions(
+        &self,
+        leaf_positions: &[usize],
+        cap_height: usize,
+    ) -> Vec<usize> {
         // BTreeSet always maintains elements in ascending order (smaller indices first), regardless of insertion order.
         let mut auth_path_set = BTreeSet::<usize>::new();
         let mut obtainable: BTreeSet<usize> = leaf_positions.iter().cloned().collect();
 
-        // Number of levels in tree
-        let num_levels = (self.node_count() + 1).ilog2();
-
-        // Iter lefevel-by-level from leaves to root.
-        for _ in 0..num_levels - 1 {
+        // Climb level-by-level from the leaves up to the cap level. A full
+        // tree has `depth` levels above the leaves; a cap of height `c` stops
+        // `c` levels early.
+        let levels_to_climb = self.depth() - cap_height.min(self.depth());
+        for _ in 0..levels_to_climb {
             let mut next_obtainable = BTreeSet::new();
 
             for &pos in &obtainable {

crypto/crypto/src/merkle_tree/proof.rs

@@ -21,7 +21,18 @@ pub struct Proof<T: PartialEq + Eq> {
 
 impl<T: PartialEq + Eq> Proof<T> {
     /// Verifies a Merkle inclusion proof for the value contained at leaf index.
-    pub fn verify<B>(&self, root_hash: &B::Node, mut index: usize, value: &B::Data) -> bool
+    pub fn verify<B>(&self, root_hash: &B::Node, index: usize, value: &B::Data) -> bool
+    where
+        B: IsMerkleTreeBackend<Node = T>,
+    {
+        self.verify_capped::<B>(core::slice::from_ref(root_hash), index, value)
+    }
+
+    /// Verifies a Merkle inclusion proof against a Merkle cap (`2^cap_height`
+    /// nodes). The path is folded up to the cap level and the result is checked
+    /// for membership in the cap. A single-node cap is equivalent to
+    /// [`verify`](Self::verify).
+    pub fn verify_capped<B>(&self, cap: &[B::Node], mut index: usize, value: &B::Data) -> bool
     where
         B: IsMerkleTreeBackend<Node = T>,
     {
@@ -37,7 +48,9 @@ impl<T: PartialEq + Eq> Proof<T> {
             index >>= 1;
         }
 
-        root_hash == &hashed_value
+        // After folding `merkle_path.len()` levels, `index` is the position of
+        // the reconstructed node within the cap.
+        cap.get(index) == Some(&hashed_value)
     }
 }
 
@@ -102,12 +115,35 @@ impl<T: PartialEq + Eq + Clone> BatchProof<T> {
         values: &[B::Data],
         num_leaves: usize,
     ) -> bool
+    where
+        B: IsMerkleTreeBackend<Node = T>,
+    {
+        self.verify_capped::<B>(
+            core::slice::from_ref(root_hash),
+            pos_list,
+            values,
+            num_leaves,
+        )
+    }
+
+    /// Batch verification against a Merkle cap (`2^cap_height` nodes). A
+    /// single-node cap is equivalent to [`verify`](Self::verify).
+    pub fn verify_capped<B>(
+        &self,
+        cap: &[B::Node],
+        pos_list: &[usize],
+        values: &[B::Data],
+        num_leaves: usize,
+    ) -> bool
     where
         B: IsMerkleTreeBackend<Node = T>,
     {
         if pos_list.len() != values.len() || pos_list.is_empty() {
             return false;
         }
+        if !cap.len().is_power_of_two() {
+            return false;
+        }
 
         // Index of the first leaf as it is ordered in the tree struct (from top to bottom).
         let first_leaf_index = num_leaves - 1;
@@ -136,9 +172,11 @@ impl<T: PartialEq + Eq + Clone> BatchProof<T> {
 
         let mut proof_path_iter = self.path.iter();
 
-        let num_levels = (2 * num_leaves).ilog2();
-        // Process level by level, from bottom to top, same as `get_batch_auth_path_positions`.
-        for _ in 0..num_levels - 1 {
+        let depth = num_leaves.trailing_zeros() as usize;
+        let cap_height = cap.len().trailing_zeros() as usize;
+        // Process level by level, from the leaves up to the cap level, same as
+        // `get_batch_auth_path_positions`.
+        for _ in 0..depth - cap_height {
             let mut next_level_known_nodes: BTreeMap<usize, T> = BTreeMap::new();
 
             // Process each known node from right to left to match the order of the proof.
@@ -178,9 +216,16 @@ impl<T: PartialEq + Eq + Clone> BatchProof<T> {
             current_level_known_nodes = next_level_known_nodes;
         }
 
-        // Verify: root computed correctly and all proof nodes consumed.
+        // Verify: every reconstructed cap-level node is in the cap, and all
+        // proof nodes were consumed.
+        let cap_start = cap.len() - 1;
         proof_path_iter.next().is_none()
-            && current_level_known_nodes.len() == 1
-            && (current_level_known_nodes.get(&0) == Some(root_hash))
+            && !current_level_known_nodes.is_empty()
+            && current_level_known_nodes.iter().all(|(tree_index, node)| {
+                tree_index
+                    .checked_sub(cap_start)
+                    .and_then(|cap_index| cap.get(cap_index))
+                    == Some(node)
+            })
     }
 }