From 4f60edafcdb3143a45c6930b88802307d4e76a01 Mon Sep 17 00:00:00 2001
From: Andreas Fackler <AndreasFackler@gmx.de>
Date: Mon, 11 Jun 2018 10:10:32 +0200
Subject: [PATCH] Traverse tree only once for proof generation.

---
 src/merkletree.rs |  8 +-----
 src/proof.rs      | 65 ++++++++++++++++++++++-------------------------
 src/tests.rs      |  3 +++
 src/tree.rs       | 29 ---------------------
 4 files changed, 34 insertions(+), 71 deletions(-)

diff --git a/src/merkletree.rs b/src/merkletree.rs
index 58435a0..a175448 100644
--- a/src/merkletree.rs
+++ b/src/merkletree.rs
@@ -166,14 +166,8 @@ impl<T> MerkleTree<T> {
         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()))
+            .map(|(lemma, value)| Proof::new(self.algorithm, root_hash, lemma, value.clone()))
     }
 
     /// Creates an `Iterator` over the values contained in this Merkle tree.
diff --git a/src/proof.rs b/src/proof.rs
index ede0474..936afab 100644
--- a/src/proof.rs
+++ b/src/proof.rs
@@ -194,30 +194,54 @@ 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> {
+    /// `tree`. If `idx >= count`, `None` is returned. Otherwise it returns
+    /// the new `Lemma` and the `idx`-th value.
+    pub fn new_by_index<T>(tree: &Tree<T>, idx: usize, count: usize) -> Option<(Lemma, &T)> {
         if idx >= count {
             return None;
         }
         match *tree {
             Tree::Empty { .. } => None,
 
-            Tree::Leaf { ref hash, .. } => {
+            Tree::Leaf {
+                ref hash,
+                ref value,
+                ..
+            } => {
                 if count != 1 {
                     return None;
                 }
-                Some(Lemma {
+                let lemma = Lemma {
                     node_hash: hash.clone(),
                     sibling_hash: None,
                     sub_lemma: None,
-                })
+                };
+                Some((lemma, value))
             }
 
             Tree::Node {
                 ref hash,
                 ref left,
                 ref right,
-            } => Lemma::new_tree_proof_by_index(hash, idx, count, left, right),
+            } => {
+                let left_count = count.next_power_of_two() / 2;
+                let (sub_lem_val, sibling_hash);
+                if idx < left_count {
+                    sub_lem_val = Lemma::new_by_index(left, idx, left_count);
+                    sibling_hash = Positioned::Right(right.hash().clone());
+                } else {
+                    sub_lem_val = Lemma::new_by_index(right, idx - left_count, count - left_count);
+                    sibling_hash = Positioned::Left(left.hash().clone());
+                }
+                sub_lem_val.map(|(sub_lemma, value)| {
+                    let lemma = Lemma {
+                        node_hash: hash.clone(),
+                        sibling_hash: Some(sibling_hash),
+                        sub_lemma: Some(Box::new(sub_lemma)),
+                    };
+                    (lemma, value)
+                })
+            }
         }
     }
 
@@ -293,35 +317,6 @@ 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.
diff --git a/src/tests.rs b/src/tests.rs
index 51f8413..c1e1c47 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -164,6 +164,9 @@ fn test_nth_proof() {
             assert!(proof.validate(&root_hash));
             assert_eq!(i, proof.index(tree.count()));
         }
+
+        assert!(tree.gen_nth_proof(count).is_none());
+        assert!(tree.gen_nth_proof(count + 1000).is_none());
     }
 }
 
diff --git a/src/tree.rs b/src/tree.rs
index 43a9cef..e65fb98 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -61,35 +61,6 @@ 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`.