Skip to content

Commit

Permalink
added support for multiproof generation within Prover
Browse files Browse the repository at this point in the history
  • Loading branch information
noahfigueras committed Jul 23, 2024
1 parent 84ef2b7 commit d52246b
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 4 deletions.
2 changes: 1 addition & 1 deletion ssz-rs/src/merkleization/multiproofs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ fn get_path_indices(tree_index: GeneralizedIndex) -> Vec<GeneralizedIndex> {
result
}

fn get_helper_indices(indices: &[GeneralizedIndex]) -> Vec<GeneralizedIndex> {
pub fn get_helper_indices(indices: &[GeneralizedIndex]) -> Vec<GeneralizedIndex> {
let mut all_helper_indices = HashSet::new();
let mut all_path_indices = HashSet::new();

Expand Down
69 changes: 66 additions & 3 deletions ssz-rs/src/merkleization/proofs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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<GeneralizedIndex> =
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<Node>,
pub branch: Vec<Node>,
pub indices: Vec<GeneralizedIndex>,
}

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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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<u8, 16>,
c: u8,
}
let data = TestContainer {
a: 18745094,
b: List::<u8, 16>::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());
}
}

0 comments on commit d52246b

Please sign in to comment.