test(skeleton): get_random_u256

crates/committer/src/patricia_merkle_tree.rs

@@ -4,3 +4,6 @@ pub mod node_data;
 pub mod original_skeleton_tree;
 pub mod types;
 pub mod updated_skeleton_tree;
+
+#[cfg(test)]
+mod test_utils;

crates/committer/src/patricia_merkle_tree/test_utils.rs

@@ -0,0 +1,49 @@
+use ethnum::U256;
+use rand::Rng;
+use rstest::rstest;
+
+/// Generates a random U256 number between low and high (exclusive).
+/// Panics if low > high.
+pub(crate) fn get_random_u256(low: U256, high: U256) -> U256 {
+    assert!(low < high);
+    let high_of_low = low.high();
+    let high_of_high = high.high();
+
+    let delta = high - low;
+    if delta <= u128::MAX {
+        let delta = u128::try_from(delta).unwrap();
+        return low + rand::thread_rng().gen_range(0..delta);
+    }
+
+    // Randomize the high 128 bits in the extracted range, and the low 128 bits in their entire
+    // domain until the result is in range.
+    // As high-low>u128::MAX, the expected number of samples until the loops breaks is bound from
+    // above by 3 (as either:
+    //  1. high_of_high > high_of_low + 1, and there is a 1/3 chance to get a valid result for high
+    //  bits in (high_of_low, high_of_high).
+    //  2. high_of_high == high_of_low + 1, and every possible low 128 bits value is valid either
+    // when the high bits equal high_of_high, or when they equal high_of_low).
+    let randomize = || {
+        U256::from_words(
+            rand::thread_rng().gen_range(*high_of_low..=*high_of_high),
+            rand::thread_rng().gen_range(0..=u128::MAX),
+        )
+    };
+    let mut result = randomize();
+    while result < low || result >= high {
+        result = randomize();
+    }
+    result
+}
+
+#[rstest]
+#[should_panic]
+#[case(U256::ZERO, U256::ZERO)]
+#[case(U256::ZERO, U256::ONE)]
+#[case(U256::ONE, U256::ONE << 128)]
+#[case((U256::ONE<<128)-U256::ONE, U256::ONE << 128)]
+#[case(U256::ONE<<128, (U256::ONE << 128)+U256::ONE)]
+fn test_get_random_u256(#[case] low: U256, #[case] high: U256) {
+    let r = get_random_u256(low, high);
+    assert!(low <= r && r < high);
+}

crates/committer/src/patricia_merkle_tree/types_test.rs

@@ -1,6 +1,7 @@
 use crate::block_committer::input::{ContractAddress, StarknetStorageKey};
 use crate::felt::Felt;
 use crate::patricia_merkle_tree::node_data::inner_node::{EdgePath, EdgePathLength, PathToBottom};
+use crate::patricia_merkle_tree::test_utils::get_random_u256;
 use crate::patricia_merkle_tree::types::NodeIndex;
 use crate::patricia_merkle_tree::types::TreeHeight;
 
@@ -58,12 +59,6 @@ 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)]
@@ -80,7 +75,10 @@ fn test_get_lca(#[case] node_index: u8, #[case] other: u8, #[case] expected: u8)
 
 #[rstest]
 fn test_get_lca_big() {
-    let lca = NodeIndex(get_random_u256(5));
+    let lca = NodeIndex::new(get_random_u256(
+        U256::ZERO,
+        (NodeIndex::MAX_INDEX >> 1).into(),
+    ));
 
     let left_child = lca << 1;
     let right_child = left_child + 1.into();