Add MerkleTree::gen_nth_proof and Proof::index.

src/merkletree.rs

@@ -160,6 +160,22 @@ impl<T> MerkleTree<T> {
             .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<Proof<T>>
+    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<T> {
         self.root.iter()

src/proof.rs

@@ -153,6 +153,11 @@ impl<T> Proof<T> {
 
         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<T>(tree: &Tree<T>, needle: &[u8]) -> Option<Lemma> {
         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<T>(tree: &Tree<T>, idx: usize, count: usize) -> Option<Lemma> {
+        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<Lemma> {
         if *hash == *needle {
             Some(Lemma {
@@ -245,6 +292,35 @@ impl Lemma {
             },
         }
     }
+
+    fn new_tree_proof_by_index<T>(
+        hash: &[u8],
+        idx: usize,
+        count: usize,
+        left: &Tree<T>,
+        right: &Tree<T>,
+    ) -> Option<Lemma> {
+        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.

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::<Vec<_>>();
+        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::<Vec<_>>();

src/tree.rs

@@ -61,6 +61,31 @@ impl<T> Tree<T> {
     pub fn iter(&self) -> LeavesIterator<T> {
         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`.