diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs index 7b94bc6b..644f73f6 100644 --- a/evm/src/all_stark.rs +++ b/evm/src/all_stark.rs @@ -220,12 +220,17 @@ mod tests { fn make_keccak_trace( num_keccak_perms: usize, keccak_stark: &KeccakStark, + config: &StarkConfig, rng: &mut R, ) -> Vec> { let keccak_inputs = (0..num_keccak_perms) .map(|_| [0u64; NUM_INPUTS].map(|_| rng.gen())) .collect_vec(); - keccak_stark.generate_trace(keccak_inputs, &mut TimingTree::default()) + keccak_stark.generate_trace( + keccak_inputs, + config.fri_config.num_cap_elements(), + &mut TimingTree::default(), + ) } fn make_keccak_memory_trace( @@ -708,7 +713,8 @@ mod tests { let mut rng = thread_rng(); let num_keccak_perms = 2; - let keccak_trace = make_keccak_trace(num_keccak_perms, &all_stark.keccak_stark, &mut rng); + let keccak_trace = + make_keccak_trace(num_keccak_perms, &all_stark.keccak_stark, config, &mut rng); let keccak_memory_trace = make_keccak_memory_trace(&all_stark.keccak_memory_stark, config); let logic_trace = make_logic_trace(num_logic_rows, &all_stark.logic_stark, config, &mut rng); diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs index 0ed4cbaf..e77769e4 100644 --- a/evm/src/keccak/keccak_stark.rs +++ b/evm/src/keccak/keccak_stark.rs @@ -49,11 +49,14 @@ pub struct KeccakStark { impl, const D: usize> KeccakStark { /// Generate the rows of the trace. Note that this does not generate the permuted columns used /// in our lookup arguments, as those are computed after transposing to column-wise form. - pub(crate) fn generate_trace_rows( + fn generate_trace_rows( &self, inputs: Vec<[u64; NUM_INPUTS]>, + min_rows: usize, ) -> Vec<[F; NUM_COLUMNS]> { - let num_rows = (inputs.len() * NUM_ROUNDS).next_power_of_two(); + let num_rows = (inputs.len() * NUM_ROUNDS) + .max(min_rows) + .next_power_of_two(); let mut rows = Vec::with_capacity(num_rows); for input in inputs.iter() { rows.extend(self.generate_trace_rows_for_perm(*input)); @@ -202,13 +205,14 @@ impl, const D: usize> KeccakStark { pub fn generate_trace( &self, inputs: Vec<[u64; NUM_INPUTS]>, + min_rows: usize, timing: &mut TimingTree, ) -> Vec> { // Generate the witness, except for permuted columns in the lookup argument. let trace_rows = timed!( timing, "generate trace rows", - self.generate_trace_rows(inputs) + self.generate_trace_rows(inputs, min_rows) ); let trace_polys = timed!( timing, @@ -596,7 +600,7 @@ mod tests { f: Default::default(), }; - let rows = stark.generate_trace_rows(vec![input.try_into().unwrap()]); + let rows = stark.generate_trace_rows(vec![input.try_into().unwrap()], 8); let last_row = rows[NUM_ROUNDS - 1]; let output = (0..NUM_INPUTS) .map(|i| { @@ -635,7 +639,7 @@ mod tests { let trace_poly_values = timed!( timing, "generate trace", - stark.generate_trace(input.try_into().unwrap(), &mut timing) + stark.generate_trace(input.try_into().unwrap(), 8, &mut timing) ); // TODO: Cloning this isn't great; consider having `from_values` accept a reference, diff --git a/evm/src/witness/traces.rs b/evm/src/witness/traces.rs index ef754e5e..bd6016e0 100644 --- a/evm/src/witness/traces.rs +++ b/evm/src/witness/traces.rs @@ -112,7 +112,10 @@ impl Traces { let cpu_rows = cpu.into_iter().map(|x| x.into()).collect(); let cpu_trace = trace_rows_to_poly_values(cpu_rows); - let keccak_trace = all_stark.keccak_stark.generate_trace(keccak_inputs, timing); + let keccak_trace = + all_stark + .keccak_stark + .generate_trace(keccak_inputs, cap_elements, timing); let keccak_memory_trace = all_stark.keccak_memory_stark.generate_trace( keccak_memory_inputs, cap_elements,