wip finish impl Prove for Vector

ssz-rs/src/boolean.rs

@@ -2,8 +2,8 @@ use crate::{
     de::{Deserialize, DeserializeError},
     lib::*,
     merkleization::{
-        proofs::{prove_primitive, ProofAndWitness, Prove},
-        GeneralizedIndex, GeneralizedIndexable, HashTreeRoot, MerkleizationError, Node,
+        proofs::{prove_primitive, ProofAndWitness, Prove, Prover},
+        GeneralizedIndexable, HashTreeRoot, MerkleizationError, Node,
     },
     ser::{Serialize, SerializeError},
     Serializable, SimpleSerialize,
@@ -63,8 +63,8 @@ impl GeneralizedIndexable for bool {
 }
 
 impl Prove for bool {
-    fn prove(&mut self, index: GeneralizedIndex) -> Result<ProofAndWitness, MerkleizationError> {
-        prove_primitive(self, index)
+    fn prove(&mut self, prover: &mut Prover) -> Result<(), MerkleizationError> {
+        prove_primitive(self, prover)
     }
 }
 

ssz-rs/src/merkleization/merkleize.rs

@@ -3,7 +3,9 @@ use crate::{
     lib::*,
     merkleization::{MerkleizationError as Error, Node, BYTES_PER_CHUNK},
     ser::Serialize,
+    GeneralizedIndex,
 };
+use alloy_primitives::hex;
 use sha2::{Digest, Sha256};
 
 /// Types that can provide the root of their corresponding Merkle tree following the SSZ spec.
@@ -43,7 +45,7 @@ where
     Ok(buffer)
 }
 
-fn hash_nodes(hasher: &mut Sha256, a: &[u8], b: &[u8], out: &mut [u8]) {
+pub fn hash_nodes(hasher: &mut Sha256, a: &[u8], b: &[u8], out: &mut [u8]) {
     hasher.update(a);
     hasher.update(b);
     out.copy_from_slice(&hasher.finalize_reset());
@@ -226,18 +228,91 @@ pub(crate) fn elements_to_chunks<'a, T: HashTreeRoot + 'a>(
     Ok(chunks)
 }
 
+pub struct Tree(Vec<u8>);
+
+impl Index<GeneralizedIndex> for Tree {
+    type Output = [u8];
+
+    fn index(&self, index: GeneralizedIndex) -> &Self::Output {
+        let start = (index - 1) * BYTES_PER_CHUNK;
+        let end = index * BYTES_PER_CHUNK;
+        &self.0[start..end]
+    }
+}
+
+impl std::fmt::Debug for Tree {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        for chunk in self.0.chunks(BYTES_PER_CHUNK) {
+            let node = hex::encode(chunk);
+            f.write_str(&node)?;
+            f.write_str(",\n")?;
+        }
+        Ok(())
+    }
+}
+
+// Return the full Merkle tree of the `chunks`.
+// Invariant: `chunks.len() % BYTES_PER_CHUNK == 0`
+// Invariant: `leaf_count.next_power_of_two() == leaf_count`
+// NOTE: naive implementation, can make much more efficient
+pub fn compute_merkle_tree(
+    hasher: &mut Sha256,
+    chunks: &[u8],
+    leaf_count: usize,
+) -> Result<Tree, Error> {
+    debug_assert!(chunks.len() % BYTES_PER_CHUNK == 0);
+    debug_assert!(leaf_count.next_power_of_two() == leaf_count);
+
+    // SAFETY: checked subtraction is unnecessary,
+    // as leaf_count != 0 (0.next_power_of_two() == 1); qed
+    let node_count = 2 * leaf_count - 1;
+    // SAFETY: checked subtraction is unnecessary, as node_count >= leaf_count; qed
+    let interior_count = node_count - leaf_count;
+    let leaf_start = interior_count * BYTES_PER_CHUNK;
+
+    let mut buffer = vec![0u8; node_count * BYTES_PER_CHUNK];
+    buffer[leaf_start..leaf_start + chunks.len()].copy_from_slice(chunks);
+
+    for i in (1..node_count).rev().step_by(2) {
+        // SAFETY: checked subtraction is unnecessary, as i >= 1; qed
+        let parent_index = (i - 1) / 2;
+        let focus = &mut buffer[parent_index * BYTES_PER_CHUNK..(i + 1) * BYTES_PER_CHUNK];
+        // SAFETY: checked subtraction is unnecessary:
+        // focus.len() = (i + 1 - parent_index) * BYTES_PER_CHUNK
+        //             = ((2*i + 2 - i + 1) / 2) * BYTES_PER_CHUNK
+        //             = ((i + 3) / 2) * BYTES_PER_CHUNK
+        // and
+        // i >= 1
+        // so focus.len() >= 2 * BYTES_PER_CHUNK; qed
+        let children_index = focus.len() - 2 * BYTES_PER_CHUNK;
+        // NOTE: children.len() == 2 * BYTES_PER_CHUNK
+        let (parent, children) = focus.split_at_mut(children_index);
+        let (left, right) = children.split_at(BYTES_PER_CHUNK);
+        hash_nodes(hasher, left, right, &mut parent[..BYTES_PER_CHUNK]);
+    }
+    Ok(Tree(buffer))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate as ssz_rs;
-    use crate::prelude::*;
+    use crate::{merkleization::default_generalized_index, prelude::*};
 
     macro_rules! hex {
         ($input:expr) => {
             hex::decode($input).unwrap()
         };
     }
 
+    // Return the root of the Merklization of a binary tree formed from `chunks`.
+    fn merkleize_chunks(chunks: &[u8], leaf_count: usize) -> Result<Node, Error> {
+        let mut hasher = Sha256::new();
+        let tree = compute_merkle_tree(&mut hasher, chunks, leaf_count)?;
+        let root_index = default_generalized_index();
+        Ok(tree[root_index].try_into().expect("can produce a single root chunk"))
+    }
+
     #[test]
     fn test_packing_basic_types_simple() {
         let b = true;
@@ -286,45 +361,6 @@ mod tests {
         assert_eq!(result, expected);
     }
 
-    // Return the root of the Merklization of a binary tree formed from `chunks`.
-    // Invariant: `chunks.len() % BYTES_PER_CHUNK == 0`
-    // Invariant: `leaf_count.next_power_of_two() == leaf_count`
-    // NOTE: naive implementation, can make much more efficient
-    fn merkleize_chunks(chunks: &[u8], leaf_count: usize) -> Result<Node, MerkleizationError> {
-        debug_assert!(chunks.len() % BYTES_PER_CHUNK == 0);
-        debug_assert!(leaf_count.next_power_of_two() == leaf_count);
-
-        // SAFETY: checked subtraction is unnecessary,
-        // as leaf_count != 0 (0.next_power_of_two() == 1); qed
-        let node_count = 2 * leaf_count - 1;
-        // SAFETY: checked subtraction is unnecessary, as node_count >= leaf_count; qed
-        let interior_count = node_count - leaf_count;
-        let leaf_start = interior_count * BYTES_PER_CHUNK;
-
-        let mut hasher = Sha256::new();
-        let mut buffer = vec![0u8; node_count * BYTES_PER_CHUNK];
-        buffer[leaf_start..leaf_start + chunks.len()].copy_from_slice(chunks);
-
-        for i in (1..node_count).rev().step_by(2) {
-            // SAFETY: checked subtraction is unnecessary, as i >= 1; qed
-            let parent_index = (i - 1) / 2;
-            let focus = &mut buffer[parent_index * BYTES_PER_CHUNK..(i + 1) * BYTES_PER_CHUNK];
-            // SAFETY: checked subtraction is unnecessary:
-            // focus.len() = (i + 1 - parent_index) * BYTES_PER_CHUNK
-            //             = ((2*i + 2 - i + 1) / 2) * BYTES_PER_CHUNK
-            //             = ((i + 3) / 2) * BYTES_PER_CHUNK
-            // and
-            // i >= 1
-            // so focus.len() >= 2 * BYTES_PER_CHUNK; qed
-            let children_index = focus.len() - 2 * BYTES_PER_CHUNK;
-            // NOTE: children.len() == 2 * BYTES_PER_CHUNK
-            let (parent, children) = focus.split_at_mut(children_index);
-            let (left, right) = children.split_at(BYTES_PER_CHUNK);
-            hash_nodes(&mut hasher, left, right, &mut parent[..BYTES_PER_CHUNK]);
-        }
-        Ok(buffer[..BYTES_PER_CHUNK].try_into().expect("can produce a single root chunk"))
-    }
-
     #[test]
     fn test_naive_merkleize_chunks() {
         let chunks = vec![0u8; 2 * BYTES_PER_CHUNK];

ssz-rs/src/merkleization/mod.rs

@@ -30,6 +30,8 @@ pub enum MerkleizationError {
     InvalidPathElement(PathElement),
     /// Signals an invalid path when walking a `GeneralizedIndexable` type
     InvalidPath(Vec<PathElement>),
+    InvalidDepth,
+    InvalidIndex,
 }
 
 impl From<SerializeError> for MerkleizationError {
@@ -49,6 +51,8 @@ impl Display for MerkleizationError {
             Self::InvalidGeneralizedIndex => write!(f, "invalid generalized index"),
             Self::InvalidPathElement(element) => write!(f, "invalid path element {element:?}"),
             Self::InvalidPath(path) => write!(f, "invalid path {path:?}"),
+            Self::InvalidDepth => write!(f, "error computing depth for proof"),
+            Self::InvalidIndex => write!(f, "error computing index for proof"),
         }
     }
 }

ssz-rs/src/merkleization/proofs.rs

@@ -10,15 +10,67 @@ use sha2::{Digest, Sha256};
 
 pub type ProofAndWitness = (Proof, Node);
 
+fn get_depth(i: GeneralizedIndex) -> Result<u32, Error> {
+    log_2(i).ok_or(Error::InvalidGeneralizedIndex)
+}
+
+fn get_index(i: GeneralizedIndex, depth: u32) -> usize {
+    i % 2usize.pow(depth)
+}
+
 pub fn get_subtree_index(i: GeneralizedIndex) -> Result<usize, Error> {
-    let i_log2 = log_2(i).ok_or(Error::InvalidGeneralizedIndex)?;
-    Ok(i % 2usize.pow(i_log2))
+    let depth = get_depth(i)?;
+    Ok(get_index(i, depth))
+}
+
+#[derive(Debug)]
+pub struct Prover {
+    pub hasher: Sha256,
+    pub proof: Proof,
+    pub witness: Node,
+}
+
+impl Prover {
+    pub fn set_leaf(&mut self, leaf: Node) {
+        self.proof.leaf = leaf;
+    }
+
+    // Adds a node to the Merkle proof's branch.
+    // Assumes nodes are provided going from the bottom of the tree to the top.
+    pub fn extend_branch(&mut self, node: Node) {
+        self.proof.branch.push(node)
+    }
+
+    pub fn set_witness(&mut self, witness: Node) {
+        self.witness = witness;
+    }
+
+    pub fn compute_depth_and_index(&self, i: GeneralizedIndex) -> Result<(u32, usize), Error> {
+        let depth = get_depth(i)?;
+        Ok((depth, get_index(i, depth)))
+    }
+}
+
+impl From<Prover> for ProofAndWitness {
+    fn from(value: Prover) -> Self {
+        (value.proof, value.witness)
+    }
+}
+
+impl From<GeneralizedIndex> for Prover {
+    fn from(index: GeneralizedIndex) -> Self {
+        Self {
+            hasher: Sha256::new(),
+            proof: Proof { leaf: Default::default(), branch: vec![], index },
+            witness: Default::default(),
+        }
+    }
 }
 
 /// Types that can produce Merkle proofs against themselves given a `GeneralizedIndex`.
 pub trait Prove {
     /// Provide a Merkle proof of the node in this type's merkle tree corresponding to the `index`.
-    fn prove(&mut self, index: GeneralizedIndex) -> Result<ProofAndWitness, Error>;
+    fn prove(&mut self, prover: &mut Prover) -> Result<(), Error>;
 }
 
 /// Produce a Merkle proof (and corresponding witness) for the type `T` at the given `path` relative
@@ -28,7 +80,9 @@ pub fn prove<T: GeneralizedIndexable + Prove>(
     path: Path,
 ) -> Result<ProofAndWitness, Error> {
     let index = T::generalized_index(path)?;
-    data.prove(index)
+    let mut prover = index.into();
+    data.prove(&mut prover)?;
+    Ok(prover.into())
 }
 
 /// Contains data necessary to verify `leaf` was included under some witness "root" node
@@ -51,15 +105,17 @@ impl Proof {
 
 pub fn prove_primitive<T: HashTreeRoot + ?Sized>(
     data: &mut T,
-    index: GeneralizedIndex,
-) -> Result<ProofAndWitness, Error> {
+    prover: &mut Prover,
+) -> Result<(), Error> {
+    let index = prover.proof.index;
     if index != default_generalized_index() {
         return Err(Error::InvalidGeneralizedIndex)
     }
 
     let root = data.hash_tree_root()?;
-    let proof = Proof { leaf: root, branch: vec![], index };
-    Ok((proof, root))
+    prover.set_leaf(root);
+    prover.set_witness(root);
+    Ok(())
 }
 
 pub fn is_valid_merkle_branch_for_generalized_index<T: AsRef<[u8]>>(

ssz-rs/src/uint.rs

@@ -3,8 +3,8 @@ use crate::{
     lib::*,
     merkleization::{
         pack_bytes,
-        proofs::{prove_primitive, ProofAndWitness, Prove},
-        GeneralizedIndex, GeneralizedIndexable, HashTreeRoot, MerkleizationError, Node,
+        proofs::{prove_primitive, ProofAndWitness, Prove, Prover},
+        GeneralizedIndexable, HashTreeRoot, MerkleizationError, Node,
     },
     ser::{Serialize, SerializeError},
     Serializable, SimpleSerialize, BITS_PER_BYTE,
@@ -76,11 +76,8 @@ macro_rules! define_uint {
         }
 
         impl Prove for $uint {
-            fn prove(
-                &mut self,
-                index: GeneralizedIndex,
-            ) -> Result<ProofAndWitness, MerkleizationError> {
-                prove_primitive(self, index)
+            fn prove(&mut self, prover: &mut Prover) -> Result<(), MerkleizationError> {
+                prove_primitive(self, prover)
             }
         }
 
@@ -156,8 +153,8 @@ impl GeneralizedIndexable for U256 {
 }
 
 impl Prove for U256 {
-    fn prove(&mut self, index: GeneralizedIndex) -> Result<ProofAndWitness, MerkleizationError> {
-        prove_primitive(self, index)
+    fn prove(&mut self, prover: &mut Prover) -> Result<(), MerkleizationError> {
+        prove_primitive(self, prover)
     }
 }