Implement push_leaf for MerkleTree (#574)

* Prototype for far_right push

* Add push method and tests

* Modify beacon_chain_builder for interop to use push instead of create

* Add Push method to MerkleTree

* Cargo fmt

* Remove redundant tests

* Fix typo

* Rename push to push_leaf

* Fix clippy warnings

* Add DepthTooSmall enum variant

* Avoid cloning in MerkleTree::push_leaf

* Add quickcheck test for push_leaf

* Cargo fmt updated

* Return err instead of using unwrap()

* Use enumerate instead of hard indexing

* Use if let and return string on error

* Fix typo in deposit_leave

* Fix cargo fmt

beacon_node/beacon_chain/src/beacon_chain_builder.rs

@@ -210,22 +210,19 @@ fn interop_genesis_state<T: EthSpec>(
         .collect::<Vec<_>>();
 
     let mut proofs = vec![];
-    for i in 1..=deposit_root_leaves.len() {
-        // Note: this implementation is not so efficient.
-        //
-        // If `MerkleTree` had a push method, we could just build one tree and sample it instead of
-        // rebuilding the tree for each deposit.
-        let tree = MerkleTree::create(
-            &deposit_root_leaves[0..i],
-            spec.deposit_contract_tree_depth as usize,
-        );
+    let depth = spec.deposit_contract_tree_depth as usize;
+    let mut tree = MerkleTree::create(&[], depth);
+    for (i, deposit_leaf) in deposit_root_leaves.iter().enumerate() {
+        if let Err(_) = tree.push_leaf(*deposit_leaf, depth) {
+            return Err(String::from("Failed to push leaf"));
+        }
 
-        let (_, mut proof) = tree.generate_proof(i - 1, spec.deposit_contract_tree_depth as usize);
-        proof.push(Hash256::from_slice(&int_to_bytes32(i)));
+        let (_, mut proof) = tree.generate_proof(i, depth);
+        proof.push(Hash256::from_slice(&int_to_bytes32(i + 1)));
 
         assert_eq!(
             proof.len(),
-            spec.deposit_contract_tree_depth as usize + 1,
+            depth + 1,
             "Deposit proof should be correct len"
         );
 

eth2/utils/merkle_proof/src/lib.rs

@@ -29,7 +29,7 @@ lazy_static! {
 ///
 /// Efficiently represents a Merkle tree of fixed depth where only the first N
 /// indices are populated by non-zero leaves (perfect for the deposit contract tree).
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub enum MerkleTree {
     /// Leaf node with the hash of its content.
     Leaf(H256),
@@ -41,6 +41,18 @@ pub enum MerkleTree {
     Zero(usize),
 }
 
+#[derive(Debug, PartialEq)]
+pub enum MerkleTreeError {
+    // Trying to push in a leaf
+    LeafReached,
+    // No more space in the MerkleTree
+    MerkleTreeFull,
+    // MerkleTree is invalid
+    Invalid,
+    // Incorrect Depth provided
+    DepthTooSmall,
+}
+
 impl MerkleTree {
     /// Create a new Merkle tree from a list of leaves and a fixed depth.
     pub fn create(leaves: &[H256], depth: usize) -> Self {
@@ -73,6 +85,62 @@ impl MerkleTree {
         }
     }
 
+    /// Push an element in the MerkleTree.
+    /// MerkleTree and depth must be correct, as the algorithm expects valid data.
+    pub fn push_leaf(&mut self, elem: H256, depth: usize) -> Result<(), MerkleTreeError> {
+        use std::mem;
+        use MerkleTree::*;
+
+        if depth == 0 {
+            return Err(MerkleTreeError::DepthTooSmall);
+        }
+
+        match self {
+            Leaf(_) => return Err(MerkleTreeError::LeafReached),
+            Zero(_) => {
+                mem::replace(self, MerkleTree::create(&[elem], depth));
+            }
+            Node(ref mut hash, ref mut left, ref mut right) => {
+                let left: &mut MerkleTree = &mut *left;
+                let right: &mut MerkleTree = &mut *right;
+                match (&*left, &*right) {
+                    // Tree is full
+                    (Leaf(_), Leaf(_)) => return Err(MerkleTreeError::MerkleTreeFull),
+                    // There is a right node so insert in right node
+                    (Node(_, _, _), Node(_, _, _)) => {
+                        if let Err(e) = right.push_leaf(elem, depth - 1) {
+                            return Err(e);
+                        }
+                    }
+                    // Both branches are zero, insert in left one
+                    (Zero(_), Zero(_)) => {
+                        mem::replace(left, MerkleTree::create(&[elem], depth - 1));
+                    }
+                    // Leaf on left branch and zero on right branch, insert on right side
+                    (Leaf(_), Zero(_)) => {
+                        mem::replace(right, MerkleTree::create(&[elem], depth - 1));
+                    }
+                    // Try inserting on the left node -> if it fails because it is full, insert in right side.
+                    (Node(_, _, _), Zero(_)) => {
+                        match left.push_leaf(elem, depth - 1) {
+                            Ok(_) => (),
+                            // Left node is full, insert in right node
+                            Err(MerkleTreeError::MerkleTreeFull) => {
+                                mem::replace(right, MerkleTree::create(&[elem], depth - 1));
+                            }
+                            Err(e) => return Err(e),
+                        };
+                    }
+                    // All other possibilities are invalid MerkleTrees
+                    (_, _) => return Err(MerkleTreeError::Invalid),
+                };
+                *hash = hash_concat(left.hash(), right.hash());
+            }
+        }
+
+        Ok(())
+    }
+
     /// Retrieve the root hash of this Merkle tree.
     pub fn hash(&self) -> H256 {
         match *self {
@@ -213,6 +281,25 @@ mod tests {
         TestResult::from_bool(proofs_ok)
     }
 
+    #[quickcheck]
+    fn quickcheck_push_leaf_and_verify(int_leaves: Vec<u64>, depth: usize) -> TestResult {
+        if depth == 0 || depth > MAX_TREE_DEPTH || int_leaves.len() > 2usize.pow(depth as u32) {
+            return TestResult::discard();
+        }
+
+        let leaves: Vec<_> = int_leaves.into_iter().map(H256::from_low_u64_be).collect();
+
+        let mut merkle_tree = MerkleTree::create(&[], depth);
+
+        let proofs_ok = leaves.into_iter().enumerate().all(|(i, leaf)| {
+            assert_eq!(merkle_tree.push_leaf(leaf, depth), Ok(()));
+            let (stored_leaf, branch) = merkle_tree.generate_proof(i, depth);
+            stored_leaf == leaf && verify_merkle_proof(leaf, &branch, depth, i, merkle_tree.hash())
+        });
+
+        TestResult::from_bool(proofs_ok)
+    }
+
     #[test]
     fn sparse_zero_correct() {
         let depth = 2;
@@ -328,4 +415,45 @@ mod tests {
         assert!(verify_merkle_proof(leaf, &[], 0, 0, leaf));
         assert!(!verify_merkle_proof(leaf, &[], 0, 7, junk));
     }
+
+    #[test]
+    fn push_complete_example() {
+        let depth = 2;
+        let mut tree = MerkleTree::create(&[], depth);
+
+        let leaf_b00 = H256::from([0xAA; 32]);
+
+        let res = tree.push_leaf(leaf_b00, 0);
+        assert_eq!(res, Err(MerkleTreeError::DepthTooSmall));
+        let expected_tree = MerkleTree::create(&[], depth);
+        assert_eq!(tree.hash(), expected_tree.hash());
+
+        tree.push_leaf(leaf_b00, depth)
+            .expect("Pushing in empty tree failed");
+        let expected_tree = MerkleTree::create(&[leaf_b00], depth);
+        assert_eq!(tree.hash(), expected_tree.hash());
+
+        let leaf_b01 = H256::from([0xBB; 32]);
+        tree.push_leaf(leaf_b01, depth)
+            .expect("Pushing in left then right node failed");
+        let expected_tree = MerkleTree::create(&[leaf_b00, leaf_b01], depth);
+        assert_eq!(tree.hash(), expected_tree.hash());
+
+        let leaf_b10 = H256::from([0xCC; 32]);
+        tree.push_leaf(leaf_b10, depth)
+            .expect("Pushing in right then left node failed");
+        let expected_tree = MerkleTree::create(&[leaf_b00, leaf_b01, leaf_b10], depth);
+        assert_eq!(tree.hash(), expected_tree.hash());
+
+        let leaf_b11 = H256::from([0xDD; 32]);
+        tree.push_leaf(leaf_b11, depth)
+            .expect("Pushing in outtermost leaf failed");
+        let expected_tree = MerkleTree::create(&[leaf_b00, leaf_b01, leaf_b10, leaf_b11], depth);
+        assert_eq!(tree.hash(), expected_tree.hash());
+
+        let leaf_b12 = H256::from([0xEE; 32]);
+        let res = tree.push_leaf(leaf_b12, depth);
+        assert_eq!(res, Err(MerkleTreeError::MerkleTreeFull));
+        assert_eq!(tree.hash(), expected_tree.hash());
+    }
 }