merkle_proof: implement tree construction

Plus QuickCheck tests!

eth2/utils/merkle_proof/Cargo.toml

@@ -7,3 +7,8 @@ edition = "2018"
 [dependencies]
 ethereum-types = "0.6"
 eth2_hashing = { path = "../eth2_hashing" }
+lazy_static = "1.3.0"
+
+[dev-dependencies]
+quickcheck = "0.8"
+quickcheck_macros = "0.8"

eth2/utils/merkle_proof/src/lib.rs

@@ -1,6 +1,138 @@
+#[macro_use]
+extern crate lazy_static;
+
 use eth2_hashing::hash;
 use ethereum_types::H256;
 
+const MAX_TREE_DEPTH: usize = 32;
+const EMPTY_SLICE: &[H256] = &[];
+
+lazy_static! {
+    /// Cached zero hashes where `ZERO_HASHES[i]` is the hash of a Merkle tree with 2^i zero leaves.
+    static ref ZERO_HASHES: Vec<H256> = {
+        let mut hashes = vec![H256::from([0; 32]); MAX_TREE_DEPTH + 1];
+
+        for i in 0..MAX_TREE_DEPTH {
+            hashes[i + 1] = hash_concat(hashes[i], hashes[i]);
+        }
+
+        hashes
+    };
+
+    /// Zero nodes to act as "synthetic" left and right subtrees of other zero nodes.
+    static ref ZERO_NODES: Vec<MerkleTree> = {
+        (0..MAX_TREE_DEPTH + 1).map(MerkleTree::Zero).collect()
+    };
+}
+
+/// Right-sparse Merkle tree.
+///
+/// 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)]
+pub enum MerkleTree {
+    /// Leaf node with the hash of its content.
+    Leaf(H256),
+    /// Internal node with hash, left subtree and right subtree.
+    Node(H256, Box<Self>, Box<Self>),
+    /// Zero subtree of a given depth.
+    ///
+    /// It represents a Merkle tree of 2^depth zero leaves.
+    Zero(usize),
+}
+
+impl MerkleTree {
+    /// Create a new Merkle tree from a list of leaves and a fixed depth.
+    pub fn create(leaves: &[H256], depth: usize) -> Self {
+        use MerkleTree::*;
+
+        if leaves.is_empty() {
+            return Zero(depth);
+        }
+
+        match depth {
+            0 => {
+                debug_assert_eq!(leaves.len(), 1);
+                Leaf(leaves[0])
+            }
+            _ => {
+                // Split leaves into left and right subtrees
+                let subtree_capacity = 2usize.pow(depth as u32 - 1);
+                let (left_leaves, right_leaves) = if leaves.len() <= subtree_capacity {
+                    (leaves, EMPTY_SLICE)
+                } else {
+                    leaves.split_at(subtree_capacity)
+                };
+
+                let left_subtree = MerkleTree::create(left_leaves, depth - 1);
+                let right_subtree = MerkleTree::create(right_leaves, depth - 1);
+                let hash = hash_concat(left_subtree.hash(), right_subtree.hash());
+
+                Node(hash, Box::new(left_subtree), Box::new(right_subtree))
+            }
+        }
+    }
+
+    /// Retrieve the root hash of this Merkle tree.
+    pub fn hash(&self) -> H256 {
+        match *self {
+            MerkleTree::Leaf(h) => h,
+            MerkleTree::Node(h, _, _) => h,
+            MerkleTree::Zero(depth) => ZERO_HASHES[depth],
+        }
+    }
+
+    /// Get a reference to the left and right subtrees if they exist.
+    pub fn left_and_right_branches(&self) -> Option<(&Self, &Self)> {
+        match *self {
+            MerkleTree::Leaf(_) | MerkleTree::Zero(0) => None,
+            MerkleTree::Node(_, ref l, ref r) => Some((l, r)),
+            MerkleTree::Zero(depth) => Some((&ZERO_NODES[depth - 1], &ZERO_NODES[depth - 1])),
+        }
+    }
+
+    /// Is this Merkle tree a leaf?
+    pub fn is_leaf(&self) -> bool {
+        match self {
+            MerkleTree::Leaf(_) => true,
+            _ => false,
+        }
+    }
+
+    /// Return the leaf at `index` and a Merkle proof of its inclusion.
+    ///
+    /// The Merkle proof is in "bottom-up" order, starting with a leaf node
+    /// and moving up the tree. Its length will be exactly equal to `depth`.
+    pub fn generate_proof(&self, index: usize, depth: usize) -> (H256, Vec<H256>) {
+        let mut proof = vec![];
+        let mut current_node = self;
+        let mut current_depth = depth;
+        while current_depth > 0 {
+            let ith_bit = (index >> (current_depth - 1)) & 0x01;
+            // Note: unwrap is safe because leaves are only ever constructed at depth == 0.
+            let (left, right) = current_node.left_and_right_branches().unwrap();
+
+            // Go right, include the left branch in the proof.
+            if ith_bit == 1 {
+                proof.push(left.hash());
+                current_node = right;
+            } else {
+                proof.push(right.hash());
+                current_node = left;
+            }
+            current_depth -= 1;
+        }
+
+        debug_assert_eq!(proof.len(), depth);
+        debug_assert!(current_node.is_leaf());
+
+        // Put proof in bottom-up order.
+        proof.reverse();
+
+        (current_node.hash(), proof)
+    }
+}
+
 /// Verify a proof that `leaf` exists at `index` in a Merkle tree rooted at `root`.
 ///
 /// The `branch` argument is the main component of the proof: it should be a list of internal
