chore: add get_path_to_descendant to NodeIndex

crates/committer/src/felt.rs

@@ -54,6 +54,7 @@ impl Felt {
     pub(crate) const ONE: Felt = Felt(StarknetTypesFelt::ONE);
     pub(crate) const TWO: Felt = Felt(StarknetTypesFelt::TWO);
     pub(crate) const THREE: Felt = Felt(StarknetTypesFelt::THREE);
+    pub(crate) const MAX: Felt = Felt(StarknetTypesFelt::MAX);
 
     pub fn from_bytes_be_slice(bytes: &[u8]) -> Self {
         Self(StarknetTypesFelt::from_bytes_be_slice(bytes))

crates/committer/src/patricia_merkle_tree.rs

@@ -1,3 +1,4 @@
+pub mod errors;
 pub mod filled_tree;
 pub mod node_data;
 pub mod original_skeleton_tree;

crates/committer/src/patricia_merkle_tree/errors.rs

@@ -0,0 +1,12 @@
+use std::fmt::Debug;
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub(crate) enum TypesError<T: Sized + Debug> {
+    #[error("Failed to convert type {from:?} to {to}. Reason: {reason}.")]
+    ConversionError {
+        from: T,
+        to: &'static str,
+        reason: &'static str,
+    },
+}

crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree.rs

@@ -2,6 +2,7 @@ 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::LeafDataImpl;
 use crate::patricia_merkle_tree::original_skeleton_tree::tree::{
     OriginalSkeletonTreeImpl, OriginalSkeletonTreeResult,
@@ -25,6 +26,23 @@ impl OriginalSkeletonTreeImpl {
         TreeHeight(self.tree_height.0 - 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.
+    /// * Descendants of the given index.
+    /// * Non-empty list.
+    fn get_path_to_lca(&self, root_index: &NodeIndex, leaf_indices: &[NodeIndex]) -> PathToBottom {
+        if leaf_indices.is_empty() {
+            panic!("Unexpected empty array.");
+        }
+        let lca = if leaf_indices.len() == 1 {
+            leaf_indices[0]
+        } else {
+            leaf_indices[0].get_lca(leaf_indices.last().expect("Unexpected empty array"))
+        };
+        root_index.get_path_to_descendant(lca)
+    }
+
     /// Returns whether a root of a subtree has leaves on both sides. Assumes:
     /// * The leaf indices array is sorted.
     /// * All leaves are descendants of the root.

crates/committer/src/patricia_merkle_tree/original_skeleton_tree/compute_updated_skeleton_tree_test.rs

@@ -1,3 +1,5 @@
+use crate::felt::Felt;
+use crate::patricia_merkle_tree::node_data::inner_node::{EdgePath, EdgePathLength, PathToBottom};
 use crate::patricia_merkle_tree::types::NodeIndex;
 use std::collections::HashMap;
 
@@ -61,3 +63,39 @@ fn test_has_leaves_on_both_sides_assertions(
     let root_index = NodeIndex::new(root_index.into());
     skeleton_tree.has_leaves_on_both_sides(&root_index, &leaf_indices);
 }
+
+#[rstest]
+#[case::small_tree_single_leaf(3, 1, vec![U256::from(8_u8)], PathToBottom {path:EdgePath(Felt::ZERO), length:EdgePathLength(3)})]
+#[case::small_tree_few_leaves(
+    3, 1, vec![U256::from(12_u8),U256::from(13_u8),U256::from(14_u8)], PathToBottom {path:EdgePath(Felt::ONE), length:EdgePathLength(1)})]
+#[case::small_tree_few_leaves2(
+    3, 1, vec![U256::from(12_u8),U256::from(13_u8)], PathToBottom {path:EdgePath(Felt::from(2_u8)), length:EdgePathLength(2)})]
+#[case::large_tree_positive_consecutive_indices_of_different_sides(
+    251,
+    1,
+    vec![(U256::from(3u8) << 250) - U256::ONE, U256::from(3u8) << 250],
+    PathToBottom {path:EdgePath(Felt::ZERO), length:EdgePathLength(0)})]
+#[case::large_tree_positive_consecutive_indices(
+    251,
+    3<<126,
+    vec![U256::from(3u8) << 250, (U256::from(3u8) << 250)+ U256::ONE],
+    PathToBottom {path:EdgePath(Felt::ZERO), length:EdgePathLength(123)})]
+fn test_get_path_to_lca(
+    #[case] tree_height: u8,
+    #[case] root_index: u128,
+    #[case] leaf_indices: Vec<U256>,
+    #[case] expected: PathToBottom,
+) {
+    let skeleton_tree = empty_skeleton(tree_height);
+    let root_index = NodeIndex::new(root_index.into());
+    assert_eq!(
+        skeleton_tree.get_path_to_lca(
+            &root_index,
+            &leaf_indices
+                .iter()
+                .map(|index: &ethnum::U256| NodeIndex::new(*index))
+                .collect::<Vec<_>>()[..]
+        ),
+        expected
+    );
+}

