|
| 1 | +use std::collections::HashMap; |
| 2 | + |
| 3 | +use apollo_starknet_os_program::OS_PROGRAM_BYTES; |
| 4 | +use cairo_vm::types::builtin_name::BuiltinName; |
| 5 | +use cairo_vm::types::layout_name::LayoutName; |
| 6 | +use cairo_vm::types::relocatable::MaybeRelocatable; |
| 7 | +use cairo_vm::vm::runners::cairo_runner::ExecutionResources; |
| 8 | +use rand::rngs::StdRng; |
| 9 | +use rand::{RngExt, SeedableRng}; |
| 10 | +use starknet_types_core::felt::Felt; |
| 11 | + |
| 12 | +use crate::test_utils::cairo_runner::{ |
| 13 | + initialize_cairo_runner, |
| 14 | + run_cairo_0_entrypoint, |
| 15 | + EndpointArg, |
| 16 | + EntryPointRunnerConfig, |
| 17 | + ImplicitArg, |
| 18 | + ValueArg, |
| 19 | +}; |
| 20 | +use crate::test_utils::{SHA256_BATCH_RESOURCES_CONSTANT, SHA256_BATCH_RESOURCES_LINEAR}; |
| 21 | + |
| 22 | +/// SHA-256 block compression: takes 8 u32 state words and 16 u32 message words, returns the |
| 23 | +/// new 8 u32 state. Wraps `sha2::compress256` (message words are big-endian per the SHA-256 spec). |
| 24 | +fn sha_256_update_state(state: &[u32; 8], message: &[u32; 16]) -> [u32; 8] { |
| 25 | + let block = sha2::digest::generic_array::GenericArray::from_exact_iter( |
| 26 | + message.iter().flat_map(|word| word.to_be_bytes()), |
| 27 | + ) |
| 28 | + .expect("message is exactly 64 bytes"); |
| 29 | + let mut new_state = *state; |
| 30 | + sha2::compress256(&mut new_state, &[block]); |
| 31 | + new_state |
| 32 | +} |
| 33 | + |
| 34 | +fn run_finalize_sha256(number_of_blocks: usize) -> ExecutionResources { |
| 35 | + // Build the SHA-256 instance array. Each instance is 32 felts: |
| 36 | + // [message (16) | initial_state (8) | output_state (8)]. |
| 37 | + let mut rng = StdRng::seed_from_u64(42); |
| 38 | + let mut input: Vec<u32> = Vec::new(); |
| 39 | + for _ in 0..number_of_blocks { |
| 40 | + let message: [u32; 16] = rng.random(); |
| 41 | + let state: [u32; 8] = rng.random(); |
| 42 | + input.extend_from_slice(&message); |
| 43 | + input.extend_from_slice(&state); |
| 44 | + let output_state = sha_256_update_state(&state, &message); |
| 45 | + input.extend_from_slice(&output_state); |
| 46 | + } |
| 47 | + |
| 48 | + let runner_config = EntryPointRunnerConfig { |
| 49 | + layout: LayoutName::starknet, |
| 50 | + add_main_prefix_to_entrypoint: false, |
| 51 | + ..Default::default() |
| 52 | + }; |
| 53 | + let implicit_args = [ |
| 54 | + ImplicitArg::Builtin(BuiltinName::range_check), |
| 55 | + ImplicitArg::Builtin(BuiltinName::bitwise), |
| 56 | + ]; |
| 57 | + let (mut cairo_runner, program, entrypoint) = initialize_cairo_runner( |
| 58 | + &runner_config, |
| 59 | + OS_PROGRAM_BYTES, |
| 60 | + "starkware.cairo.common.cairo_sha256.sha256_utils.finalize_sha256", |
| 61 | + &implicit_args, |
| 62 | + HashMap::new(), |
| 63 | + ) |
| 64 | + .unwrap(); |
| 65 | + |
| 66 | + let sha256_start = cairo_runner |
| 67 | + .vm |
| 68 | + .gen_arg( |
| 69 | + &input.iter().map(|&word| MaybeRelocatable::Int(Felt::from(word))).collect::<Vec<_>>(), |
| 70 | + ) |
| 71 | + .unwrap() |
| 72 | + .get_relocatable() |
| 73 | + .unwrap(); |
| 74 | + let sha256_end = (sha256_start + input.len()).unwrap(); |
| 75 | + |
| 76 | + let explicit_args = [ |
| 77 | + EndpointArg::Value(ValueArg::Single(MaybeRelocatable::RelocatableValue(sha256_start))), |
| 78 | + EndpointArg::Value(ValueArg::Single(MaybeRelocatable::RelocatableValue(sha256_end))), |
| 79 | + ]; |
| 80 | + run_cairo_0_entrypoint( |
| 81 | + entrypoint, |
| 82 | + &explicit_args, |
| 83 | + &implicit_args, |
| 84 | + None, |
| 85 | + &mut cairo_runner, |
| 86 | + &program, |
| 87 | + &runner_config, |
| 88 | + &[], |
| 89 | + ) |
| 90 | + .unwrap(); |
| 91 | + |
| 92 | + cairo_runner.get_execution_resources().unwrap() |
| 93 | +} |
| 94 | + |
| 95 | +/// Tests that the `finalize_sha256` Cairo function from the OS program consumes the expected |
| 96 | +/// resources. |
| 97 | +#[test] |
| 98 | +fn test_finalize_sha256() { |
| 99 | + // Sha256 batching resources has a factor that is linear in the number of rounds, and a constant |
| 100 | + // factor. Sample the execution at two points to compute both factors. |
| 101 | + let blocks_in_round = 7_usize; |
| 102 | + let number_of_blocks_1 = 8_usize; |
| 103 | + let number_of_rounds_1 = (number_of_blocks_1 - 1) / blocks_in_round + 1; |
| 104 | + let number_of_blocks_2 = number_of_blocks_1 + blocks_in_round; |
| 105 | + let number_of_rounds_2 = (number_of_blocks_2 - 1) / blocks_in_round + 1; |
| 106 | + let resources_1 = run_finalize_sha256(number_of_blocks_1); |
| 107 | + let resources_2 = run_finalize_sha256(number_of_blocks_2); |
| 108 | + |
| 109 | + assert_eq!(number_of_rounds_2 - number_of_rounds_1, 1); |
| 110 | + let linear_factor = (&resources_2 - &resources_1).filter_unused_builtins(); |
| 111 | + let constant_factor = |
| 112 | + (&resources_1 - &(&linear_factor * number_of_rounds_1)).filter_unused_builtins(); |
| 113 | + |
| 114 | + assert_eq!(linear_factor, *SHA256_BATCH_RESOURCES_LINEAR); |
| 115 | + assert_eq!(constant_factor, *SHA256_BATCH_RESOURCES_CONSTANT); |
| 116 | +} |
0 commit comments