blockifier: move logic inside felt_size_count

Author: avivg-starkware

crates/blockifier/src/execution/contract_class.rs

@@ -71,9 +71,30 @@ pub struct FeltSizeCount {
 }
 
 impl FeltSizeCount {
+    // Constants that define how felts are encoded into u32s for BLAKE hashing.
+    // Number of `u32` words in a single BLAKE input message block.
+    pub(crate) const U32_WORDS_PER_MESSAGE: usize = 16;
+    // Number of `u32` words a large felt expands into.
+    pub(crate) const U32_WORDS_PER_LARGE_FELT: usize = 8;
+    // Number of `u32` words a small felt expands into.
+    pub(crate) const U32_WORDS_PER_SMALL_FELT: usize = 2;
+
     pub fn n_felts(&self) -> usize {
         self.small + self.large
     }
+
+    /// Returns the total number of `u32` words required to encode all felts
+    /// according to the small/large felt encoding scheme.
+    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).

crates/blockifier/src/execution/execution_utils.rs

@@ -373,16 +373,6 @@ pub fn poseidon_hash_many_cost(data_length: usize) -> ExecutionResources {
     }
 }
 
-// Constants that define how felts are encoded into u32s for BLAKE hashing.
-mod blake_encoding {
-    /// Number of u32s in a Blake input message.
-    pub const N_U32S_MESSAGE: usize = 16;
-
-    /// Number of u32s a felt is encoded into.
-    pub const N_U32S_BIG_FELT: usize = 8;
-    pub const N_U32S_SMALL_FELT: usize = 2;
-}
-
 // Constants used for estimating the cost of BLAKE hashing inside Starknet OS.
 // These values are based on empirical measurement by running
 // `encode_felt252_data_and_calc_blake_hash` on various combinations of big and small felts.
@@ -404,20 +394,8 @@ 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(blake_encoding::N_U32S_BIG_FELT)
-        .expect("Overflow computing big felts u32s");
-    let small_u32s = n_small_felts
-        .checked_mul(blake_encoding::N_U32S_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 % blake_encoding::N_U32S_MESSAGE;
+    let rem_u32s = n_u32s % FeltSizeCount::U32_WORDS_PER_MESSAGE;
     if rem_u32s == 0 {
         blake_estimation::BASE_STEPS_FULL_MSG
     } else {
@@ -441,27 +419,19 @@ fn felts_steps(n_big_felts: usize, n_small_felts: usize) -> usize {
 /// Estimates the number of VM steps needed to hash the given felts with Blake in Starknet OS.
 /// Each small felt unpacks into 2 u32s, and each big felt into 8 u32s.
 /// Adds a base cost depending on whether the total fits exactly into full 16-u32 messages.
-fn compute_blake_hash_steps(n_big_felts: usize, n_small_felts: usize) -> usize {
-    let total_u32s = total_u32s_from_felts(n_big_felts, n_small_felts);
+fn compute_blake_hash_steps(felt_size_groups: &FeltSizeCount) -> usize {
+    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;
     }
 
     let base_steps = base_steps_for_blake_hash(total_u32s);
-    let felt_steps = felts_steps(n_big_felts, n_small_felts);
+    let felt_steps = felts_steps(felt_size_groups.large, felt_size_groups.small);
 
     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(n_big_felts: usize, n_small_felts: usize) -> usize {
-    // Count the total number of u32s to be hashed.
-    let total_u32s = total_u32s_from_felts(n_big_felts, n_small_felts);
-    total_u32s.div_ceil(blake_encoding::N_U32S_MESSAGE)
-}
-
 /// Estimates resource usage for `encode_felt252_data_and_calc_blake_hash` in the Starknet OS.
 ///
 /// # Encoding Details
@@ -475,7 +445,7 @@ fn count_blake_opcode(n_big_felts: usize, n_small_felts: usize) -> usize {
 pub fn encode_and_blake_hash_resources(
     felt_size_groups: &FeltSizeCount,
 ) -> EstimatedExecutionResources {
-    let n_steps = compute_blake_hash_steps(felt_size_groups.large, felt_size_groups.small);
+    let n_steps = compute_blake_hash_steps(felt_size_groups);
     let builtin_instance_counter = match felt_size_groups.n_felts() {
         // The empty case does not use builtins at all.
         0 => HashMap::new(),
@@ -490,7 +460,7 @@ pub fn encode_and_blake_hash_resources(
 
     EstimatedExecutionResources::V2Hash {
         resources,
-        blake_count: count_blake_opcode(felt_size_groups.large, felt_size_groups.small),
+        blake_count: felt_size_groups.blake_opcode_count(),
     }
 }
 

crates/blockifier/src/execution/execution_utils_test.rs

@@ -4,11 +4,9 @@ use pretty_assertions::assert_eq;
 use starknet_types_core::felt::Felt;
 
 use crate::execution::contract_class::FeltSizeCount;
-use crate::execution::execution_utils::blake_encoding::{N_U32S_BIG_FELT, N_U32S_SMALL_FELT};
 use crate::execution::execution_utils::blake_estimation::STEPS_EMPTY_INPUT;
 use crate::execution::execution_utils::{
     compute_blake_hash_steps,
-    count_blake_opcode,
     encode_and_blake_hash_resources,
 };
 
@@ -23,19 +21,19 @@ fn test_u32_constants() {
     let big_u32s = encode_felts_to_u32s(vec![big_felt]);
 
     // Blake estimation constants should match the actual encoding.
-    assert_eq!(small_u32s.len(), N_U32S_SMALL_FELT);
-    assert_eq!(big_u32s.len(), N_U32S_BIG_FELT);
+    assert_eq!(small_u32s.len(), FeltSizeCount::U32_WORDS_PER_SMALL_FELT);
+    assert_eq!(big_u32s.len(), FeltSizeCount::U32_WORDS_PER_LARGE_FELT);
 }
 
 /// Test the edge case of hashing an empty array of felt values.
 #[test]
 fn test_zero_inputs() {
     // logic was written.
-    let steps = compute_blake_hash_steps(0, 0);
+    let steps = compute_blake_hash_steps(&FeltSizeCount { large: 0, small: 0 });
     assert_eq!(steps, STEPS_EMPTY_INPUT, "Unexpected base step cost for zero inputs");
 
     // No opcodes should be emitted.
-    let opcodes = count_blake_opcode(0, 0);
+    let opcodes = FeltSizeCount { large: 0, small: 0 }.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 +45,7 @@ fn test_zero_inputs() {
 
 // TODO(AvivG): Add tests for:
 // - `compute_blake_hash_steps` 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