Merge pull request #2 from NikVolf/review-fixes

Address offline review

README.md

@@ -1,10 +1,10 @@
 # zcash_mmr
 
-Special implementation of merkle mountain ranges (MMR) for ZCash!
+Special implementation of Merkle mountain ranges (MMR) for Zcash!
 
 [![Build Status](https://travis-ci.org/NikVolf/zcash-mmr.svg?branch=master)](https://travis-ci.org/NikVolf/zcash-mmr)
 
-The main design goals of this mmr implementation are
+The main design goals of this MMR implementation are
 
 - Allow zero-cache and avoid db callbacks. As it is implemented, calling side must just smartly pre-load MMR nodes from the database (about log2(tree length) for append, twice as much for deletion).
 
@@ -14,7 +14,7 @@ The main design goals of this mmr implementation are
 
 # License
 
-`zcash_mmr` is primarily distributed under the terms of both the MIT
+`zcash_mmr` is distributed under the terms of both the MIT
 license and the Apache License (Version 2.0), at your choice.
 
 See LICENSE-APACHE, and LICENSE-MIT for details.

src/entry.rs

@@ -29,7 +29,7 @@ impl Entry {
 
     /// Number of leaves under this node.
     pub fn leaf_count(&self) -> u64 {
-        self.data.end_height - self.data.start_height + 1
+        self.data.end_height - (self.data.start_height - 1)
     }
 
     /// Is this node a leaf.

src/node_data.rs

@@ -3,11 +3,24 @@ use bigint::U256;
 use blake2::Params as Blake2Params;
 
 /// Maximum serialized size of the node metadata.
-pub const MAX_NODE_DATA_SIZE: usize = 32 + 4 + 4 + 4 + 4 + 32 + 32 + 32 + 9 + 9 + 9; // 171
+pub const MAX_NODE_DATA_SIZE: usize =
+    32 + // subtree commitment
+    4 +  // start time
+    4 +  // end time
+    4 +  // start target
+    4 +  // end target
+    32 + // start sapling tree root
+    32 + // end sapling tree root
+    32 + // subtree total work
+    9 +  // start height (compact uint)
+    9 +  // end height (compact uint)
+    9;   // shielded tx count (compact uint)
+    // = total of 171
 
 /// Node metadata.
 #[repr(C)]
 #[derive(Debug, Clone, Default)]
+#[cfg_attr(test, derive(PartialEq))]
 pub struct NodeData {
     /// Consensus branch id, should be provided by deserializing node.
     pub consensus_branch_id: u32,
@@ -114,7 +127,7 @@ impl NodeData {
             i @ 0..=0xfc => i.into(),
             0xfd => reader.read_u16::<LittleEndian>()?.into(),
             0xfe => reader.read_u32::<LittleEndian>()?.into(),
-            _ => reader.read_u64::<LittleEndian>()?.into(),
+            _ => reader.read_u64::<LittleEndian>()?,
         };
 
         Ok(result)
@@ -181,3 +194,39 @@ impl NodeData {
         Self::read(consensus_branch_id, &mut cursor)
     }
 }
+
+#[cfg(test)]
+impl quickcheck::Arbitrary for NodeData {
+    fn arbitrary<G: quickcheck::Gen>(gen: &mut G) -> Self {
+        let mut node_data = NodeData::default();
+        node_data.consensus_branch_id = 0;
+        gen.fill_bytes(&mut node_data.subtree_commitment[..]);
+        node_data.start_time = gen.next_u32();
+        node_data.end_time = gen.next_u32();
+        node_data.start_target = gen.next_u32();
+        node_data.end_target = gen.next_u32();
+        gen.fill_bytes(&mut node_data.start_sapling_root[..]);
+        gen.fill_bytes(&mut node_data.end_sapling_root[..]);
+        let mut number = [0u8; 32];
+        gen.fill_bytes(&mut number[..]);
+        node_data.subtree_total_work = U256::from_little_endian(&number[..]);
+        node_data.start_height = gen.next_u64();
+        node_data.end_height = gen.next_u64();
+        node_data.shielded_tx = gen.next_u64();
+
+        node_data
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::NodeData;
+    use quickcheck::{quickcheck, TestResult};
+
+
+    quickcheck! {
+        fn serialization_round_trip(node_data: NodeData) -> TestResult {
+            TestResult::from_bool(NodeData::from_bytes(0, &node_data.to_bytes()).unwrap() == node_data)
+        }
+    }
+}

src/tree.rs

@@ -8,18 +8,20 @@ use crate::{Entry, EntryLink, NodeData, Error, EntryKind};
 /// With only some of the leaves/nodes pre-loaded / pre-generated.
 /// Exact amount of the loaded data can be calculated by the constructing party,
 /// depending on the length of the tree and maximum amount of operations that are going
-/// to happen after construction.
+/// to happen after construction. `Tree` should not be used as self-contained data structure,
+/// since it's internal state can grow indefinitely after serial operations.
+/// Intended use of this `Tree` is to instantiate it based on partially loaded data (see example
+/// how to pick right nodes from the array representation of MMR Tree), perform several operations
+/// (append-s/delete-s) and then drop it.
 pub struct Tree {
     stored: HashMap<u32, Entry>,
 
-    generated: HashMap<u32, Entry>,
+    // This can grow indefinitely if `Tree` is misused as a self-contained data structure
+    generated: Vec<Entry>,
 
     // number of persistent(!) tree entries
     stored_count: u32,
 
-    // number of virtual nodes generated
-    generated_count: u32,
-
     root: EntryLink,
 }
 
@@ -28,7 +30,7 @@ impl Tree {
     pub fn resolve_link(&self, link: EntryLink) -> Result<IndexedNode, Error> {
         match link {
             EntryLink::Generated(index) => {
-                let node = self.generated.get(&index).ok_or(Error::ExpectedInMemory(link))?;
+                let node = self.generated.get(index as usize).ok_or(Error::ExpectedInMemory(link))?;
                 Ok(IndexedNode {
                     node,
                     link,
@@ -46,20 +48,19 @@ impl Tree {
 
     fn push(&mut self, data: Entry) -> EntryLink {
         let idx = self.stored_count;
-        self.stored_count = self.stored_count + 1;
+        self.stored_count += 1;
         self.stored.insert(idx, data);
         EntryLink::Stored(idx)
     }
 
     fn push_generated(&mut self, data: Entry) -> EntryLink {
-        let idx = self.generated_count;
+        self.generated.push(data);
-        self.generated_count = self.generated_count + 1;
+        EntryLink::Generated(self.generated.len() as u32 - 1)
-        self.generated.insert(idx, data);
-        EntryLink::Generated(idx)
     }
 
-    /// Populate tree with plain list of the leaves/nodes. Mostly for tests,
+    /// Populate tree with plain list of the leaves/nodes. For now, only for tests,
-    /// since this `Tree` structure is for partially loaded tree.
+    /// since this `Tree` structure is for partially loaded tree (but it might change)
+    #[cfg(test)]
     pub fn populate(loaded: Vec<Entry>, root: EntryLink) -> Self {
         let mut result = Tree::invalid();
         result.stored_count = loaded.len() as u32;
@@ -71,27 +72,34 @@ impl Tree {
         result
     }
 
+    // Empty tree with invalid root
     fn invalid() -> Self {
         Tree {
             root: EntryLink::Generated(0),
             generated: Default::default(),
             stored: Default::default(),
-            generated_count: 0,
             stored_count: 0,
         }
     }
 
     /// New view into the the tree array representation
     ///
-    /// `length` is total length of the array representation
+    /// `length` is total length of the array representation (is generally not a sum of
+    ///     peaks.len + extra.len)
     /// `peaks` is peaks of the mmr tree
     /// `extra` is some extra nodes that calculated to be required during next one or more
     /// operations on the tree.
+    ///
+    /// # Panics
+    ///
+    /// Will panic if `peaks` is empty.
     pub fn new(
         length: u32,
         peaks: Vec<(u32, Entry)>,
         extra: Vec<(u32, Entry)>,
     ) -> Self {
+        assert!(peaks.len() > 0);
+
         let mut result = Tree::invalid();
 
         result.stored_count = length;
@@ -156,6 +164,9 @@ impl Tree {
         let mut merge_stack = Vec::new();
         merge_stack.push(new_leaf_link);
 
+        // Scan the peaks right-to-left, merging together equal-sized adjacent
+        // complete subtrees. After this, merge_stack only contains peaks of
+        // unequal-sized subtrees.
         while let Some(next_peak) = peaks.pop() {
             let next_merge = merge_stack.pop().expect("there should be at least one, initial or re-pushed");
 
@@ -170,12 +181,15 @@ impl Tree {
                 merge_stack.push(link);
                 appended.push(link);
                 continue;
-            }
+            } else {
                 merge_stack.push(next_merge);
                 merge_stack.push(next_peak);
             }
+        }
 
         let mut new_root = merge_stack.pop().expect("Loop above cannot reduce the merge_stack");
+        // Scan the peaks left-to-right, producing new generated nodes that
+        // connect the subtrees
         while let Some(next_child) = merge_stack.pop() {
             new_root = self.push_generated(
                 combine_nodes(
@@ -191,7 +205,7 @@ impl Tree {
     }
 
     #[cfg(test)]
-    fn for_children<F: FnMut(EntryLink, EntryLink)>(&mut self, node: EntryLink, mut f: F) {
+    fn for_children<F: Fn(EntryLink, EntryLink)>(&self, node: EntryLink, f: F) {
         let (left, right) = {
             let link = self.resolve_link(node).expect("Failed to resolve link in test");
             (
@@ -209,7 +223,8 @@ impl Tree {
 
     /// Truncate one leaf from the end of the tree.
     ///
-    /// Returns actual number of nodes that has to be removed from the array representation.
+    /// Returns actual number of nodes that should be removed by the caller
+    /// from the end of the array representation.
     pub fn truncate_leaf(&mut self) -> Result<u32, Error> {
         let root = {
             let (leaves, root_left_child) = {
@@ -244,7 +259,7 @@ impl Tree {
             }
         }
 
-        let mut new_root = *peaks.iter().nth(0).expect("At lest 2 elements in peaks");
+        let mut new_root = *peaks.get(0).expect("At lest 1 elements in peaks");
 
         for next_peak in peaks.into_iter().skip(1) {
             new_root = self.push_generated(
@@ -270,10 +285,15 @@ impl Tree {
     /// Link to the root node
     pub fn root(&self) -> EntryLink { self.root }
 
-    /// Reference to the root ndoe
+    /// Reference to the root node.
     pub fn root_node(&self) -> Result<IndexedNode, Error> {
         self.resolve_link(self.root)
     }
+
+    /// If this tree is empty.
+    pub fn is_empty(&self) -> bool {
+        self.stored_count == 0
+    }
 }
 
 /// Reference to the node with link attached.
@@ -340,29 +360,12 @@ mod tests {
         }
     }
 
-    fn node(start_height: u64, end_height: u64) -> NodeData {
-        NodeData {
-            consensus_branch_id: 1,
-            subtree_commitment: [0u8; 32],
-            start_time: 0,
-            end_time: 0,
-            start_target: 0,
-            end_target: 0,
-            start_sapling_root: [0u8; 32],
-            end_sapling_root: [0u8; 32],
-            subtree_total_work: 0.into(),
-            start_height: start_height,
-            end_height: end_height,
-            shielded_tx: 7,
-        }
-    }
-
     fn initial() -> Tree {
         let node1: Entry = leaf(1).into();
         let node2: Entry = leaf(2).into();
 
         let node3 = Entry {
-            data: node(1, 2),
+            data: NodeData::combine(&node1.data, &node2.data),
             kind: EntryKind::Leaf,
         };
 
@@ -482,7 +485,7 @@ mod tests {
         //           (0)  (1) (3)   (4)  (7)
         //
         // new tree:
-        //                     (---8g---)
+        //                     (---10g--)
         //                    /          \
         //                 ( 6 )          \
         //                /     \          \
@@ -491,7 +494,7 @@ mod tests {
         //           (0)  (1) (3)   (4)  (7)  (8)
         //
         // so (7) is added as real leaf
-        // and new root, (8g) is generated one
+        // and new root, (10g) is generated one
         assert_eq!(new_root.data.end_height, 6);
         assert_eq!(appended.len(), 2);
         assert_matches!(tree.root(), EntryLink::Generated(_));
@@ -530,13 +533,16 @@ mod tests {
         //             /  \     /   \     /   \     \
         //           (0)  (1) (3)   (4) (7)   (8)  (10)
         //
-        // so (7) is added as real leaf
+        // so (10) is added as real leaf
-        // and new root, (8g) is generated one
+        // and new root, (12g) is generated one
         assert_eq!(new_root.data.end_height, 7);
         assert_eq!(appended.len(), 1);
         assert_matches!(tree.root(), EntryLink::Generated(_));
         tree.for_children(tree.root(), |l, r| {
             assert_matches!(l, EntryLink::Generated(_));
+            tree.for_children(l, |l, r|
+                assert_matches!((l, r), (EntryLink::Stored(6), EntryLink::Stored(9)))
+            );
             assert_matches!(r, EntryLink::Stored(10));
         });
     }
@@ -544,7 +550,7 @@ mod tests {
     #[test]
     fn truncate_simple() {
         let mut tree = generated(9);
-        tree.truncate_leaf().expect("Failed to truncate");
+        let total_truncated = tree.truncate_leaf().expect("Failed to truncate");
 
         // initial tree:
         //
@@ -571,6 +577,7 @@ mod tests {
         // and new root, (14) is a stored one now
 
         assert_matches!(tree.root(), EntryLink::Stored(14));
+        assert_eq!(total_truncated, 1);
         assert_eq!(tree.len(), 15);
     }
 
@@ -606,19 +613,11 @@ mod tests {
 
         assert_matches!(tree.root(), EntryLink::Generated(_));
 
-        // left is 14 and right is 15
+        tree.for_children(tree.root(),|left, right|
-        let (left_root_child, right_root_child) = {
-            let root = tree.root_node().expect("Failed to resolve");
-
-            (
-                root.left().expect("Expected node"),
-                root.right().expect("Expected node"),
-            )
-        };
-
             assert_matches!(
-            (left_root_child, right_root_child),
+                (left, right),
                 (EntryLink::Stored(14), EntryLink::Stored(15))
+            )
         );
 
         // two stored nodes should leave us (leaf 16 and no longer needed node 17)
@@ -703,7 +702,7 @@ mod tests {
                 }
 
                 TestResult::from_bool(
-                    if number & number - 1 == 0 {
+                    if number & (number - 1) == 0 {
                         if let EntryLink::Stored(_) = tree.root() { true }
                         else { false }
                     } else {