From a1e14cc3690df63053ef7887d1154eb25273dbb9 Mon Sep 17 00:00:00 2001 From: pscott <30843220+pscott@users.noreply.github.com> Date: Tue, 5 Nov 2019 02:29:07 +0100 Subject: [PATCH] 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_chain/src/beacon_chain_builder.rs | 21 ++- eth2/utils/merkle_proof/src/lib.rs | 130 +++++++++++++++++- 2 files changed, 138 insertions(+), 13 deletions(-) diff --git a/beacon_node/beacon_chain/src/beacon_chain_builder.rs b/beacon_node/beacon_chain/src/beacon_chain_builder.rs index 949fb856b..2ee695061 100644 --- a/beacon_node/beacon_chain/src/beacon_chain_builder.rs +++ b/beacon_node/beacon_chain/src/beacon_chain_builder.rs @@ -210,22 +210,19 @@ fn interop_genesis_state( .collect::>(); 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" ); diff --git a/eth2/utils/merkle_proof/src/lib.rs b/eth2/utils/merkle_proof/src/lib.rs index 13f67c584..785072eb4 100644 --- a/eth2/utils/merkle_proof/src/lib.rs +++ b/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, 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()); + } }