diff --git a/ssz-rs/src/merkleization/multiproofs.rs b/ssz-rs/src/merkleization/multiproofs.rs index 1882c13e..bf1f3ebd 100644 --- a/ssz-rs/src/merkleization/multiproofs.rs +++ b/ssz-rs/src/merkleization/multiproofs.rs @@ -30,7 +30,7 @@ fn get_path_indices(tree_index: GeneralizedIndex) -> Vec { result } -fn get_helper_indices(indices: &[GeneralizedIndex]) -> Vec { +pub fn get_helper_indices(indices: &[GeneralizedIndex]) -> Vec { let mut all_helper_indices = HashSet::new(); let mut all_path_indices = HashSet::new(); diff --git a/ssz-rs/src/merkleization/proofs.rs b/ssz-rs/src/merkleization/proofs.rs index abbc7244..a5629cd3 100644 --- a/ssz-rs/src/merkleization/proofs.rs +++ b/ssz-rs/src/merkleization/proofs.rs @@ -3,8 +3,8 @@ pub use crate::merkleization::generalized_index::log_2; use crate::{ lib::*, merkleization::{ - compute_merkle_tree, GeneralizedIndex, GeneralizedIndexable, MerkleizationError as Error, - Node, Path, + compute_merkle_tree, multiproofs, GeneralizedIndex, GeneralizedIndexable, + MerkleizationError as Error, Node, Path, }, }; use sha2::{Digest, Sha256}; @@ -165,6 +165,50 @@ pub trait Prove: GeneralizedIndexable { prover.compute_proof(self)?; Ok(prover.into()) } + + /// Compute a Multi Merkle proof of `Self` at the type's `paths`, along with the root of the + /// Merkle tree as a witness value. + fn multi_prove(&self, paths: &[Path]) -> Result<(MultiProof, Node), Error> { + let indices: Vec = + paths.iter().map(|x| Self::generalized_index(x).unwrap()).collect(); + let helpers = multiproofs::get_helper_indices(&indices); + let mut proof = MultiProof { leaves: vec![], branch: vec![], indices: vec![] }; + let mut witness: Node = Node::ZERO; + + for index in indices.iter() { + let mut prover = Prover::from(*index); + prover.compute_proof(self)?; + proof.leaves.push(prover.proof.leaf); + proof.indices.push(prover.proof.index); + witness = prover.witness; + } + + for helper in helpers.iter() { + let mut prover = Prover::from(*helper); + prover.compute_proof(self)?; + proof.branch.push(prover.proof.leaf); + } + + Ok((proof, witness)) + } +} + +/// Contains data necessary to verify `leaf` was included under some witness "root" node +/// at the generalized position `index`. +#[derive(Debug, PartialEq, Eq)] +pub struct MultiProof { + pub leaves: Vec, + pub branch: Vec, + pub indices: Vec, +} + +impl MultiProof { + /// Verify `self` against the provided `root` witness node. + /// This `root` is the hash tree root of the SSZ object that produced the proof. + /// See `Prover` for further information. + pub fn verify(&self, root: Node) -> Result<(), Error> { + multiproofs::verify_merkle_multiproof(&self.leaves, &self.branch, &self.indices, root) + } } /// Contains data necessary to verify `leaf` was included under some witness "root" node @@ -207,7 +251,7 @@ pub fn is_valid_merkle_branch( root: Node, ) -> Result<(), Error> { if branch.len() != depth { - return Err(Error::InvalidProof) + return Err(Error::InvalidProof); } let mut derived_root = leaf; @@ -336,4 +380,23 @@ pub(crate) mod tests { let data = false; compute_and_verify_proof_for_path(&data, &[]); } + + #[test] + fn test_generate_multi_proofs() { + #[derive(PartialEq, Eq, Debug, Default, SimpleSerialize)] + struct TestContainer { + a: u64, + b: List, + c: u8, + } + let data = TestContainer { + a: 18745094, + b: List::::try_from(vec![200, 50, 40, 80]).unwrap(), + c: 240, + }; + + let (proof, witness) = + data.multi_prove(&[&["a".into()], &["b".into(), 0.into()], &["c".into()]]).unwrap(); + assert!(proof.verify(witness).is_ok()); + } }