diff --git a/src/iterator.rs b/src/iterator.rs index 3bd0d21..a58638b 100644 --- a/src/iterator.rs +++ b/src/iterator.rs @@ -206,9 +206,7 @@ where let mut current_node_key = NodeKey::new_empty_path(version); let mut current_node = reader.get_node(¤t_node_key)?; - let total_leaves = current_node - .leaf_count() - .ok_or_else(|| format_err!("Leaf counts not available."))?; + let total_leaves = current_node.leaf_count(); if start_idx >= total_leaves { return Ok(Self { reader, @@ -260,7 +258,7 @@ where target_leaf_idx: usize, ) -> Result<(Nibble, &'a Child)> { for (nibble, child) in internal_node.children_sorted() { - let child_leaf_count = child.leaf_count().expect("Leaf count available."); + let child_leaf_count = child.leaf_count(); // n.b. The index is 0-based, so to reach leaf at N, N previous ones need to be skipped. if *leaves_skipped + child_leaf_count <= target_leaf_idx { *leaves_skipped += child_leaf_count; diff --git a/src/node_type.rs b/src/node_type.rs index a98aeb3..88de3b7 100644 --- a/src/node_type.rs +++ b/src/node_type.rs @@ -114,12 +114,7 @@ impl NodeKey { )] pub enum NodeType { Leaf, - /// A internal node that haven't been finished the leaf count migration, i.e. None or not all - /// of the children leaf counts are known. - InternalLegacy, - Internal { - leaf_count: usize, - }, + Internal { leaf_count: usize }, } #[cfg(any(test))] @@ -130,7 +125,6 @@ impl Arbitrary for NodeType { fn arbitrary_with(_args: ()) -> Self::Strategy { prop_oneof![ Just(NodeType::Leaf), - Just(NodeType::InternalLegacy), (2..100usize).prop_map(|leaf_count| NodeType::Internal { leaf_count }) ] .boxed() @@ -157,7 +151,7 @@ pub struct Child { /// from the storage. Used by `[`NodeKey::gen_child_node_key`]. pub version: Version, /// Indicates if the child is a leaf, or if it's an internal node, the total number of leaves - /// under it (though it can be unknown during migration). + /// under it. pub node_type: NodeType, } @@ -174,11 +168,10 @@ impl Child { matches!(self.node_type, NodeType::Leaf) } - pub fn leaf_count(&self) -> Option { + pub fn leaf_count(&self) -> usize { match self.node_type { - NodeType::Leaf => Some(1), - NodeType::InternalLegacy => None, - NodeType::Internal { leaf_count } => Some(leaf_count), + NodeType::Leaf => 1, + NodeType::Internal { leaf_count } => leaf_count, } } } @@ -317,7 +310,7 @@ pub struct InternalNode { /// Up to 16 children. children: Children, /// Total number of leaves under this internal node - leaf_count: Option, + leaf_count: usize, } impl SparseMerkleInternalNode { @@ -434,26 +427,22 @@ impl InternalNode { } } - fn sum_leaf_count(children: &Children) -> Option { + fn sum_leaf_count(children: &Children) -> usize { let mut leaf_count = 0; for child in children.values() { - if let Some(n) = child.leaf_count() { - leaf_count += n; - } else { - return None; - } + let n = child.leaf_count(); + leaf_count += n; } - Some(leaf_count) + leaf_count } - pub fn leaf_count(&self) -> Option { + pub fn leaf_count(&self) -> usize { self.leaf_count } pub fn node_type(&self) -> NodeType { - match self.leaf_count { - Some(leaf_count) => NodeType::Internal { leaf_count }, - None => NodeType::InternalLegacy, + NodeType::Internal { + leaf_count: self.leaf_count, } } @@ -956,10 +945,10 @@ impl Node { } /// Returns leaf count if known - pub(crate) fn leaf_count(&self) -> Option { + pub(crate) fn leaf_count(&self) -> usize { match self { - Node::Null => Some(0), - Node::Leaf(_) => Some(1), + Node::Null => 0, + Node::Leaf(_) => 1, Node::Internal(internal_node) => internal_node.leaf_count, } } diff --git a/src/restore.rs b/src/restore.rs index 67cccec..9b57d67 100644 --- a/src/restore.rs +++ b/src/restore.rs @@ -39,7 +39,7 @@ enum ChildInfo { /// hash of an internal node after we see all the keys that share the same prefix. Internal { hash: Option<[u8; 32]>, - leaf_count: Option, + leaf_count: usize, }, /// This child is a leaf node. @@ -53,9 +53,7 @@ impl ChildInfo { Self::Internal { hash, leaf_count } => Child::new( hash.expect("Must have been initialized."), version, - leaf_count - .map(|n| NodeType::Internal { leaf_count: n }) - .unwrap_or(NodeType::InternalLegacy), + NodeType::Internal { leaf_count }, ), Self::Leaf { node } => Child::new(node.hash::(), version, NodeType::Leaf), } @@ -270,7 +268,7 @@ impl JellyfishMerkleRestore { index, ChildInfo::Internal { hash: None, - leaf_count: None, + leaf_count: 0, }, ); } @@ -394,7 +392,7 @@ impl JellyfishMerkleRestore { child_index, ChildInfo::Internal { hash: None, - leaf_count: None, + leaf_count: 0, }, ); @@ -414,7 +412,7 @@ impl JellyfishMerkleRestore { u8::from(next_nibble) as usize, ChildInfo::Internal { hash: None, - leaf_count: None, + leaf_count: 0, }, ); self.partial_nodes.push(internal_info); @@ -522,8 +520,8 @@ impl JellyfishMerkleRestore { ref mut leaf_count, }) => { assert_eq!(hash.replace(node_hash), None); - assert!(leaf_count.is_none()); - node_leaf_count.map(|n| leaf_count.replace(n)); + assert_eq!(*leaf_count, 0); + *leaf_count = node_leaf_count; } _ => panic!( "Must have at least one child and the rightmost child must not be a leaf." diff --git a/src/tests/helper.rs b/src/tests/helper.rs index f7dc53e..51e0b4c 100644 --- a/src/tests/helper.rs +++ b/src/tests/helper.rs @@ -874,5 +874,5 @@ pub fn test_get_leaf_count(keys: HashSet) { let kvs = keys.into_iter().map(|k| (k, vec![])).collect(); let (db, version) = init_mock_db::(&kvs); let tree = JellyfishMerkleTree::<_, H>::new(&db); - assert_eq!(tree.get_leaf_count(version).unwrap().unwrap(), kvs.len()) + assert_eq!(tree.get_leaf_count(version).unwrap(), kvs.len()) } diff --git a/src/tree.rs b/src/tree.rs index fba6ae4..9d9b339 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -35,7 +35,6 @@ pub type Sha256Jmt<'a, R> = JellyfishMerkleTree<'a, R, sha2::Sha256>; /// and a [`SimpleHasher`] `H`. See [`crate`] for description. pub struct JellyfishMerkleTree<'a, R, H: SimpleHasher> { reader: &'a R, - leaf_count_migration: bool, _phantom_hasher: PhantomData, } @@ -51,15 +50,6 @@ where pub fn new(reader: &'a R) -> Self { Self { reader, - leaf_count_migration: true, - _phantom_hasher: Default::default(), - } - } - - pub fn new_migration(reader: &'a R, leaf_count_migration: bool) -> Self { - Self { - reader, - leaf_count_migration, _phantom_hasher: Default::default(), } } @@ -1390,16 +1380,8 @@ where } // TODO: should this be public? seems coupled to tests? - pub fn get_leaf_count(&self, version: Version) -> Result> { - if self.leaf_count_migration { - self.get_root_node(version).map(|n| n.leaf_count()) - } else { - // When all children of an internal node are leaves, the leaf count is accessible - // even if the migration haven't started. In fact, in such a case, there's no difference - // in the old and new serialization format. Forcing it None here just to make the tests - // straightforward. - Ok(None) - } + pub fn get_leaf_count(&self, version: Version) -> Result { + self.get_root_node(version).map(|n| n.leaf_count()) } }