From 9fa214040d04c02fb3ec8b5721f986eff546a95f Mon Sep 17 00:00:00 2001 From: Tzahi Taub Date: Wed, 8 May 2024 14:30:01 +0300 Subject: [PATCH] chore: tight range checks in NodeIndex --- .../patricia_merkle_tree/filled_tree/tree.rs | 2 +- .../compute_updated_skeleton_tree.rs | 3 +- .../compute_updated_skeleton_tree_test.rs | 10 +-- .../original_skeleton_tree/create_tree.rs | 11 ++- .../src/patricia_merkle_tree/types.rs | 80 +++++++++++++------ .../updated_skeleton_tree/tree.rs | 13 +-- .../updated_skeleton_tree/tree_test.rs | 2 +- 7 files changed, 74 insertions(+), 47 deletions(-) diff --git a/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs b/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs index 23716c1d..e856fddb 100644 --- a/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs +++ b/crates/committer/src/patricia_merkle_tree/filled_tree/tree.rs @@ -48,7 +48,7 @@ impl FilledTree for FilledTreeImpl { } fn get_root_hash(&self) -> Result { - match self.tree_map.get(&NodeIndex::root_index()) { + match self.tree_map.get(&NodeIndex::ROOT) { Some(root_node) => Ok(root_node.hash), None => Err(FilledTreeError::MissingRoot), } 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 c6c3685f..f2bb0600 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 @@ -53,6 +53,7 @@ impl OriginalSkeletonTreeImpl { assert_child(*last_leaf); let child_direction_mask = U256::ONE << (root_height.0 - 1); - (first_leaf.0 & child_direction_mask) != (last_leaf.0 & child_direction_mask) + (U256::from(first_leaf) & child_direction_mask) + != (U256::from(*last_leaf) & child_direction_mask) } } diff --git a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree_test.rs b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree_test.rs index 36789772..b4ab93cc 100644 --- a/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree_test.rs +++ b/crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree_test.rs @@ -22,17 +22,17 @@ fn empty_skeleton(height: u8) -> OriginalSkeletonTreeImpl { #[case::large_tree_farthest_leaves( 251, 1, - vec![NodeIndex(U256::ONE << 251),NodeIndex((U256::ONE << 252) - U256::ONE)], + vec![NodeIndex::ROOT << 251, NodeIndex::MAX_INDEX], true)] #[case::large_tree_positive_consecutive_indices_of_different_sides( 251, 1, - vec![NodeIndex((U256::from(3u8) << 250) - U256::ONE), NodeIndex(U256::from(3u8) << 250)], + vec![NodeIndex::new((U256::from(3u8) << 250) - U256::ONE), NodeIndex::new(U256::from(3u8) << 250)], true)] #[case::large_tree_negative_one_shift_of_positive_case( 251, 1, - vec![NodeIndex(U256::from(3u8) << 250), NodeIndex((U256::from(3u8) << 250)+ U256::ONE)], + vec![NodeIndex::new(U256::from(3u8) << 250), NodeIndex::new((U256::from(3u8) << 250)+ U256::ONE)], false)] fn test_has_leaves_on_both_sides( #[case] tree_height: u8, @@ -41,7 +41,7 @@ fn test_has_leaves_on_both_sides( #[case] expected: bool, ) { let skeleton_tree = empty_skeleton(tree_height); - let root_index = NodeIndex(root_index.into()); + let root_index = NodeIndex::new(root_index.into()); assert_eq!( skeleton_tree.has_leaves_on_both_sides(&root_index, &leaf_indices), expected @@ -58,6 +58,6 @@ fn test_has_leaves_on_both_sides_assertions( #[case] leaf_indices: Vec, ) { let skeleton_tree = empty_skeleton(tree_height); - let root_index = NodeIndex(root_index.into()); + let root_index = NodeIndex::new(root_index.into()); skeleton_tree.has_leaves_on_both_sides(&root_index, &leaf_indices); } 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 a81147d2..a751d3f5 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 @@ -19,7 +19,6 @@ use crate::storage::storage_trait::Storage; use crate::storage::storage_trait::StorageKey; use crate::storage::storage_trait::StoragePrefix; use bisection::{bisect_left, bisect_right}; -use ethnum::U256; use std::collections::HashMap; #[cfg(test)] #[path = "create_tree_test.rs"] @@ -42,7 +41,7 @@ impl<'a> SubTree<'a> { ) -> (&'a [NodeIndex], &'a [NodeIndex]) { let height = self.get_height(total_tree_height); let leftmost_index_in_right_subtree = - ((self.root_index << 1) + NodeIndex(U256::ONE)) << (height.0 - 1); + ((self.root_index << 1) + NodeIndex::ROOT) << (height.0 - 1); let mid = bisect_left(self.sorted_leaf_indices, &leftmost_index_in_right_subtree); ( &self.sorted_leaf_indices[..mid], @@ -65,7 +64,7 @@ impl<'a> SubTree<'a> { self.get_height(total_tree_height) - TreeHeight(path_to_bottom.length.0); let leftmost_in_subtree = bottom_index << bottom_height.0; let rightmost_in_subtree = - leftmost_in_subtree + (NodeIndex(U256::ONE) << bottom_height.0) - NodeIndex(U256::ONE); + leftmost_in_subtree + (NodeIndex::ROOT << bottom_height.0) - NodeIndex::ROOT; let bottom_leaves = &self.sorted_leaf_indices[bisect_left(self.sorted_leaf_indices, &leftmost_in_subtree) ..bisect_right(self.sorted_leaf_indices, &rightmost_in_subtree)]; @@ -84,7 +83,7 @@ impl<'a> SubTree<'a> { total_tree_height: &TreeHeight, ) -> (Self, Self) { let (left_leaves, right_leaves) = self.split_leaves(total_tree_height); - let left_root_index = self.root_index * 2; + let left_root_index = self.root_index * 2.into(); ( SubTree { sorted_leaf_indices: left_leaves, @@ -93,7 +92,7 @@ impl<'a> SubTree<'a> { }, SubTree { sorted_leaf_indices: right_leaves, - root_index: left_root_index + NodeIndex(U256::ONE), + root_index: left_root_index + NodeIndex::ROOT, root_hash: right_hash, }, ) @@ -216,7 +215,7 @@ impl OriginalSkeletonTreeImpl { ) -> OriginalSkeletonTreeResult { let main_subtree = SubTree { sorted_leaf_indices, - root_index: NodeIndex::root_index(), + root_index: NodeIndex::ROOT, root_hash, }; let mut skeleton_tree = Self { diff --git a/crates/committer/src/patricia_merkle_tree/types.rs b/crates/committer/src/patricia_merkle_tree/types.rs index bc1d4242..bc3e88f4 100644 --- a/crates/committer/src/patricia_merkle_tree/types.rs +++ b/crates/committer/src/patricia_merkle_tree/types.rs @@ -9,27 +9,36 @@ use ethnum::U256; #[path = "types_test.rs"] pub mod types_test; -#[allow(dead_code)] +#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Sub)] +pub struct TreeHeight(pub u8); + +impl TreeHeight { + pub const MAX_HEIGHT: u8 = 251; +} #[derive( - Clone, - Copy, - Debug, - PartialEq, - Eq, - Hash, - derive_more::Add, - derive_more::Mul, - derive_more::Sub, - PartialOrd, - Ord, + Clone, Copy, Debug, PartialEq, Eq, Hash, derive_more::BitAnd, derive_more::Sub, PartialOrd, Ord, )] -pub(crate) struct NodeIndex(pub U256); +pub(crate) struct NodeIndex(U256); -#[allow(dead_code)] // Wraps a U256. Maximal possible value is the largest index in a tree of height 251 (2 ^ 252 - 1). impl NodeIndex { - pub(crate) fn root_index() -> NodeIndex { - NodeIndex(U256::ONE) + pub const BITS: u8 = TreeHeight::MAX_HEIGHT + 1; + + /// [NodeIndex] constant that represents the root index. + pub const ROOT: Self = Self(U256::ONE); + + /// [NodeIndex] constant that represents the largest index in a tree. + #[allow(clippy::as_conversions)] + pub const MAX_INDEX: Self = Self(U256::from_words( + u128::MAX >> (U256::BITS - Self::BITS as u32), + u128::MAX, + )); + + pub(crate) fn new(index: U256) -> Self { + if index > Self::MAX_INDEX.0 { + panic!("Index {index} is too large."); + } + Self(index) } // TODO(Amos, 1/5/2024): Move to EdgePath. @@ -41,10 +50,15 @@ impl NodeIndex { (index << length.0) + Self::from_felt_value(&path.0) } - pub(crate) fn bit_length(&self) -> u8 { - (U256::BITS - self.0.leading_zeros()) + /// Returns the number of leading zeroes when represented with Self::BITS bits. + pub(crate) fn leading_zeros(&self) -> u8 { + (self.0.leading_zeros() - (U256::BITS - u32::from(Self::BITS))) .try_into() - .expect("Failed to convert to u8.") + .expect("Leading zeroes are unexpectedly larger than a u8.") + } + + pub(crate) fn bit_length(&self) -> u8 { + Self::BITS - self.leading_zeros() } pub(crate) fn from_starknet_storage_key( @@ -74,12 +88,28 @@ impl NodeIndex { } } +impl std::ops::Add for NodeIndex { + type Output = Self; + + fn add(self, rhs: Self) -> Self { + Self::new(self.0 + rhs.0) + } +} + +impl std::ops::Mul for NodeIndex { + type Output = Self; + + fn mul(self, rhs: Self) -> Self { + Self::new(self.0 * rhs.0) + } +} + impl std::ops::Shl for NodeIndex { type Output = Self; /// Returns the index of the left descendant (child for rhs=1) of the node. fn shl(self, rhs: u8) -> Self::Output { - NodeIndex(self.0 << rhs) + Self::new(self.0 << rhs) } } @@ -88,7 +118,7 @@ impl std::ops::Shr for NodeIndex { /// Returns the index of the ancestor (parent for rhs=1) of the node. fn shr(self, rhs: u8) -> Self::Output { - NodeIndex(self.0 >> rhs) + Self::new(self.0 >> rhs) } } @@ -98,6 +128,8 @@ impl From for NodeIndex { } } -#[allow(dead_code)] -#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Sub)] -pub struct TreeHeight(pub u8); +impl From for U256 { + fn from(value: NodeIndex) -> Self { + value.0 + } +} diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs index 7f4d2453..5557ac8d 100644 --- a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs +++ b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree.rs @@ -1,5 +1,3 @@ -use ethnum::U256; - use std::collections::HashMap; use std::sync::{Arc, Mutex}; @@ -119,8 +117,8 @@ impl let node = self.get_node(index)?; match node { UpdatedSkeletonNode::Binary => { - let left_index = index * 2; - let right_index = left_index + NodeIndex(U256::ONE); + let left_index = index * 2.into(); + let right_index = left_index + NodeIndex::ROOT; let (left_hash, right_hash) = tokio::join!( self.compute_filled_tree_rec::(left_index, Arc::clone(&output_map)), @@ -171,11 +169,8 @@ impl Up // 2. Fill in the hash values. let filled_tree_map = Arc::new(self.initialize_with_placeholders()); - self.compute_filled_tree_rec::( - NodeIndex::root_index(), - Arc::clone(&filled_tree_map), - ) - .await?; + self.compute_filled_tree_rec::(NodeIndex::ROOT, Arc::clone(&filled_tree_map)) + .await?; // Create and return a new FilledTreeImpl from the hashmap. Ok(FilledTreeImpl::new(Self::remove_arc_mutex_and_option( diff --git a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs index 3f808f8b..81cc5446 100644 --- a/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs +++ b/crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/tree_test.rs @@ -21,7 +21,7 @@ use crate::patricia_merkle_tree::updated_skeleton_tree::tree::{ async fn test_filled_tree_sanity() { let mut skeleton_tree: HashMap> = HashMap::new(); skeleton_tree.insert( - NodeIndex::root_index(), + NodeIndex::ROOT, UpdatedSkeletonNode::Leaf(LeafDataImpl::CompiledClassHash(ClassHash(Felt::ONE))), ); let updated_skeleton_tree = UpdatedSkeletonTreeImpl { skeleton_tree };