Skip to content

Commit

Permalink
jmt: remove leaf migration (#106)
Browse files Browse the repository at this point in the history 
* jmt: remove leaf migration

* jmt: propagate api changes

---------

Co-authored-by: Preston Evans <32944016+preston-evans98@users.noreply.github.com>
  • Loading branch information
@erwanor @preston-evans98
erwanor and preston-evans98 authored Nov 22, 2023
1 parent d8c854e commit 041ad5c
Show file tree
Hide file tree
Showing 6 changed files with 48 additions and 124 deletions.
6 changes: 2 additions & 4 deletions src/iterator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
Expand Down
64 changes: 21 additions & 43 deletions src/node_type.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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))]
Expand Down Expand Up @@ -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))]
Expand All @@ -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()
Expand All @@ -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,
}

Expand All @@ -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,
}
}
}
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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()
Expand All @@ -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,
}
}

Expand Down Expand Up @@ -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,
}
}
Expand Down
37 changes: 10 additions & 27 deletions src/restore.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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),
}
Expand Down Expand Up @@ -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`:
Expand All @@ -103,10 +97,7 @@ impl InternalInfo {
}
}

(
self.node_key,
InternalNode::new_migration(children, leaf_count_migration),
)
(self.node_key, InternalNode::new(children))
}
}

Expand Down Expand Up @@ -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>,
}

Expand All @@ -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) =
Expand All @@ -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(),
})
}
Expand All @@ -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,
Expand All @@ -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(),
})
}
Expand Down Expand Up @@ -284,7 +268,7 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
index,
ChildInfo::Internal {
hash: None,
leaf_count: None,
leaf_count: 0,
},
);
}
Expand Down Expand Up @@ -408,7 +392,7 @@ impl<H: SimpleHasher> JellyfishMerkleRestore<H> {
child_index,
ChildInfo::Internal {
hash: None,
leaf_count: None,
leaf_count: 0,
},
);

Expand All @@ -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);
Expand Down Expand Up @@ -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>();
Expand All @@ -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."
Expand Down
2 changes: 1 addition & 1 deletion src/tests/helper.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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())
}
30 changes: 8 additions & 22 deletions src/tests/restore.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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(),
Expand Down Expand Up @@ -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()
};
Expand Down
Loading

0 comments on commit 041ad5c

Please sign in to comment.