diff --git a/crates/core/executor/src/record.rs b/crates/core/executor/src/record.rs index 937ab7609..bd210caec 100644 --- a/crates/core/executor/src/record.rs +++ b/crates/core/executor/src/record.rs @@ -17,7 +17,7 @@ use crate::{ MemInstrEvent, MemoryInitializeFinalizeEvent, MemoryLocalEvent, MemoryRecordEnum, MiscEvent, PrecompileEvent, PrecompileEvents, SyscallEvent, }, - syscalls::SyscallCode, + syscalls::{precompiles::keccak::sponge::GENERAL_BLOCK_SIZE_U32S, SyscallCode}, MipsAirId, Opcode, Program, }; @@ -158,25 +158,52 @@ impl ExecutionRecord { _ => opts.deferred, }; - let chunks = events.chunks_exact(threshold); - if last { - let remainder = chunks.remainder().to_vec(); - if !remainder.is_empty() { - let mut execution_record = ExecutionRecord::new(self.program.clone()); - execution_record.precompile_events.insert(syscall_code, remainder); - shards.push(execution_record); + let mut shards_input = Vec::new(); + let remainder = if syscall_code == SyscallCode::KECCAK_SPONGE { + let mut current_shard = Vec::new(); + let mut current_len = 0; + + for (syscall_event, event) in events { + if let PrecompileEvent::KeccakSponge(event) = &event { + // Here, input_len_u32s must be a multiple of GENERAL_BLOCK_SIZE_U32S. + let input_len = event.input_len_u32s as usize / GENERAL_BLOCK_SIZE_U32S; + + if current_len + input_len > threshold && !current_shard.is_empty() { + let mut record = ExecutionRecord::new(self.program.clone()); + record.precompile_events.insert(syscall_code, current_shard); + shards_input.push(record); + current_shard = Vec::new(); + current_len = 0; + } + current_len += input_len; + } + current_shard.push((syscall_event, event)); } + + current_shard } else { - self.precompile_events.insert(syscall_code, chunks.remainder().to_vec()); + let chunks = events.chunks_exact(threshold); + let remainder = chunks.remainder().to_vec(); + + for chunk in chunks { + let mut record = ExecutionRecord::new(self.program.clone()); + record.precompile_events.insert(syscall_code, chunk.to_vec()); + shards_input.push(record); + } + + remainder + }; + if !remainder.is_empty() { + if last { + let mut record = ExecutionRecord::new(self.program.clone()); + record.precompile_events.insert(syscall_code, remainder); + shards_input.push(record); + } else { + self.precompile_events.insert(syscall_code, remainder); + } } - let mut event_shards = chunks - .map(|chunk| { - let mut execution_record = ExecutionRecord::new(self.program.clone()); - execution_record.precompile_events.insert(syscall_code, chunk.to_vec()); - execution_record - }) - .collect::>(); - shards.append(&mut event_shards); + + shards.extend(shards_input); } if last { diff --git a/crates/core/executor/src/syscalls/mod.rs b/crates/core/executor/src/syscalls/mod.rs index 0c1beb583..6dbd2cb9e 100644 --- a/crates/core/executor/src/syscalls/mod.rs +++ b/crates/core/executor/src/syscalls/mod.rs @@ -6,7 +6,7 @@ mod context; mod deferred; mod halt; mod hint; -mod precompiles; +pub(crate) mod precompiles; mod unconstrained; mod verify; mod write; diff --git a/examples/keccak-precompile/host/src/main.rs b/examples/keccak-precompile/host/src/main.rs index 20a413522..90ac75e53 100644 --- a/examples/keccak-precompile/host/src/main.rs +++ b/examples/keccak-precompile/host/src/main.rs @@ -7,7 +7,7 @@ fn prove_keccak_rust() { let mut stdin = ZKMStdin::new(); // load input - let args = env::var("ARGS").unwrap_or("data-to-hash".to_string()); + let args = env::var("ARGS").unwrap_or("a2279f61f227e186fe9a3b9ba199d8747ddfe637153fc3bd0151def164a925ad data-to-hash".to_string()); // assume the first arg is the hash output(which is a public input), and the second is the input. let args: Vec<&str> = args.split_whitespace().collect(); assert_eq!(args.len(), 2);