From 713237c8e9d03221b9ee05c73360fc145ea93fba Mon Sep 17 00:00:00 2001 From: Erwan Date: Mon, 9 Oct 2023 16:00:00 -0400 Subject: [PATCH 1/2] jmt: remove leaf migration --- src/node_type.rs | 21 +++++---------------- src/restore.rs | 21 +++------------------ src/tests/restore.rs | 30 ++++++++---------------------- src/tree.rs | 11 ++++------- 4 files changed, 20 insertions(+), 63 deletions(-) diff --git a/src/node_type.rs b/src/node_type.rs index c3d304e..a98aeb3 100644 --- a/src/node_type.rs +++ b/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))] @@ -318,8 +318,6 @@ pub struct InternalNode { children: Children, /// Total number of leaves under this internal node leaf_count: Option, - /// serialize leaf counts - leaf_count_migration: bool, } impl SparseMerkleInternalNode { @@ -415,19 +413,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 { // 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,11 +428,10 @@ 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 { diff --git a/src/restore.rs b/src/restore.rs index 5f2ab1e..67cccec 100644 --- a/src/restore.rs +++ b/src/restore.rs @@ -88,11 +88,7 @@ impl InternalInfo { /// Converts `self` to an internal node, assuming all of its children are already known and /// fully initialized. - fn into_internal_node( - mut self, - version: Version, - leaf_count_migration: bool, - ) -> (NodeKey, InternalNode) { + fn into_internal_node(mut self, version: Version) -> (NodeKey, InternalNode) { let mut children = Children::new(); // Calling `into_iter` on an array is equivalent to calling `iter`: @@ -103,10 +99,7 @@ impl InternalInfo { } } - ( - self.node_key, - InternalNode::new_migration(children, leaf_count_migration), - ) + (self.node_key, InternalNode::new(children)) } } @@ -165,9 +158,6 @@ pub struct JellyfishMerkleRestore { /// 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, } @@ -176,7 +166,6 @@ impl JellyfishMerkleRestore { store: Arc, version: Version, expected_root_hash: RootHash, - leaf_count_migration: bool, ) -> Result { let tree_reader = Arc::clone(&store); let (partial_nodes, previous_leaf) = @@ -203,7 +192,6 @@ impl JellyfishMerkleRestore { previous_leaf, num_keys_received: 0, expected_root_hash, - leaf_count_migration, _phantom_hasher: Default::default(), }) } @@ -212,7 +200,6 @@ impl JellyfishMerkleRestore { store: Arc, version: Version, expected_root_hash: RootHash, - leaf_count_migration: bool, ) -> Result { Ok(Self { store, @@ -222,7 +209,6 @@ impl JellyfishMerkleRestore { previous_leaf: None, num_keys_received: 0, expected_root_hash, - leaf_count_migration, _phantom_hasher: Default::default(), }) } @@ -512,8 +498,7 @@ impl JellyfishMerkleRestore { 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::(self.version, self.leaf_count_migration); + let (node_key, internal_node) = last_node.into_internal_node::(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::(); diff --git a/src/tests/restore.rs b/src/tests/restore.rs index 2a716c6..9da711d 100644 --- a/src/tests/restore.rs +++ b/src/tests/restore.rs @@ -25,13 +25,9 @@ fn test_restore_with_interruption( let restore_db = Arc::new(MockTreeStore::default()); { - let mut restore = JellyfishMerkleRestore::::new( - Arc::clone(&restore_db), - version, - expected_root_hash, - true, /* leaf_count_migraion */ - ) - .unwrap(); + let mut restore = + JellyfishMerkleRestore::::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( .filter(|(k, _v)| *k > rightmost_key) .collect(); - let mut restore = JellyfishMerkleRestore::::new( - Arc::clone(&restore_db), - version, - expected_root_hash, - true, /* leaf_count_migration */ - ) - .unwrap(); + let mut restore = + JellyfishMerkleRestore::::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( let expected_root_hash = tree.get_root_hash(source_version).unwrap(); let mut restore = if try_resume { - JellyfishMerkleRestore::::new( - Arc::clone(target_db), - target_version, - expected_root_hash, - true, /* account_count_migration */ - ) - .unwrap() + JellyfishMerkleRestore::::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() }; diff --git a/src/tree.rs b/src/tree.rs index 1ef7436..fba6ae4 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -193,8 +193,7 @@ where ), ); } - let new_internal_node = - InternalNode::new_migration(children, self.leaf_count_migration); + let new_internal_node = InternalNode::new(children); node_key.set_version(version); @@ -300,8 +299,7 @@ where tree_cache.put_node(existing_leaf_node_key, existing_leaf_node.into())?; } - let new_internal_node = - InternalNode::new_migration(children, self.leaf_count_migration); + let new_internal_node = InternalNode::new(children); tree_cache.put_node(node_key.clone(), new_internal_node.clone().into())?; Ok((node_key, new_internal_node.into())) @@ -343,8 +341,7 @@ where ), ); } - let new_internal_node = - InternalNode::new_migration(children, self.leaf_count_migration); + let new_internal_node = InternalNode::new(children); tree_cache.put_node(node_key.clone(), new_internal_node.clone().into())?; Ok((node_key, new_internal_node.into())) @@ -910,7 +907,7 @@ where Child::new(new_leaf_node.hash::(), version, NodeType::Leaf), ); - let internal_node = InternalNode::new_migration(children, self.leaf_count_migration); + let internal_node = InternalNode::new(children); let mut next_internal_node = internal_node.clone(); tree_cache.put_node(node_key.clone(), internal_node.into())?; From 5dec4bfaa5c9c3365c047cee34eee232c2a7ef33 Mon Sep 17 00:00:00 2001 From: Erwan Date: Mon, 9 Oct 2023 16:09:22 -0400 Subject: [PATCH 2/2] jmt: propagate api changes --- src/iterator.rs | 6 ++---- src/node_type.rs | 43 ++++++++++++++++--------------------------- src/restore.rs | 16 +++++++--------- src/tests/helper.rs | 2 +- src/tree.rs | 22 ++-------------------- 5 files changed, 28 insertions(+), 61 deletions(-) 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()) } }