chore(skeleton): common split_leaves function

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;

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<L: LeafData + std::clone::Clone> OriginalSkeletonTreeImpl<L> {
         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<L: LeafData + std::clone::Clone> OriginalSkeletonTreeImpl<L> {
         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())
     }
 }

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 {

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..]]
+}

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<U256> {
+    let size_u256: U256 = size.try_into().unwrap();
+    assert!(jump > 0 && start + jump * size_u256 < U256::MAX);
+    let mut ret: Vec<U256> = 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::<Vec<_>>();
+    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(),
+    )
+}

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(