From 443f45f4306c104ce1a9e5c02b3eb6b7fb070168 Mon Sep 17 00:00:00 2001
From: NikVolf <nikvolf@gmail.com>
Date: Wed, 25 Sep 2019 09:50:59 +0200
Subject: [PATCH] dry example a bit and reduce api

---
 examples/long.rs | 55 +++++++++++++++++-------------------------------
 src/entry.rs     | 18 +++++++++++++---
 2 files changed, 34 insertions(+), 39 deletions(-)

diff --git a/examples/long.rs b/examples/long.rs
index 85f2c13..951d226 100644
--- a/examples/long.rs
+++ b/examples/long.rs
@@ -3,6 +3,22 @@ use zcash_mmr::{Entry, EntryLink, NodeData, Tree};
 #[path= "lib/shared.rs"]
 mod share;
 
+fn draft(into: &mut Vec<(u32, Entry)>, vec: &Vec<NodeData>, peak_pos: usize, h: u32) {
+    let node_data = vec[peak_pos-1].clone();
+    let peak: Entry = match h {
+        0 => node_data.into(),
+        _ => Entry::new(
+            node_data,
+            EntryLink::Stored((peak_pos - (1 << h) - 1) as u32),
+            EntryLink::Stored((peak_pos - 2) as u32),
+        ),
+    };
+
+    println!("Entry #{}: {}", into.len(), peak);
+
+    into.push(((peak_pos-1) as u32, peak));
+}
+
 fn prepare_tree(vec: &Vec<NodeData>) -> Tree {
 
     assert!(vec.len() > 0);
@@ -25,19 +41,7 @@ fn prepare_tree(vec: &Vec<NodeData>) -> Tree {
         }
 
         if peak_pos <= vec.len() {
-            let mut peak: Entry = vec[peak_pos-1].clone().into();
-            if h != 0 {
-                let left_idx = (peak_pos - (1<<h) - 1) as u32;
-                let right_idx = (peak_pos - 2) as u32;
-
-                peak.update_siblings(
-                    EntryLink::Stored(left_idx),
-                    EntryLink::Stored(right_idx),
-                );
-
-                println!("Peak #{}: ({}, {})", nodes.len(), left_idx, right_idx);
-            }
-            nodes.push(((peak_pos-1) as u32, peak));
+            draft(&mut nodes, vec, peak_pos, h);
 
             // save to be used in next loop
             last_peak_pos = peak_pos;
@@ -63,30 +67,10 @@ fn prepare_tree(vec: &Vec<NodeData>) -> Tree {
         h = h - 1;
 
         // drafting left child
-        let mut peak: Entry = vec[left_pos-1].clone().into();
-        if h != 0 {
-            let left_idx = (left_pos - (1<<h) - 1) as u32;
-            let right_idx = (left_pos - 2) as u32;
-
-            peak.update_siblings(
-                EntryLink::Stored(left_idx),
-                EntryLink::Stored(right_idx),
-            );
-        }
-        extra.push(((left_pos-1) as u32, peak));
+        draft(&mut extra, vec, left_pos, h);
 
         // drafting right child
-        let mut peak: Entry = vec[right_pos-1].clone().into();
-        if h != 0 {
-            let left_idx = (right_pos - (1<<h) - 1) as u32;
-            let right_idx = (right_pos - 2) as u32;
-
-            peak.update_siblings(
-                EntryLink::Stored(left_idx),
-                EntryLink::Stored(right_idx),
-            );
-        }
-        extra.push(((right_pos-1) as u32, peak));
+        draft(&mut extra, vec, right_pos, h);
 
         // continuing on right slope
         peak_pos = right_pos;
@@ -94,7 +78,6 @@ fn prepare_tree(vec: &Vec<NodeData>) -> Tree {
 
     println!("Total extra of {} required for deletion!", extra.len());
 
-
     Tree::new(vec.len() as u32, nodes, extra)
 }
 
diff --git a/src/entry.rs b/src/entry.rs
index 04490b3..b208555 100644
--- a/src/entry.rs
+++ b/src/entry.rs
@@ -13,9 +13,12 @@ pub struct Entry {
 }
 
 impl Entry {
-    /// Update siblings of the entry (to promote it from leaf to node)
-    pub fn update_siblings(&mut self, left: EntryLink, right: EntryLink) {
-        self.kind = EntryKind::Node(left, right);
+    /// New entry of type node.
+    pub fn new(data: NodeData, left: EntryLink, right: EntryLink) -> Self {
+        Entry {
+            kind: EntryKind::Node(left, right),
+            data,
+        }
     }
 
     /// Returns if is this node complete (has total of 2^N leaves)
@@ -107,3 +110,12 @@ impl From<NodeData> for Entry {
         Entry { kind: EntryKind::Leaf, data: s }
     }
 }
+
+impl std::fmt::Display for Entry {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self.kind {
+            EntryKind::Node(l, r) => write!(f, "node({}, {}, ..)", l, r),
+            EntryKind::Leaf => write!(f, "leaf(..)"),
+        }
+    }
+}