blockifier: move logic inside felt_size_count

Author: avivg-starkware

crates/blockifier/src/execution/contract_class.rs

@@ -84,6 +84,19 @@ impl FeltSizeCount {
     pub(crate) fn n_felts(&self) -> usize {
         self.small + self.large
     }
+
+    /// Returns the total number of `u32` words required to encode all felts
+    /// according to encode_felts_to_u32s func.
+    pub(crate) fn encoded_u32_len(&self) -> usize {
+        self.large * Self::U32_WORDS_PER_LARGE_FELT + self.small * Self::U32_WORDS_PER_SMALL_FELT
+    }
+
+    /// Returns the number of BLAKE opcodes required to hash the felts.
+    /// Each BLAKE opcode processes one message block of [`U32_WORDS_PER_MESSAGE`] `u32`s
+    /// (partial messages are padded).
+    pub(crate) fn blake_opcode_count(&self) -> usize {
+        self.encoded_u32_len().div_ceil(Self::U32_WORDS_PER_MESSAGE)
+    }
 }
 
 /// Counts felts in bytecode by size (small < 2^63, large >= 2^63).
@@ -104,6 +117,25 @@ impl From<&[BigUintAsHex]> for FeltSizeCount {
     }
 }
 
+/// Counts felts in bytecode by size (small < 2^63, large >= 2^63).
+// TODO(AvivG): Remove this.
+impl From<&Vec<Felt>> for FeltSizeCount {
+    fn from(felts: &Vec<Felt>) -> Self {
+        // TODO(AvivG): use blake2s::SMALL_THRESHOLD.
+        const SMALL_THRESHOLD: Felt = Felt::from_hex_unchecked("8000000000000000");
+
+        let (small, large) = felts.iter().fold((0, 0), |(small_count, large_count), x| {
+            if *x < SMALL_THRESHOLD {
+                (small_count + 1, large_count)
+            } else {
+                (small_count, large_count + 1)
+            }
+        });
+
+        FeltSizeCount { small, large }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum NestedFeltCounts {
     Leaf(usize, FeltSizeCount), // (leaf length, felt size groups)

crates/blockifier/src/execution/contract_class_test.rs

@@ -2,13 +2,15 @@ use std::collections::HashSet;
 use std::sync::Arc;
 
 use assert_matches::assert_matches;
+use blake2s::encode_felts_to_u32s;
 use blockifier_test_utils::cairo_versions::{CairoVersion, RunnableCairo1};
 use blockifier_test_utils::contracts::FeatureContract;
 use cairo_lang_starknet_classes::NestedIntList;
 use cairo_lang_utils::bigint::BigUintAsHex;
 use rstest::rstest;
 use starknet_api::contract_class::compiled_class_hash::{HashVersion, HashableCompiledClass};
 use starknet_api::contract_class::ContractClass;
+use starknet_types_core::felt::Felt;
 
 use crate::execution::contract_class::{
     CompiledClassV1,
@@ -136,3 +138,16 @@ fn test_create_bytecode_segment_felt_sizes(
     let result = NestedFeltCounts::new(&bytecode_segment_lengths, &bytecode);
     assert_eq!(result, expected_structure);
 }
+
+#[rstest]
+#[case::boundary_small_felt(vec![Felt::from((1u64 << 63) - 1)])]
+#[case::boundary_at_2_63(vec![Felt::from(1u64 << 63)])]
+#[case::very_large_felt(vec![Felt::from_hex("0x800000000000011000000000000000000000000000000000000000000000000").unwrap()])]
+#[case::mixed_small_large(vec![Felt::from(42), Felt::from(1u64 << 63), Felt::from(1337)])]
+#[case::many_large(vec![Felt::from(1u64 << 63); 100])]
+fn test_encode_u32s_lengths(#[case] test_data: Vec<Felt>) {
+    let actual_u32s = encode_felts_to_u32s(test_data.clone());
+    let estimated_u32s_len = FeltSizeCount::from(&test_data).encoded_u32_len();
+
+    assert_eq!(actual_u32s.len(), estimated_u32s_len);
+}

crates/blockifier/src/execution/execution_utils.rs

@@ -394,18 +394,6 @@ mod blake_estimation {
     pub const STEPS_EMPTY_INPUT: usize = 170;
 }
 
-/// Calculates the total number of u32s required to encode the given number of big and small felts.
-/// Big felts encode to 8 u32s each, small felts encode to 2 u32s each.
-fn total_u32s_from_felts(n_big_felts: usize, n_small_felts: usize) -> usize {
-    let big_u32s = n_big_felts
-        .checked_mul(FeltSizeCount::U32_WORDS_PER_LARGE_FELT)
-        .expect("Overflow computing big felts u32s");
-    let small_u32s = n_small_felts
-        .checked_mul(FeltSizeCount::U32_WORDS_PER_SMALL_FELT)
-        .expect("Overflow computing small felts u32s");
-    big_u32s.checked_add(small_u32s).expect("Overflow computing total u32s")
-}
-
 fn base_steps_for_blake_hash(n_u32s: usize) -> usize {
     let rem_u32s = n_u32s % FeltSizeCount::U32_WORDS_PER_MESSAGE;
     if rem_u32s == 0 {
@@ -434,7 +422,7 @@ fn felts_steps(n_big_felts: usize, n_small_felts: usize) -> usize {
 fn estimate_steps_of_encode_felt252_data_and_calc_blake_hash(
     felt_size_groups: &FeltSizeCount,
 ) -> usize {
-    let total_u32s = total_u32s_from_felts(felt_size_groups.large, felt_size_groups.small);
+    let total_u32s = felt_size_groups.encoded_u32_len();
     if total_u32s == 0 {
         // The empty input case is a special case.
         return blake_estimation::STEPS_EMPTY_INPUT;
@@ -446,14 +434,6 @@ fn estimate_steps_of_encode_felt252_data_and_calc_blake_hash(
     base_steps.checked_add(felt_steps).expect("Overflow computing total Blake hash steps")
 }
 
-/// Returns the number of BLAKE opcodes needed to hash the given felts.
-/// Each BLAKE opcode processes 16 u32s (partial messages are padded).
-fn count_blake_opcode(felt_size_groups: &FeltSizeCount) -> usize {
-    // Count the total number of u32s to be hashed.
-    let total_u32s = total_u32s_from_felts(felt_size_groups.large, felt_size_groups.small);
-    total_u32s.div_ceil(FeltSizeCount::U32_WORDS_PER_MESSAGE)
-}
-
 /// Estimates resource usage for `encode_felt252_data_and_calc_blake_hash` in the Starknet OS.
 ///
 /// # Encoding Details
@@ -482,7 +462,7 @@ pub fn encode_and_blake_hash_resources(
 
     EstimatedExecutionResources::V2Hash {
         resources,
-        blake_count: count_blake_opcode(felt_size_groups),
+        blake_count: felt_size_groups.blake_opcode_count(),
     }
 }
 

crates/blockifier/src/execution/execution_utils_test.rs

@@ -6,7 +6,6 @@ use starknet_types_core::felt::Felt;
 use crate::execution::contract_class::FeltSizeCount;
 use crate::execution::execution_utils::blake_estimation::STEPS_EMPTY_INPUT;
 use crate::execution::execution_utils::{
-    count_blake_opcode,
     encode_and_blake_hash_resources,
     estimate_steps_of_encode_felt252_data_and_calc_blake_hash,
 };
@@ -35,7 +34,7 @@ fn test_zero_inputs() {
     assert_eq!(steps, STEPS_EMPTY_INPUT, "Unexpected base step cost for zero inputs");
 
     // No opcodes should be emitted.
-    let opcodes = count_blake_opcode(&FeltSizeCount::default());
+    let opcodes = FeltSizeCount::default().blake_opcode_count();
     assert_eq!(opcodes, 0, "Expected zero BLAKE opcodes for zero inputs");
 
     // Should result in base cost only (no opcode cost).
@@ -47,7 +46,7 @@ fn test_zero_inputs() {
 
 // TODO(AvivG): Add tests for:
 // - `estimate_steps_of_encode_felt252_data_and_calc_blake_hash` simple cases (felts input).
-// - `count_blake_opcode` simple cases (felts input).
+// - `blake_opcode_count` simple cases (felts input).
 // - `cost_of_encode_felt252_data_and_calc_blake_hash` simple cases (felts input) (including partial
 //   remainder).
 // - `cost_of_encode_felt252_data_and_calc_blake_hash` compare against actual execution resources