starknet_os: generalize sha resource test to include sha512

crates/starknet_os/src/resource_utils_test.rs

@@ -5,6 +5,7 @@ use cairo_vm::types::builtin_name::BuiltinName;
 use cairo_vm::types::layout_name::LayoutName;
 use cairo_vm::types::relocatable::MaybeRelocatable;
 use cairo_vm::vm::runners::cairo_runner::ExecutionResources;
+use rand::distr::{Distribution, StandardUniform};
 use rand::rngs::StdRng;
 use rand::{RngExt, SeedableRng};
 use starknet_types_core::felt::Felt;
@@ -21,6 +22,9 @@ use crate::test_utils::{
     SHA256_BATCH_RESOURCES_CONSTANT,
     SHA256_BATCH_RESOURCES_LINEAR,
     SHA256_BATCH_SIZE,
+    SHA512_BATCH_RESOURCES_CONSTANT,
+    SHA512_BATCH_RESOURCES_LINEAR,
+    SHA512_BATCH_SIZE,
 };
 
 /// SHA-256 block compression: takes 8 u32 state words and 16 u32 message words, returns the
@@ -35,17 +39,39 @@ fn sha_256_update_state(state: &[u32; 8], message: &[u32; 16]) -> [u32; 8] {
     new_state
 }
 
