jmt: remove leaf migration

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

@@ -14,7 +14,7 @@ use crate::SimpleHasher;
 use alloc::format;
 use alloc::vec::Vec;
 use alloc::{boxed::Box, vec};
-use anyhow::{ensure, Context, Result};
+use anyhow::Context;
 use borsh::{BorshDeserialize, BorshSerialize};
 use num_derive::{FromPrimitive, ToPrimitive};
 #[cfg(any(test))]
@@ -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,9 +310,7 @@ pub struct InternalNode {
     /// Up to 16 children.
     children: Children,
     /// Total number of leaves under this internal node
-    leaf_count: Option<usize>,
-    /// serialize leaf counts
-    leaf_count_migration: bool,
+    leaf_count: usize,
 }
 
 impl SparseMerkleInternalNode {
@@ -415,19 +406,11 @@ fn has_child(
 impl InternalNode {
     /// Creates a new Internal node.
     pub fn new(children: Children) -> Self {
-        Self::new_migration(children, true /* leaf_count_migration */)
-    }
-
-    pub fn new_migration(children: Children, leaf_count_migration: bool) -> Self {
-        Self::new_impl(children, leaf_count_migration).expect("Input children are logical.")
-    }
-
-    pub fn new_impl(children: Children, leaf_count_migration: bool) -> Result<Self> {
         // Assert the internal node must have >= 1 children. If it only has one child, it cannot be
         // a leaf node. Otherwise, the leaf node should be a child of this internal node's parent.
-        ensure!(!children.is_empty(), "Children must not be empty");
+        assert!(!children.is_empty(), "Children must not be empty");
         if children.num_children() == 1 {
-            ensure!(
+            assert!(
                 !children
                     .values()
                     .next()
@@ -438,33 +421,28 @@ impl InternalNode {
         }
 
         let leaf_count = Self::sum_leaf_count(&children);
-        Ok(Self {
+        Self {
             children,
             leaf_count,
-            leaf_count_migration,
-        })
+        }
     }
 
-    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,
         }
     }
 
@@ -967,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),
         }
@@ -88,11 +86,7 @@ impl InternalInfo {
 
     /// Converts `self` to an internal node, assuming all of its children are already known and
     /// fully initialized.
-    fn into_internal_node<H: SimpleHasher>(
-        mut self,
-        version: Version,
-        leaf_count_migration: bool,
-    ) -> (NodeKey, InternalNode) {
+    fn into_internal_node<H: SimpleHasher>(mut self, version: Version) -> (NodeKey, InternalNode) {
         let mut children = Children::new();
 
         // Calling `into_iter` on an array is equivalent to calling `iter`:
@@ -103,10 +97,7 @@ impl InternalInfo {
             }
         }
 
-        (
-            self.node_key,
-            InternalNode::new_migration(children, leaf_count_migration),
-        )
+        (self.node_key, InternalNode::new(children))
     }
 }
 
@@ -165,9 +156,6 @@ pub struct JellyfishMerkleRestore<H: SimpleHasher> {
     /// When the restoration process finishes, we expect the tree to have this root hash.
     expected_root_hash: RootHash,
 
-    /// Whether to use the new internal node format where leaf counts are written.
-    leaf_count_migration: bool,
-
     _phantom_hasher: PhantomData<H>,
 }
 
@@ -176,7 +164,6 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
         store: Arc<D>,
         version: Version,
         expected_root_hash: RootHash,
-        leaf_count_migration: bool,
     ) -> Result<Self> {
         let tree_reader = Arc::clone(&store);
         let (partial_nodes, previous_leaf) =
@@ -203,7 +190,6 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
             previous_leaf,
             num_keys_received: 0,
             expected_root_hash,
-            leaf_count_migration,
             _phantom_hasher: Default::default(),
         })
     }
@@ -212,7 +198,6 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
         store: Arc<D>,
         version: Version,
         expected_root_hash: RootHash,
-        leaf_count_migration: bool,
     ) -> Result<Self> {
         Ok(Self {
             store,
@@ -222,7 +207,6 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
             previous_leaf: None,
             num_keys_received: 0,
             expected_root_hash,
-            leaf_count_migration,
             _phantom_hasher: Default::default(),
         })
     }
@@ -284,7 +268,7 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
                     index,
                     ChildInfo::Internal {
                         hash: None,
-                        leaf_count: None,
+                        leaf_count: 0,
                     },
                 );
             }
@@ -408,7 +392,7 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
             child_index,
             ChildInfo::Internal {
                 hash: None,
-                leaf_count: None,
+                leaf_count: 0,
             },
         );
 
@@ -428,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);
@@ -512,8 +496,7 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
     fn freeze_internal_nodes(&mut self, num_remaining_nodes: usize) {
         while self.partial_nodes.len() > num_remaining_nodes {
             let last_node = self.partial_nodes.pop().expect("This node must exist.");
-            let (node_key, internal_node) =
-                last_node.into_internal_node::<H>(self.version, self.leaf_count_migration);
+            let (node_key, internal_node) = last_node.into_internal_node::<H>(self.version);
             // Keep the hash of this node before moving it into `frozen_nodes`, so we can update
             // its parent later.
             let node_hash = internal_node.hash::<H>();
@@ -537,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/tests/restore.rs

@@ -25,13 +25,9 @@ fn test_restore_with_interruption<H: SimpleHasher>(
 
     let restore_db = Arc::new(MockTreeStore::default());
     {
-        let mut restore = JellyfishMerkleRestore::<H>::new(
-            Arc::clone(&restore_db),
-            version,
-            expected_root_hash,
-            true, /* leaf_count_migraion */
-        )
-        .unwrap();
+        let mut restore =
+            JellyfishMerkleRestore::<H>::new(Arc::clone(&restore_db), version, expected_root_hash)
+                .unwrap();
         let proof = tree
             .get_range_proof(batch1.last().map(|(key, _value)| *key).unwrap(), version)
             .unwrap();
@@ -55,13 +51,9 @@ fn test_restore_with_interruption<H: SimpleHasher>(
             .filter(|(k, _v)| *k > rightmost_key)
             .collect();
 
-        let mut restore = JellyfishMerkleRestore::<H>::new(
-            Arc::clone(&restore_db),
-            version,
-            expected_root_hash,
-            true, /* leaf_count_migration */
-        )
-        .unwrap();
+        let mut restore =
+            JellyfishMerkleRestore::<H>::new(Arc::clone(&restore_db), version, expected_root_hash)
+                .unwrap();
         let proof = tree
             .get_range_proof(
                 remaining_accounts.last().map(|(key, _value)| *key).unwrap(),
@@ -182,19 +174,13 @@ fn restore_without_interruption<H: SimpleHasher>(
     let expected_root_hash = tree.get_root_hash(source_version).unwrap();
 
     let mut restore = if try_resume {
-        JellyfishMerkleRestore::<H>::new(
-            Arc::clone(target_db),
-            target_version,
-            expected_root_hash,
-            true, /* account_count_migration */
-        )
-        .unwrap()
+        JellyfishMerkleRestore::<H>::new(Arc::clone(target_db), target_version, expected_root_hash)
+            .unwrap()
     } else {
         JellyfishMerkleRestore::new_overwrite(
             Arc::clone(target_db),
             target_version,
             expected_root_hash,
-            true, /* account_count_migration */
         )
         .unwrap()
     };