From b1e87fb47789b17f4b2bbe228687d75faf6b7daa Mon Sep 17 00:00:00 2001 From: Tzahi Taub Date: Wed, 8 May 2024 17:10:48 +0300 Subject: [PATCH] chore: get_lca for NodeIndex --- Cargo.lock | 48 +++++++++++++++++++ Cargo.toml | 1 + crates/committer/Cargo.toml | 1 + .../src/patricia_merkle_tree/types.rs | 23 +++++++++ .../src/patricia_merkle_tree/types_test.rs | 40 ++++++++++++++++ 5 files changed, 113 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 6660fcc7..8f4c817a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -204,6 +204,7 @@ dependencies = [ "ethnum", "hex", "pretty_assertions", + "rand", "rstest", "serde", "serde_json", @@ -395,6 +396,17 @@ dependencies = [ "version_check", ] +[[package]] +name = "getrandom" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94b22e06ecb0110981051723910cbf0b5f5e09a2062dd7663334ee79a9d1286c" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "gimli" version = "0.28.1" @@ -590,6 +602,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "ppv-lite86" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" + [[package]] name = "pretty_assertions" version = "1.4.0" @@ -624,6 +642,36 @@ version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + [[package]] name = "redox_syscall" version = "0.4.1" diff --git a/Cargo.toml b/Cargo.toml index 8fdae4d2..d7725657 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ derive_more = "0.99.17" ethnum = "1.5.0" hex = "0.4" pretty_assertions = "1.2.1" +rand = "0.8.5" rstest = "0.17.0" serde = { version = "1.0.197", features = ["derive"] } serde_json = "1.0.116" diff --git a/crates/committer/Cargo.toml b/crates/committer/Cargo.toml index c203bd5d..55c3d722 100644 --- a/crates/committer/Cargo.toml +++ b/crates/committer/Cargo.toml @@ -11,6 +11,7 @@ workspace = true [dev-dependencies] pretty_assertions.workspace = true +rand.workspace = true rstest.workspace = true [dependencies] diff --git a/crates/committer/src/patricia_merkle_tree/types.rs b/crates/committer/src/patricia_merkle_tree/types.rs index bc3e88f4..53f415d2 100644 --- a/crates/committer/src/patricia_merkle_tree/types.rs +++ b/crates/committer/src/patricia_merkle_tree/types.rs @@ -61,6 +61,29 @@ impl NodeIndex { Self::BITS - self.leading_zeros() } + #[allow(dead_code)] + /// Get the LCA (Lowest Common Ancestor) of the two nodes. + pub(crate) fn get_lca(&self, other: &NodeIndex) -> NodeIndex { + if self == other { + return *self; + } + + let bit_length = self.bit_length(); + let other_bit_length = other.bit_length(); + // Bring self to the level of other. + let adapted_self = if self < other { + *self << (other_bit_length - bit_length) + } else { + *self >> (bit_length - other_bit_length) + }; + + let xor = adapted_self.0 ^ other.0; + // The length of the remainder after removing the common prefix of the two nodes. + let post_common_prefix_len = NodeIndex(xor).bit_length(); + let lca = adapted_self.0 >> post_common_prefix_len; + NodeIndex(lca) + } + pub(crate) fn from_starknet_storage_key( key: &StarknetStorageKey, tree_height: &TreeHeight, diff --git a/crates/committer/src/patricia_merkle_tree/types_test.rs b/crates/committer/src/patricia_merkle_tree/types_test.rs index 0d72c07d..a1a1701f 100644 --- a/crates/committer/src/patricia_merkle_tree/types_test.rs +++ b/crates/committer/src/patricia_merkle_tree/types_test.rs @@ -3,6 +3,9 @@ use crate::felt::Felt; use crate::patricia_merkle_tree::node_data::inner_node::{EdgePath, EdgePathLength, PathToBottom}; use crate::patricia_merkle_tree::types::NodeIndex; use crate::patricia_merkle_tree::types::TreeHeight; + +use ethnum::U256; +use rand::Rng; use rstest::rstest; #[rstest] @@ -54,3 +57,40 @@ fn test_cast_to_node_index( }; assert_eq!(actual, expected_node_index.into()); } + +fn get_random_u256(min_trailing_zeros: u8) -> U256 { + let msbits: U256 = rand::thread_rng().gen_range(0..u128::MAX).into(); + let lsbits: u128 = rand::thread_rng().gen(); + ((msbits << (128 - lsbits.leading_zeros())) + lsbits) >> min_trailing_zeros +} + +#[rstest] +#[case(1, 1, 1)] +#[case(2, 5, 2)] +#[case(5, 2, 2)] +#[case(8, 10, 2)] +#[case(9, 12, 1)] +fn test_get_lca(#[case] node_index: u8, #[case] other: u8, #[case] expected: u8) { + let root_index = NodeIndex(node_index.into()); + let other_index = NodeIndex(other.into()); + let lca = root_index.get_lca(&other_index); + let expected = NodeIndex(expected.into()); + assert_eq!(lca, expected); +} + +#[rstest] +fn test_get_lca_big() { + let lca = NodeIndex(get_random_u256(5)); + + let left_child = lca << 1; + let right_child = left_child + 1.into(); + let random_extension = |index: NodeIndex| { + let extension_bits = index.leading_zeros(); + let extension: u128 = rand::thread_rng().gen_range(0..(1 << extension_bits)); + (index << extension_bits) + NodeIndex(U256::from(extension)) + }; + + let left_descendant = random_extension(left_child); + let right_descendant = random_extension(right_child); + assert_eq!(left_descendant.get_lca(&right_descendant), lca); +}