diff --git a/src/merkletree.rs b/src/merkletree.rs index c2e3b00..60eeb49 100644 --- a/src/merkletree.rs +++ b/src/merkletree.rs @@ -160,6 +160,22 @@ impl MerkleTree { .map(|lemma| Proof::new(self.algorithm, root_hash, lemma, value)) } + /// Generate an inclusion proof for the `n`-th leaf value. + pub fn gen_nth_proof(&self, n: usize) -> Option> + where + T: Hashable + Clone + { + let root_hash = self.root_hash().clone(); + let value = if let Some(value) = self.root.nth_leaf(n, self.count) { + value + } else { + return None; + }; + + Lemma::new_by_index(&self.root, n, self.count) + .map(|lemma| Proof::new(self.algorithm, root_hash, lemma, value.clone())) + } + /// Creates an `Iterator` over the values contained in this Merkle tree. pub fn iter(&self) -> LeavesIterator { self.root.iter() diff --git a/src/proof.rs b/src/proof.rs index 2fb7cee..48a8431 100644 --- a/src/proof.rs +++ b/src/proof.rs @@ -153,6 +153,11 @@ impl Proof { self.lemma.validate(self.algorithm) } + + /// Returns the index of this proof's value, given the total number of items in the tree. + pub fn index(&self, count: usize) -> usize { + self.lemma.index(count) + } } /// A `Lemma` holds the hash of a node, the hash of its sibling node, @@ -170,7 +175,8 @@ pub struct Lemma { } impl Lemma { - /// Attempts to generate a proof that the a value with hash `needle` is a member of the given `tree`. + /// Attempts to generate a proof that the a value with hash `needle` is a + /// member of the given `tree`. pub fn new(tree: &Tree, needle: &[u8]) -> Option { match *tree { Tree::Empty { .. } => None, @@ -185,6 +191,47 @@ impl Lemma { } } + /// Attempts to generate a proof that the `idx`-th leaf is a member of + /// the given tree. The `count` must equal the number of leaves in the + /// `tree`. If `idx >= count`, `None` is returned. + pub fn new_by_index(tree: &Tree, idx: usize, count: usize) -> Option { + if idx >= count { + return None; + } + match *tree { + Tree::Empty { .. } => None, + + Tree::Leaf { ref hash, .. } => { + if count != 1 { + return None; + } + Some(Lemma { + node_hash: hash.clone(), + sibling_hash: None, + sub_lemma: None, + }) + } + + Tree::Node { + ref hash, + ref left, + ref right, + } => Lemma::new_tree_proof_by_index(hash, idx, count, left, right), + } + } + + /// Returns the index of this lemma's value, given the total number of items in the tree. + pub fn index(&self, count: usize) -> usize { + let left_count = count.next_power_of_two() / 2; + match (self.sub_lemma.as_ref(), self.sibling_hash.as_ref()) { + (None, Some(&Positioned::Right(_))) | (None, None) => 0, + (None, Some(&Positioned::Left(_))) => 1, + (Some(l), None) => l.index(count), + (Some(l), Some(&Positioned::Left(_))) => left_count + l.index(count - left_count), + (Some(l), Some(&Positioned::Right(_))) => l.index(left_count), + } + } + fn new_leaf_proof(hash: &[u8], needle: &[u8]) -> Option { if *hash == *needle { Some(Lemma { @@ -245,6 +292,35 @@ impl Lemma { }, } } + + fn new_tree_proof_by_index( + hash: &[u8], + idx: usize, + count: usize, + left: &Tree, + right: &Tree, + ) -> Option { + let left_count = count.next_power_of_two() / 2; + Lemma::new_by_index(left, idx, left_count) + .map(|lemma| { + let right_hash = right.hash().clone(); + let sub_lemma = Some(Positioned::Right(right_hash)); + (lemma, sub_lemma) + }) + .or_else(|| { + let sub_lemma = Lemma::new_by_index(right, idx - left_count, count - left_count); + sub_lemma.map(|lemma| { + let left_hash = left.hash().clone(); + let sub_lemma = Some(Positioned::Left(left_hash)); + (lemma, sub_lemma) + }) + }) + .map(|(sub_lemma, sibling_hash)| Lemma { + node_hash: hash.into(), + sibling_hash, + sub_lemma: Some(Box::new(sub_lemma)), + }) + } } /// Tags a value so that we know from which branch of a `Tree` (if any) it was found. diff --git a/src/tests.rs b/src/tests.rs index 390259f..81e9bfa 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -150,6 +150,23 @@ fn test_wrong_proof() { } } +#[test] +fn test_nth_proof() { + // Calculation depends on the total count. Try a few numbers: odd, even, powers of two... + for &count in [1, 2, 3, 10, 15, 16, 17, 22].iter() { + let values = (1..(count + 1)).map(|x| vec![x as u8]).collect::>(); + let tree = MerkleTree::from_vec(digest, values.clone()); + let root_hash = tree.root_hash(); + + for i in 0..count { + let proof = tree.gen_nth_proof(i).expect("gen proof by index"); + assert_eq!(vec![i as u8 + 1], proof.value); + assert!(proof.validate(&root_hash)); + assert_eq!(i, proof.index(tree.count())); + } + } +} + #[test] fn test_mutate_proof_first_lemma() { let values = (1..10).map(|x| vec![x]).collect::>(); diff --git a/src/tree.rs b/src/tree.rs index e65fb98..6294b59 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -61,6 +61,31 @@ impl Tree { pub fn iter(&self) -> LeavesIterator { LeavesIterator::new(self) } + + /// Returns the `n`-th leaf value, given the total `count`. + pub fn nth_leaf(&self, n: usize, count: usize) -> Option<&T> { + if n >= count { + return None; + } + let left_count = count.next_power_of_two() >> 1; + match *self { + Tree::Empty { .. } => None, + Tree::Leaf { ref value, .. } => { + if n == 0 { + Some(value) + } else { + None + } + } + Tree::Node { ref left, ref right, .. } => { + if n < left_count { + left.nth_leaf(n, left_count) + } else { + right.nth_leaf(n - left_count, count - left_count) + } + } + } + } } /// An borrowing iterator over the leaves of a `Tree`.