@@ -46,15 +178,66 @@ fn concat(mut vec1: Vec<u8>, mut vec2: Vec<u8>) -> Vec<u8> {
     vec1
 }
 
+/// Compute the hash of two other hashes concatenated.
+fn hash_concat(h1: H256, h2: H256) -> H256 {
+    H256::from_slice(&hash(&concat(
+        h1.as_bytes().to_vec(),
+        h2.as_bytes().to_vec(),
+    )))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
+    use quickcheck::TestResult;
+    use quickcheck_macros::quickcheck;
 
-    fn hash_concat(h1: H256, h2: H256) -> H256 {
-        H256::from_slice(&hash(&concat(
-            h1.as_bytes().to_vec(),
-            h2.as_bytes().to_vec(),
-        )))
+    /// Check that we can:
+    /// 1. Build a MerkleTree from arbitrary leaves and an arbitrary depth.
+    /// 2. Generate valid proofs for all of the leaves of this MerkleTree.
+    #[quickcheck]
+    fn quickcheck_create_and_verify(int_leaves: Vec<u64>, depth: usize) -> TestResult {
+        if 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 merkle_tree = MerkleTree::create(&leaves, depth);
+        let merkle_root = merkle_tree.hash();
+
+        let proofs_ok = (0..leaves.len()).into_iter().all(|i| {
+            let (leaf, branch) = merkle_tree.generate_proof(i, depth);
+            leaf == leaves[i] && verify_merkle_proof(leaf, &branch, depth, i, merkle_root)
+        });
+
+        TestResult::from_bool(proofs_ok)
+    }
+
+    #[test]
+    fn sparse_zero_correct() {
+        let depth = 2;
+        let zero = H256::from([0x00; 32]);
+        let dense_tree = MerkleTree::create(&[zero, zero, zero, zero], depth);
+        let sparse_tree = MerkleTree::create(&[], depth);
+        assert_eq!(dense_tree.hash(), sparse_tree.hash());
+    }
+
+    #[test]
+    fn create_small_example() {
+        // Construct a small merkle tree manually and check that it's consistent with
+        // the MerkleTree type.
+        let leaf_b00 = H256::from([0xAA; 32]);
+        let leaf_b01 = H256::from([0xBB; 32]);
+        let leaf_b10 = H256::from([0xCC; 32]);
+        let leaf_b11 = H256::from([0xDD; 32]);
+
+        let node_b0x = hash_concat(leaf_b00, leaf_b01);
+        let node_b1x = hash_concat(leaf_b10, leaf_b11);
+
+        let root = hash_concat(node_b0x, node_b1x);
+
+        let tree = MerkleTree::create(&[leaf_b00, leaf_b01, leaf_b10, leaf_b11], 2);
+        assert_eq!(tree.hash(), root);
     }
 
     #[test]