blockifier: unify from felts and from bytecode to counts

Author: avivg-starkware

crates/blockifier/src/execution/contract_class.rs

@@ -76,22 +76,22 @@ impl FeltSizeCount {
     pub fn n_felts(&self) -> usize {
         self.small + self.large
     }
-}
 
-/// Counts felts in bytecode by size (small < 2^63, large >= 2^63).
-impl From<&[BigUintAsHex]> for FeltSizeCount {
-    fn from(bytecode: &[BigUintAsHex]) -> Self {
+    /// Builds a `FeltSizeCount` by mapping each element to a `Felt` with `to_felt`,
+    /// then classifying it according to its size using `SMALL_THRESHOLD`.
+    pub fn from_slice<T, F>(items: &[T], to_felt: F) -> Self
+    where
+        F: Fn(&T) -> Felt,
+    {
         // TODO(AvivG): use blake2s::SMALL_THRESHOLD.
         const SMALL_THRESHOLD: Felt = Felt::from_hex_unchecked("8000000000000000");
 
-        let (small, large) = bytecode.iter().fold((0, 0), |(small_count, large_count), x| {
-            if Felt::from(&x.value) < SMALL_THRESHOLD {
-                (small_count + 1, large_count)
-            } else {
-                (small_count, large_count + 1)
-            }
-        });
+        let mut small = 0;
+        let mut large = 0;
 
+        for x in items {
+            if to_felt(x) < SMALL_THRESHOLD { small += 1 } else { large += 1 }
+        }
         FeltSizeCount { small, large }
     }
 }
@@ -156,7 +156,10 @@ impl NestedFeltCounts {
 
         match bytecode_segment_lengths {
             NestedIntList::Leaf(len) => {
-                let felt_size_groups = FeltSizeCount::from(&bytecode[..*len]);
+                let felt_size_groups =
+                    FeltSizeCount::from_slice(&bytecode[..*len], |x: &BigUintAsHex| {
+                        Felt::from(&x.value)
+                    });
                 (NestedFeltCounts::Leaf(*len, felt_size_groups), *len)
             }
             NestedIntList::Node(segments_vec) => {

crates/starknet_os/src/hints/hint_implementation/blake2s/blake2s_test.rs

@@ -19,27 +19,9 @@ use crate::test_utils::cairo_runner::{
     ValueArg,
 };
 
-/// Counts the number of small and big felts in the data.
-// TODO(AvivG): use From<&[Felt]> for FeltSizeCount.
-fn data_to_felt_count(data: &[Felt]) -> FeltSizeCount {
-    // TODO(AvivG): Use `Blake2Felt252::SMALL_THRESHOLD` when exposed.
-    const SMALL_THRESHOLD: Felt = Felt::from_hex_unchecked("8000000000000000"); // 2^63
-
-    let felt_size_groups = FeltSizeCount::default();
-
-    data.iter().fold(felt_size_groups, |mut felt_size_groups, felt| {
-        if *felt >= SMALL_THRESHOLD {
-            felt_size_groups.large += 1;
-        } else {
-            felt_size_groups.small += 1;
-        }
-        felt_size_groups
-    })
-}
-
 /// Return the estimated execution resources for Blake2s hashing.
 fn estimated_encode_and_blake_hash_execution_resources(data: &[Felt]) -> ExecutionResources {
-    let felt_size_groups = data_to_felt_count(data);
+    let felt_size_groups = FeltSizeCount::from_slice(data, |&felt| felt);
     let estimated = encode_and_blake_hash_resources(&felt_size_groups);
 
     let mut resources = estimated.resources().clone();