-fn run_finalize_sha256(number_of_blocks: usize) -> ExecutionResources {
-    // Build the SHA-256 instance array. Each instance is 32 felts:
+/// SHA-512 block compression: takes 8 u64 state words and 16 u64 message words, returns the
+/// new 8 u64 state. Wraps `sha2::compress512` (message words are big-endian per the SHA-512 spec).
+fn sha_512_update_state(state: &[u64; 8], message: &[u64; 16]) -> [u64; 8] {
+    let block = sha2::digest::generic_array::GenericArray::from_exact_iter(
+        message.iter().flat_map(|word| word.to_be_bytes()),
+    )
+    .expect("message is exactly 128 bytes");
+    let mut new_state = *state;
+    sha2::compress512(&mut new_state, &[block]);
+    new_state
+}
+
+/// Use T=u32 for sha256, T=u64 for sha512.
+fn run_finalize_sha<T>(
+    number_of_blocks: usize,
+    cairo_finalize_fn: &str,
+    sha_update_state_fn: fn(&[T; 8], &[T; 16]) -> [T; 8],
+) -> ExecutionResources
+where
+    T: Clone + Copy,
+    StandardUniform: Distribution<T> + Distribution<[T; 8]> + Distribution<[T; 16]>,
+    Felt: From<T>,
+{
+    // Build the SHA instance array. Each instance is 32 felts:
     // [message (16) | initial_state (8) | output_state (8)].
     let mut rng = StdRng::seed_from_u64(42);
-    let mut input: Vec<u32> = Vec::new();
+    let mut input: Vec<T> = Vec::new();
     for _ in 0..number_of_blocks {
-        let message: [u32; 16] = rng.random();
-        let state: [u32; 8] = rng.random();
+        let message: [T; 16] = rng.random();
+        let state: [T; 8] = rng.random();
         input.extend_from_slice(&message);
         input.extend_from_slice(&state);
-        let output_state = sha_256_update_state(&state, &message);
+        let output_state = sha_update_state_fn(&state, &message);
         input.extend_from_slice(&output_state);
     }
 
@@ -61,25 +87,25 @@ fn run_finalize_sha256(number_of_blocks: usize) -> ExecutionResources {
     let (mut cairo_runner, program, entrypoint) = initialize_cairo_runner(
         &runner_config,
         OS_PROGRAM_BYTES,
-        "starkware.cairo.common.cairo_sha256.sha256_utils.finalize_sha256",
+        cairo_finalize_fn,
         &implicit_args,
         HashMap::new(),
     )
     .unwrap();
 
-    let sha256_start = cairo_runner
+    let sha_start = cairo_runner
         .vm
         .gen_arg(
             &input.iter().map(|&word| MaybeRelocatable::Int(Felt::from(word))).collect::<Vec<_>>(),
         )
         .unwrap()
         .get_relocatable()
         .unwrap();
-    let sha256_end = (sha256_start + input.len()).unwrap();
+    let sha_end = (sha_start + input.len()).unwrap();
 
     let explicit_args = [
-        EndpointArg::Value(ValueArg::Single(MaybeRelocatable::RelocatableValue(sha256_start))),
-        EndpointArg::Value(ValueArg::Single(MaybeRelocatable::RelocatableValue(sha256_end))),
+        EndpointArg::Value(ValueArg::Single(MaybeRelocatable::RelocatableValue(sha_start))),
+        EndpointArg::Value(ValueArg::Single(MaybeRelocatable::RelocatableValue(sha_end))),
     ];
     run_cairo_0_entrypoint(
         entrypoint,
@@ -102,12 +128,15 @@ fn run_finalize_sha256(number_of_blocks: usize) -> ExecutionResources {
 fn test_finalize_sha256() {
     // Sha256 batching resources has a factor that is linear in the number of rounds, and a constant
     // factor. Sample the execution at two points to compute both factors.
+    let cairo_finalize_fn = "starkware.cairo.common.cairo_sha256.sha256_utils.finalize_sha256";
     let number_of_blocks_1 = 8_usize;
     let number_of_rounds_1 = (number_of_blocks_1 - 1) / SHA256_BATCH_SIZE + 1;
     let number_of_blocks_2 = number_of_blocks_1 + SHA256_BATCH_SIZE;
     let number_of_rounds_2 = (number_of_blocks_2 - 1) / SHA256_BATCH_SIZE + 1;
-    let resources_1 = run_finalize_sha256(number_of_blocks_1);
-    let resources_2 = run_finalize_sha256(number_of_blocks_2);
+    let resources_1 =
+        run_finalize_sha::<u32>(number_of_blocks_1, cairo_finalize_fn, sha_256_update_state);
+    let resources_2 =
+        run_finalize_sha::<u32>(number_of_blocks_2, cairo_finalize_fn, sha_256_update_state);
 
     assert_eq!(number_of_rounds_2 - number_of_rounds_1, 1);
     let linear_factor = (&resources_2 - &resources_1).filter_unused_builtins();
@@ -117,3 +146,28 @@ fn test_finalize_sha256() {
     assert_eq!(linear_factor, *SHA256_BATCH_RESOURCES_LINEAR);
     assert_eq!(constant_factor, *SHA256_BATCH_RESOURCES_CONSTANT);
 }
+
+/// Tests that the `finalize_sha512` Cairo function from the OS program consumes the expected
+/// resources.
+#[test]
+fn test_finalize_sha512() {
+    // Sha256 batching resources has a linear factor and a constant factor. Sample the execution at
+    // two points to compute both factors.
+    let cairo_finalize_fn = "starkware.cairo.common.cairo_sha512.sha512_utils.finalize_sha512";
+    let number_of_blocks_1 = 2_usize;
+    let number_of_rounds_1 = (number_of_blocks_1 - 1) / SHA512_BATCH_SIZE + 1;
+    let number_of_blocks_2 = number_of_blocks_1 + SHA512_BATCH_SIZE;
+    let number_of_rounds_2 = (number_of_blocks_2 - 1) / SHA512_BATCH_SIZE + 1;
+    let resources_1 =
+        run_finalize_sha::<u64>(number_of_blocks_1, cairo_finalize_fn, sha_512_update_state);
+    let resources_2 =
+        run_finalize_sha::<u64>(number_of_blocks_2, cairo_finalize_fn, sha_512_update_state);
+
+    assert_eq!(number_of_rounds_2 - number_of_rounds_1, 1);
+    let linear_factor = (&resources_2 - &resources_1).filter_unused_builtins();
+    let constant_factor =
+        (&resources_1 - &(&linear_factor * number_of_rounds_1)).filter_unused_builtins();
+
+    assert_eq!(linear_factor, *SHA512_BATCH_RESOURCES_LINEAR);
+    assert_eq!(constant_factor, *SHA512_BATCH_RESOURCES_CONSTANT);
+}

crates/starknet_os/src/test_utils.rs

@@ -47,3 +47,30 @@ pub static SHA256_BATCH_RESOURCES_CONSTANT: LazyLock<ExecutionResources> =
             3 / BUILTIN_INSTANCE_SIZES.get(&BuiltinName::range_check).unwrap(),
         )]),
     });
+
+// Resources consumed by the SHA-512 batch phase, separated into linear and constant factors.
+pub const SHA512_BATCH_SIZE: usize = 3;
+pub static SHA512_BATCH_RESOURCES_LINEAR: LazyLock<ExecutionResources> =
+    LazyLock::new(|| ExecutionResources {
+        n_steps: 13710,
+        n_memory_holes: 0,
+        builtin_instance_counter: BTreeMap::from([
+            (
+                BuiltinName::bitwise,
+                9960 / BUILTIN_INSTANCE_SIZES.get(&BuiltinName::bitwise).unwrap(),
+            ),
+            (
+                BuiltinName::range_check,
+                192 / BUILTIN_INSTANCE_SIZES.get(&BuiltinName::range_check).unwrap(),
+            ),
+        ]),
+    });
+pub static SHA512_BATCH_RESOURCES_CONSTANT: LazyLock<ExecutionResources> =
+    LazyLock::new(|| ExecutionResources {
+        n_steps: 49,
+        n_memory_holes: 0,
+        builtin_instance_counter: BTreeMap::from([(
+            BuiltinName::range_check,
+            3 / BUILTIN_INSTANCE_SIZES.get(&BuiltinName::range_check).unwrap(),
+        )]),
+    });