crates/committer/src/patricia_merkle_tree/types.rs

@@ -1,7 +1,8 @@
 use crate::block_committer::input::{ContractAddress, StarknetStorageKey};
 use crate::felt::Felt;
+use crate::patricia_merkle_tree::errors::TypesError;
 use crate::patricia_merkle_tree::filled_tree::node::ClassHash;
-use crate::patricia_merkle_tree::node_data::inner_node::PathToBottom;
+use crate::patricia_merkle_tree::node_data::inner_node::{EdgePath, EdgePathLength, PathToBottom};
 
 use ethnum::U256;
 
@@ -84,6 +85,31 @@ impl NodeIndex {
         NodeIndex(lca)
     }
 
+    /// Returns the path from the node to its given descendant.
+    /// Panics if the supposed descendant is not really a descendant.
+    pub(crate) fn get_path_to_descendant(&self, descendant: Self) -> PathToBottom {
+        let descendant_bit_length = descendant.bit_length();
+        let bit_length = self.bit_length();
+        if bit_length > descendant_bit_length {
+            panic!("The descendant is not a really descendant of the node.");
+        };
+
+        let distance = descendant_bit_length - bit_length;
+        let delta = descendant - (*self << distance);
+        if descendant >> distance != *self {
+            panic!("The descendant is not a really descendant of the node.");
+        };
+
+        PathToBottom {
+            path: EdgePath(
+                delta
+                    .try_into()
+                    .expect("Delta of two indices is unexpectedly larger than a Felt."),
+            ),
+            length: EdgePathLength(distance),
+        }
+    }
+
     pub(crate) fn from_starknet_storage_key(
         key: &StarknetStorageKey,
         tree_height: &TreeHeight,
@@ -156,3 +182,19 @@ impl From<NodeIndex> for U256 {
         value.0
     }
 }
+
+impl TryFrom<NodeIndex> for Felt {
+    type Error = TypesError<NodeIndex>;
+
+    fn try_from(value: NodeIndex) -> Result<Self, Self::Error> {
+        if value.0 > U256::from_be_bytes(Self::MAX.to_bytes_be()) {
+            return Err(TypesError::ConversionError {
+                from: value,
+                to: "Felt",
+                reason: "NodeIndex is too large",
+            });
+        }
+        let bytes = value.0.to_be_bytes();
+        Ok(Self::from_bytes_be_slice(&bytes))
+    }
+}

crates/committer/src/patricia_merkle_tree/types_test.rs

@@ -94,3 +94,47 @@ fn test_get_lca_big() {
     let right_descendant = random_extension(right_child);
     assert_eq!(left_descendant.get_lca(&right_descendant), lca);
 }
+
+#[rstest]
+#[case(3, 3, 0, 0)]
+#[case(2, 10, 2, 2)]
+#[should_panic]
+#[case(2, 3, 0, 0)]
+#[should_panic]
+#[case(2, 6, 0, 0)]
+#[should_panic]
+#[case(6, 2, 0, 0)]
+fn test_get_path_to_descendant(
+    #[case] root_index: u8,
+    #[case] descendant: u8,
+    #[case] expected_path: u8,
+    #[case] expected_length: u8,
+) {
+    let root_index = NodeIndex(root_index.into());
+    let descendant = NodeIndex(descendant.into());
+    let path_to_bottom = root_index.get_path_to_descendant(descendant);
+    assert_eq!(path_to_bottom.path, EdgePath(Felt::from(expected_path)));
+    assert_eq!(path_to_bottom.length, EdgePathLength(expected_length));
+}
+
+#[rstest]
+fn test_get_path_to_descendant_big() {
+    let root_index = NodeIndex(U256::from(rand::thread_rng().gen::<u128>()));
+    let max_bits = NodeIndex::BITS - 128;
+    let extension: u128 = rand::thread_rng().gen_range(0..1 << max_bits);
+    let extension_index = NodeIndex(U256::from(extension));
+
+    let descendant = (root_index << extension_index.bit_length()) + extension_index;
+    let path_to_bottom = root_index.get_path_to_descendant(descendant);
+    assert_eq!(path_to_bottom.path, EdgePath(Felt::from(extension)));
+    assert_eq!(
+        path_to_bottom.length,
+        EdgePathLength(extension_index.bit_length())
+    );
+}
+
+#[rstest]
+fn test_nodeindex_to_felt_conversion() {
+    let index = NodeIndex::MAX_INDEX;
+    assert!(Felt::try_from(index).is_err());
+}