diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree.rs index a805db97..b88642d9 100644 --- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree.rs +++ b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree.rs @@ -4,3 +4,4 @@ pub mod errors; pub mod node; pub mod skeleton_forest; pub mod tree; +pub mod utils; diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree.rs index ba3de154..a88eb069 100644 --- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree.rs +++ b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree.rs @@ -1,13 +1,12 @@ use std::collections::HashMap; -use ethnum::U256; - use crate::patricia_merkle_tree::node_data::inner_node::PathToBottom; use crate::patricia_merkle_tree::node_data::leaf::LeafData; use crate::patricia_merkle_tree::original_skeleton_tree::tree::{ OriginalSkeletonTreeImpl, OriginalSkeletonTreeResult, }; -use crate::patricia_merkle_tree::types::{NodeIndex, TreeHeight}; +use crate::patricia_merkle_tree::original_skeleton_tree::utils::split_leaves; +use crate::patricia_merkle_tree::types::NodeIndex; use crate::patricia_merkle_tree::updated_skeleton_tree::tree::UpdatedSkeletonTreeImpl; #[cfg(test)] @@ -22,10 +21,6 @@ impl OriginalSkeletonTreeImpl { todo!() } - fn get_node_height(&self, index: &NodeIndex) -> TreeHeight { - TreeHeight::new(u8::from(self.tree_height) - index.bit_length() + 1) - } - #[allow(dead_code)] /// Returns the path from the given root_index to the LCA of the leaves. Assumes the leaves are: /// * Sorted. @@ -51,27 +46,8 @@ impl OriginalSkeletonTreeImpl { if leaf_indices.is_empty() { return false; } - - let root_height = self.get_node_height(root_index); - let assert_child = |leaf_index: NodeIndex| { - if (leaf_index >> root_height.into()) != *root_index { - panic!("Leaf is not a descendant of the root."); - } - }; - - let first_leaf = leaf_indices[0]; - assert_child(first_leaf); - if leaf_indices.len() == 1 { - return false; - } - - let last_leaf = leaf_indices - .last() - .expect("leaf_indices unexpectedly empty."); - assert_child(*last_leaf); - - let child_direction_mask = U256::ONE << (u8::from(root_height) - 1); - (U256::from(first_leaf) & child_direction_mask) - != (U256::from(*last_leaf) & child_direction_mask) + split_leaves(&self.tree_height, root_index, leaf_indices) + .iter() + .all(|leaves_in_side| !leaves_in_side.is_empty()) } } diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs index a7ca73ff..c65cb84a 100644 --- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs +++ b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs @@ -9,6 +9,7 @@ use crate::patricia_merkle_tree::node_data::leaf::LeafData; use crate::patricia_merkle_tree::node_data::leaf::LeafDataImpl; use crate::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeImpl; use crate::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeResult; +use crate::patricia_merkle_tree::original_skeleton_tree::utils::split_leaves; use crate::patricia_merkle_tree::types::TreeHeight; use crate::patricia_merkle_tree::{ original_skeleton_tree::node::OriginalSkeletonNode, types::NodeIndex, @@ -36,17 +37,11 @@ impl<'a> SubTree<'a> { TreeHeight::new(u8::from(*total_tree_height) - (self.root_index.bit_length() - 1)) } - pub(crate) fn split_leaves( - &self, - total_tree_height: &TreeHeight, - ) -> (&'a [NodeIndex], &'a [NodeIndex]) { - let height = self.get_height(total_tree_height); - let leftmost_index_in_right_subtree = - ((self.root_index << 1) + NodeIndex::ROOT) << (u8::from(height) - 1); - let mid = bisect_left(self.sorted_leaf_indices, &leftmost_index_in_right_subtree); - ( - &self.sorted_leaf_indices[..mid], - &self.sorted_leaf_indices[mid..], + pub(crate) fn split_leaves(&self, total_tree_height: &TreeHeight) -> [&'a [NodeIndex]; 2] { + split_leaves( + total_tree_height, + &self.root_index, + self.sorted_leaf_indices, ) } @@ -83,7 +78,7 @@ impl<'a> SubTree<'a> { right_hash: HashOutput, total_tree_height: &TreeHeight, ) -> (Self, Self) { - let (left_leaves, right_leaves) = self.split_leaves(total_tree_height); + let [left_leaves, right_leaves] = self.split_leaves(total_tree_height); let left_root_index = self.root_index * 2.into(); ( SubTree { diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/utils.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/utils.rs new file mode 100644 index 00000000..d4705c68 --- /dev/null +++ b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/utils.rs @@ -0,0 +1,52 @@ +use bisection::bisect_left; + +use crate::patricia_merkle_tree::types::{NodeIndex, TreeHeight}; + +#[cfg(test)] +#[path = "utils_test.rs"] +pub mod utils_test; + +/// Returns the height of the node with the given index. +pub(crate) fn get_node_height(tree_height: &TreeHeight, index: &NodeIndex) -> TreeHeight { + TreeHeight::new(u8::from(*tree_height) - index.bit_length() + 1) +} + +/// Splits leaf_indices into two arrays according to the given root: the left child leaves and +/// the right child leaves. Assumes: +/// * The leaf indices array is sorted. +/// * All leaves are descendants of the root. +pub(crate) fn split_leaves<'a>( + tree_height: &TreeHeight, + root_index: &NodeIndex, + leaf_indices: &'a [NodeIndex], +) -> [&'a [NodeIndex]; 2] { + if leaf_indices.is_empty() { + return [&[]; 2]; + } + + let root_height = get_node_height(tree_height, root_index); + let assert_descendant = |leaf_index: &NodeIndex| { + if (*leaf_index >> u8::from(root_height)) != *root_index { + panic!( + "Leaf {leaf_index:?} is not a descendant of the root {root_index:?} \ + (root height={root_height:?})." + ); + } + }; + + let first_leaf = leaf_indices[0]; + assert_descendant(&first_leaf); + + if leaf_indices.len() > 1 { + assert_descendant( + leaf_indices + .last() + .expect("leaf_indices unexpectedly empty."), + ); + } + + let right_child_index = (*root_index << 1) + 1.into(); + let leftmost_index_in_right_subtree = right_child_index << (u8::from(root_height) - 1); + let leaves_split = bisect_left(leaf_indices, &leftmost_index_in_right_subtree); + [&leaf_indices[..leaves_split], &leaf_indices[leaves_split..]] +} diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/utils_test.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/utils_test.rs new file mode 100644 index 00000000..c71c5499 --- /dev/null +++ b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/utils_test.rs @@ -0,0 +1,72 @@ +use super::split_leaves; +use crate::patricia_merkle_tree::test_utils::get_random_u256; +use crate::patricia_merkle_tree::types::{NodeIndex, TreeHeight}; +use ethnum::{uint, U256}; +use rstest::rstest; + +/// Creates an array of increasing random U256 numbers, with jumps of up to 'jump' between two +/// consecutive numbers. +fn create_increasing_random_array(size: usize, start: U256, jump: U256) -> Vec { + let size_u256: U256 = size.try_into().unwrap(); + assert!(jump > 0 && start + jump * size_u256 < U256::MAX); + let mut ret: Vec = Vec::with_capacity(size); + let mut low = start; + for i in 0..size { + ret.push(get_random_u256(low, low + jump)); + low = ret[i] + 1; + } + ret +} + +#[rstest] +#[case::small_tree_one_side( + 3_u8, U256::ONE, &[uint!("8"), uint!("10"), uint!("11")], + &[uint!("8"), uint!("10"), uint!("11")], &[] +)] +#[case::small_tree_two_sides( + 3_u8, U256::ONE, &[uint!("8"), uint!("10"), uint!("14")], + &[uint!("8"), uint!("10")], &[uint!("14")] +)] +#[should_panic] +#[case::small_tree_wrong_height( + 3_u8, U256::ONE, &[uint!("8"), uint!("10"), uint!("16")], &[], &[] +)] +#[should_panic] +#[case::small_tree_not_descendant( + 3_u8, uint!("2"), &[uint!("8"), uint!("10"), uint!("14")], &[], &[] +)] +fn test_split_leaves( + #[case] tree_height: u8, + #[case] root_index: U256, + #[case] leaf_indices: &[U256], + #[case] expected_left: &[U256], + #[case] expected_right: &[U256], +) { + let tree_height = TreeHeight::new(tree_height); + let root_index = NodeIndex::new(root_index); + let to_node_index = |arr: &[U256]| arr.iter().map(|i| NodeIndex::new(*i)).collect::>(); + let leaf_indices = to_node_index(leaf_indices); + let expected = [to_node_index(expected_left), to_node_index(expected_right)]; + assert_eq!( + split_leaves(&tree_height, &root_index, &leaf_indices), + expected + ); +} + +#[rstest] +fn test_split_leaves_big_tree() { + let left_leaf_indices = + create_increasing_random_array(100, NodeIndex::FIRST_LEAF.into(), U256::ONE << 200); + let right_leaf_indices = create_increasing_random_array( + 100, + (U256::from(NodeIndex::FIRST_LEAF) + U256::from(NodeIndex::MAX_INDEX)) / 2 + 1, + U256::ONE << 100, + ); + test_split_leaves( + TreeHeight::MAX.into(), + NodeIndex::ROOT.into(), + &[&left_leaf_indices[..], &right_leaf_indices[..]].concat(), + left_leaf_indices.as_slice(), + right_leaf_indices.as_slice(), + ) +} diff --git a/crates/committer/src/patricia_merkle_tree/types.rs b/crates/committer/src/patricia_merkle_tree/types.rs index f7ba06fa..39db4e7c 100644 --- a/crates/committer/src/patricia_merkle_tree/types.rs +++ b/crates/committer/src/patricia_merkle_tree/types.rs @@ -14,10 +14,10 @@ pub mod types_test; pub struct TreeHeight(u8); impl TreeHeight { - pub const MAX_HEIGHT: u8 = 251; + pub const MAX: TreeHeight = TreeHeight(251); pub fn new(height: u8) -> Self { - if height > Self::MAX_HEIGHT { + if height > Self::MAX.0 { panic!("Height {height} is too large."); } Self(height) @@ -37,11 +37,15 @@ pub(crate) struct NodeIndex(U256); // Wraps a U256. Maximal possible value is the largest index in a tree of height 251 (2 ^ 252 - 1). impl NodeIndex { - pub const BITS: u8 = TreeHeight::MAX_HEIGHT + 1; + pub const BITS: u8 = TreeHeight::MAX.0 + 1; /// [NodeIndex] constant that represents the root index. pub const ROOT: Self = Self(U256::ONE); + #[allow(dead_code)] + /// [NodeIndex] constant that represents the first leaf index. + pub const FIRST_LEAF: Self = Self(U256::from_words(1_u128 << (Self::BITS - 1 - 128), 0)); + /// [NodeIndex] constant that represents the largest index in a tree. #[allow(clippy::as_conversions)] pub const MAX_INDEX: Self = Self(U256::from_words(