jmt: propagate api changes

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(&current_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;

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<usize> {
+    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<usize>,
+    leaf_count: usize,
 }
 
 impl SparseMerkleInternalNode {
@@ -434,26 +427,22 @@ impl InternalNode {
         }
     }
 
-    fn sum_leaf_count(children: &Children) -> Option<usize> {
+    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<usize> {
+    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<usize> {
+    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,
         }
     }

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<usize>,
+        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::<H>(), version, NodeType::Leaf),
         }
@@ -270,7 +268,7 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
                     index,
                     ChildInfo::Internal {
                         hash: None,
-                        leaf_count: None,
+                        leaf_count: 0,
                     },
                 );
             }
@@ -394,7 +392,7 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
             child_index,
             ChildInfo::Internal {
                 hash: None,
-                leaf_count: None,
+                leaf_count: 0,
             },
         );
 
@@ -414,7 +412,7 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
                 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<H: SimpleHasher> JellyfishMerkleRestore<H> {
                         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."

src/tests/helper.rs

@@ -874,5 +874,5 @@ pub fn test_get_leaf_count<H: SimpleHasher>(keys: HashSet<KeyHash>) {
     let kvs = keys.into_iter().map(|k| (k, vec![])).collect();
     let (db, version) = init_mock_db::<H>(&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())
 }

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<H>,
 }
 
@@ -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<Option<usize>> {
-        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<usize> {
+        self.get_root_node(version).map(|n| n.leaf_count())
     }
 }