Fix bug with cached tree hash, passes tests

eth2/utils/cached_tree_hash/src/btree_overlay.rs

@@ -29,6 +29,13 @@ impl Into<BTreeSchema> for BTreeOverlay {
     }
 }
 
+#[derive(Debug, PartialEq, Clone)]
+pub enum LeafNode {
+    DoesNotExist,
+    Exists(Range<usize>),
+    Padding,
+}
+
 #[derive(Debug, PartialEq, Clone)]
 pub struct BTreeOverlay {
     offset: usize,
@@ -82,6 +89,10 @@ impl BTreeOverlay {
         self.num_leaf_nodes().trailing_zeros() as usize
     }
 
+    pub fn internal_chunk_range(&self) -> Range<usize> {
+        self.offset..self.offset + self.num_internal_nodes()
+    }
+
     pub fn chunk_range(&self) -> Range<usize> {
         self.first_node()..self.next_node()
     }
@@ -104,13 +115,15 @@ impl BTreeOverlay {
     ///     - The specified node is internal.
     ///     - The specified node is padding.
     ///     - The specified node is OOB of the tree.
-    pub fn get_leaf_node(&self, i: usize) -> Result<Option<Range<usize>>, Error> {
-        if i >= self.num_nodes() - self.num_padding_leaves() {
-            Ok(None)
+    pub fn get_leaf_node(&self, i: usize) -> Result<LeafNode, Error> {
+        if i >= self.num_nodes() {
+            Ok(LeafNode::DoesNotExist)
+        } else if i >= self.num_nodes() - self.num_padding_leaves() {
+            Ok(LeafNode::Padding)
         } else if (i == self.num_internal_nodes()) && (self.lengths.len() == 0) {
             // If this is the first leaf node and the overlay contains zero items, return `None` as
             // this node must be padding.
-            Ok(None)
+            Ok(LeafNode::Padding)
         } else {
             let i = i - self.num_internal_nodes();
 
@@ -119,7 +132,7 @@ impl BTreeOverlay {
                 + self.lengths.iter().take(i).sum::<usize>();
             let last_node = first_node + self.lengths[i];
 
-            Ok(Some(first_node..last_node))
+            Ok(LeafNode::Exists(first_node..last_node))
         }
     }
 
@@ -237,10 +250,28 @@ mod test {
     fn get_leaf_node() {
         let tree = get_tree_a(4);
 
-        assert_eq!(tree.get_leaf_node(3), Ok(Some(3..4)));
-        assert_eq!(tree.get_leaf_node(4), Ok(Some(4..5)));
-        assert_eq!(tree.get_leaf_node(5), Ok(Some(5..6)));
-        assert_eq!(tree.get_leaf_node(6), Ok(Some(6..7)));
+        assert_eq!(tree.get_leaf_node(3), Ok(LeafNode::Exists(3..4)));
+        assert_eq!(tree.get_leaf_node(4), Ok(LeafNode::Exists(4..5)));
+        assert_eq!(tree.get_leaf_node(5), Ok(LeafNode::Exists(5..6)));
+        assert_eq!(tree.get_leaf_node(6), Ok(LeafNode::Exists(6..7)));
+        assert_eq!(tree.get_leaf_node(7), Ok(LeafNode::DoesNotExist));
+
+        let tree = get_tree_a(3);
+
+        assert_eq!(tree.get_leaf_node(3), Ok(LeafNode::Exists(3..4)));
+        assert_eq!(tree.get_leaf_node(4), Ok(LeafNode::Exists(4..5)));
+        assert_eq!(tree.get_leaf_node(5), Ok(LeafNode::Exists(5..6)));
+        assert_eq!(tree.get_leaf_node(6), Ok(LeafNode::Padding));
+        assert_eq!(tree.get_leaf_node(7), Ok(LeafNode::DoesNotExist));
+
+        let tree = get_tree_a(0);
+
+        assert_eq!(tree.get_leaf_node(0), Ok(LeafNode::Padding));
+        assert_eq!(tree.get_leaf_node(1), Ok(LeafNode::DoesNotExist));
+
+        let tree = BTreeSchema::from_lengths(0, vec![3]).into_overlay(0);
+        assert_eq!(tree.get_leaf_node(0), Ok(LeafNode::Exists(0..3)));
+        assert_eq!(tree.get_leaf_node(1), Ok(LeafNode::DoesNotExist));
     }
 
     #[test]