From 95b9fafd90d7d305e91de64c2aa0999d225f7021 Mon Sep 17 00:00:00 2001 From: Linda Guiga Date: Fri, 6 Oct 2023 10:23:09 -0400 Subject: [PATCH 1/3] Remove reg_preimage columns in KeccakStark --- evm/src/all_stark.rs | 28 ++++- evm/src/keccak/columns.rs | 14 +-- evm/src/keccak/keccak_stark.rs | 115 ++++++++----------- evm/src/keccak_sponge/keccak_sponge_stark.rs | 12 +- evm/src/witness/traces.rs | 10 +- evm/src/witness/util.rs | 8 +- 6 files changed, 95 insertions(+), 92 deletions(-) diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs index e5f631e81d..eb82b5533e 100644 --- a/evm/src/all_stark.rs +++ b/evm/src/all_stark.rs @@ -96,7 +96,8 @@ pub(crate) fn all_cross_table_lookups() -> Vec> { ctl_arithmetic(), ctl_byte_packing(), ctl_keccak_sponge(), - ctl_keccak(), + ctl_keccak_inputs(), + ctl_keccak_outputs(), ctl_logic(), ctl_memory(), ] @@ -131,16 +132,33 @@ fn ctl_byte_packing() -> CrossTableLookup { ) } -fn ctl_keccak() -> CrossTableLookup { +// We now need two different looked tables for `KeccakStark`: +// one for the inputs and one for the outputs. +// They are linked with the timestamp. +fn ctl_keccak_inputs() -> CrossTableLookup { let keccak_sponge_looking = TableWithColumns::new( Table::KeccakSponge, - keccak_sponge_stark::ctl_looking_keccak(), + keccak_sponge_stark::ctl_looking_keccak_input(), Some(keccak_sponge_stark::ctl_looking_keccak_filter()), ); let keccak_looked = TableWithColumns::new( Table::Keccak, - keccak_stark::ctl_data(), - Some(keccak_stark::ctl_filter()), + keccak_stark::ctl_data_input(), + Some(keccak_stark::ctl_filter_input()), + ); + CrossTableLookup::new(vec![keccak_sponge_looking], keccak_looked) +} + +fn ctl_keccak_outputs() -> CrossTableLookup { + let keccak_sponge_looking = TableWithColumns::new( + Table::KeccakSponge, + keccak_sponge_stark::ctl_looking_keccak_output(), + Some(keccak_sponge_stark::ctl_looking_keccak_filter()), + ); + let keccak_looked = TableWithColumns::new( + Table::Keccak, + keccak_stark::ctl_data_output(), + Some(keccak_stark::ctl_filter_output()), ); CrossTableLookup::new(vec![keccak_sponge_looking], keccak_looked) } diff --git a/evm/src/keccak/columns.rs b/evm/src/keccak/columns.rs index afd92ad9ec..d9a71af4ff 100644 --- a/evm/src/keccak/columns.rs +++ b/evm/src/keccak/columns.rs @@ -20,7 +20,7 @@ pub fn reg_input_limb(i: usize) -> Column { let y = i_u64 / 5; let x = i_u64 % 5; - let reg_low_limb = reg_preimage(x, y); + let reg_low_limb = reg_a(x, y); let is_high_limb = i % 2; Column::single(reg_low_limb + is_high_limb) } @@ -48,15 +48,11 @@ const R: [[u8; 5]; 5] = [ [27, 20, 39, 8, 14], ]; -const START_PREIMAGE: usize = NUM_ROUNDS; -/// Registers to hold the original input to a permutation, i.e. the input to the first round. -pub(crate) const fn reg_preimage(x: usize, y: usize) -> usize { - debug_assert!(x < 5); - debug_assert!(y < 5); - START_PREIMAGE + (x * 5 + y) * 2 -} +/// Column holding the timestamp, used to link inputs and outputs +/// in the `KeccakSpongeStark`. +pub(crate) const TIMESTAMP: usize = NUM_ROUNDS; -const START_A: usize = START_PREIMAGE + 5 * 5 * 2; +const START_A: usize = TIMESTAMP + 1; pub(crate) const fn reg_a(x: usize, y: usize) -> usize { debug_assert!(x < 5); debug_assert!(y < 5); diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs index 189380f541..6ce280d228 100644 --- a/evm/src/keccak/keccak_stark.rs +++ b/evm/src/keccak/keccak_stark.rs @@ -11,13 +11,13 @@ use plonky2::plonk::plonk_common::reduce_with_powers_ext_circuit; use plonky2::timed; use plonky2::util::timing::TimingTree; +use super::columns::reg_input_limb; use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; use crate::cross_table_lookup::Column; use crate::evaluation_frame::{StarkEvaluationFrame, StarkFrame}; use crate::keccak::columns::{ reg_a, reg_a_prime, reg_a_prime_prime, reg_a_prime_prime_0_0_bit, reg_a_prime_prime_prime, - reg_b, reg_c, reg_c_prime, reg_input_limb, reg_output_limb, reg_preimage, reg_step, - NUM_COLUMNS, + reg_b, reg_c, reg_c_prime, reg_output_limb, reg_step, NUM_COLUMNS, TIMESTAMP, }; use crate::keccak::constants::{rc_value, rc_value_bit}; use crate::keccak::logic::{ @@ -33,16 +33,26 @@ pub(crate) const NUM_ROUNDS: usize = 24; /// Number of 64-bit elements in the Keccak permutation input. pub(crate) const NUM_INPUTS: usize = 25; -pub fn ctl_data() -> Vec> { +pub fn ctl_data_output() -> Vec> { + let mut res: Vec<_> = Column::singles((0..2 * NUM_INPUTS).map(reg_output_limb)).collect(); + res.push(Column::single(TIMESTAMP)); + res +} + +pub fn ctl_data_input() -> Vec> { let mut res: Vec<_> = (0..2 * NUM_INPUTS).map(reg_input_limb).collect(); - res.extend(Column::singles((0..2 * NUM_INPUTS).map(reg_output_limb))); + res.push(Column::single(TIMESTAMP)); res } -pub fn ctl_filter() -> Column { +pub fn ctl_filter_output() -> Column { Column::single(reg_step(NUM_ROUNDS - 1)) } +pub fn ctl_filter_input() -> Column { + Column::single(reg_step(0)) +} + #[derive(Copy, Clone, Default)] pub struct KeccakStark { pub(crate) f: PhantomData, @@ -53,7 +63,7 @@ impl, const D: usize> KeccakStark { /// in our lookup arguments, as those are computed after transposing to column-wise form. fn generate_trace_rows( &self, - inputs: Vec<[u64; NUM_INPUTS]>, + inputs: Vec<([u64; NUM_INPUTS], usize)>, min_rows: usize, ) -> Vec<[F; NUM_COLUMNS]> { let num_rows = (inputs.len() * NUM_ROUNDS) @@ -72,20 +82,19 @@ impl, const D: usize> KeccakStark { rows } - fn generate_trace_rows_for_perm(&self, input: [u64; NUM_INPUTS]) -> Vec<[F; NUM_COLUMNS]> { + fn generate_trace_rows_for_perm( + &self, + input_and_timestamp: ([u64; NUM_INPUTS], usize), + ) -> Vec<[F; NUM_COLUMNS]> { let mut rows = vec![[F::ZERO; NUM_COLUMNS]; NUM_ROUNDS]; - - // Populate the preimage for each row. + let input = input_and_timestamp.0; + let timestamp = input_and_timestamp.1; + // Set the timestamp of the current input. + // It will be checked against the value in `KeccakSponge`. + // The timestamp is used to link the input and output of + // the same permutation together. for round in 0..24 { - for x in 0..5 { - for y in 0..5 { - let input_xy = input[y * 5 + x]; - let reg_preimage_lo = reg_preimage(x, y); - let reg_preimage_hi = reg_preimage_lo + 1; - rows[round][reg_preimage_lo] = F::from_canonical_u64(input_xy & 0xFFFFFFFF); - rows[round][reg_preimage_hi] = F::from_canonical_u64(input_xy >> 32); - } - } + rows[round][TIMESTAMP] = F::from_canonical_usize(timestamp); } // Populate the round input for the first round. @@ -220,7 +229,7 @@ impl, const D: usize> KeccakStark { pub fn generate_trace( &self, - inputs: Vec<[u64; NUM_INPUTS]>, + inputs: Vec<([u64; NUM_INPUTS], usize)>, min_rows: usize, timing: &mut TimingTree, ) -> Vec> { @@ -269,26 +278,14 @@ impl, const D: usize> Stark for KeccakStark(); + yield_constr.constraint( + sum_round_flags * not_final_step * (next_values[TIMESTAMP] - local_values[TIMESTAMP]), + ); // C'[x, z] = xor(C[x, z], C[x - 1, z], C[x + 1, z - 1]). for x in 0..5 { @@ -454,34 +451,13 @@ impl, const D: usize> Stark for KeccakStark = (0..NUM_PERMS).map(|_| rand::random()).collect(); + let input: Vec<([u64; NUM_INPUTS], usize)> = + (0..NUM_PERMS).map(|_| (rand::random(), 0)).collect(); let mut timing = TimingTree::new("prove", log::Level::Debug); let trace_poly_values = timed!( diff --git a/evm/src/keccak_sponge/keccak_sponge_stark.rs b/evm/src/keccak_sponge/keccak_sponge_stark.rs index 2ed31c1fee..b01d700edc 100644 --- a/evm/src/keccak_sponge/keccak_sponge_stark.rs +++ b/evm/src/keccak_sponge/keccak_sponge_stark.rs @@ -47,7 +47,7 @@ pub(crate) fn ctl_looked_data() -> Vec> { .collect() } -pub(crate) fn ctl_looking_keccak() -> Vec> { +pub(crate) fn ctl_looking_keccak_input() -> Vec> { let cols = KECCAK_SPONGE_COL_MAP; let mut res: Vec<_> = Column::singles( [ @@ -57,6 +57,13 @@ pub(crate) fn ctl_looking_keccak() -> Vec> { .concat(), ) .collect(); + res.push(Column::single(cols.timestamp)); + + res +} + +pub(crate) fn ctl_looking_keccak_output() -> Vec> { + let cols = KECCAK_SPONGE_COL_MAP; // We recover the 32-bit digest limbs from their corresponding bytes, // and then append them to the rest of the updated state limbs. @@ -68,9 +75,10 @@ pub(crate) fn ctl_looking_keccak() -> Vec> { ) }); - res.extend(digest_u32s); + let mut res: Vec<_> = digest_u32s.collect(); res.extend(Column::singles(&cols.partial_updated_state_u32s)); + res.push(Column::single(cols.timestamp)); res } diff --git a/evm/src/witness/traces.rs b/evm/src/witness/traces.rs index c4cf832dd5..91035fc403 100644 --- a/evm/src/witness/traces.rs +++ b/evm/src/witness/traces.rs @@ -36,7 +36,7 @@ pub(crate) struct Traces { pub(crate) cpu: Vec>, pub(crate) logic_ops: Vec, pub(crate) memory_ops: Vec, - pub(crate) keccak_inputs: Vec<[u64; keccak::keccak_stark::NUM_INPUTS]>, + pub(crate) keccak_inputs: Vec<([u64; keccak::keccak_stark::NUM_INPUTS], usize)>, pub(crate) keccak_sponge_ops: Vec, } @@ -131,18 +131,18 @@ impl Traces { self.byte_packing_ops.push(op); } - pub fn push_keccak(&mut self, input: [u64; keccak::keccak_stark::NUM_INPUTS]) { - self.keccak_inputs.push(input); + pub fn push_keccak(&mut self, input: [u64; keccak::keccak_stark::NUM_INPUTS], clock: usize) { + self.keccak_inputs.push((input, clock)); } - pub fn push_keccak_bytes(&mut self, input: [u8; KECCAK_WIDTH_BYTES]) { + pub fn push_keccak_bytes(&mut self, input: [u8; KECCAK_WIDTH_BYTES], clock: usize) { let chunks = input .chunks(size_of::()) .map(|chunk| u64::from_le_bytes(chunk.try_into().unwrap())) .collect_vec() .try_into() .unwrap(); - self.push_keccak(chunks); + self.push_keccak(chunks, clock); } pub fn push_keccak_sponge(&mut self, op: KeccakSpongeOp) { diff --git a/evm/src/witness/util.rs b/evm/src/witness/util.rs index 068a8e1113..dbe4c0ede5 100644 --- a/evm/src/witness/util.rs +++ b/evm/src/witness/util.rs @@ -229,7 +229,9 @@ pub(crate) fn keccak_sponge_log( address.increment(); } xor_into_sponge(state, &mut sponge_state, block.try_into().unwrap()); - state.traces.push_keccak_bytes(sponge_state); + state + .traces + .push_keccak_bytes(sponge_state, clock * NUM_CHANNELS); keccakf_u8s(&mut sponge_state); } @@ -254,7 +256,9 @@ pub(crate) fn keccak_sponge_log( final_block[KECCAK_RATE_BYTES - 1] = 0b10000000; } xor_into_sponge(state, &mut sponge_state, &final_block); - state.traces.push_keccak_bytes(sponge_state); + state + .traces + .push_keccak_bytes(sponge_state, clock * NUM_CHANNELS); state.traces.push_keccak_sponge(KeccakSpongeOp { base_address, From dddd833aacd6ea25cc03ae03cf695715cbadf522 Mon Sep 17 00:00:00 2001 From: Linda Guiga Date: Fri, 6 Oct 2023 15:08:10 -0400 Subject: [PATCH 2/3] Apply comments --- evm/src/all_stark.rs | 12 +++++------ evm/src/keccak/keccak_stark.rs | 22 ++++++++++---------- evm/src/keccak_sponge/keccak_sponge_stark.rs | 4 ++-- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs index eb82b5533e..079ff114c4 100644 --- a/evm/src/all_stark.rs +++ b/evm/src/all_stark.rs @@ -138,13 +138,13 @@ fn ctl_byte_packing() -> CrossTableLookup { fn ctl_keccak_inputs() -> CrossTableLookup { let keccak_sponge_looking = TableWithColumns::new( Table::KeccakSponge, - keccak_sponge_stark::ctl_looking_keccak_input(), + keccak_sponge_stark::ctl_looking_keccak_inputs(), Some(keccak_sponge_stark::ctl_looking_keccak_filter()), ); let keccak_looked = TableWithColumns::new( Table::Keccak, - keccak_stark::ctl_data_input(), - Some(keccak_stark::ctl_filter_input()), + keccak_stark::ctl_data_inputs(), + Some(keccak_stark::ctl_filter_inputs()), ); CrossTableLookup::new(vec![keccak_sponge_looking], keccak_looked) } @@ -152,13 +152,13 @@ fn ctl_keccak_inputs() -> CrossTableLookup { fn ctl_keccak_outputs() -> CrossTableLookup { let keccak_sponge_looking = TableWithColumns::new( Table::KeccakSponge, - keccak_sponge_stark::ctl_looking_keccak_output(), + keccak_sponge_stark::ctl_looking_keccak_outputs(), Some(keccak_sponge_stark::ctl_looking_keccak_filter()), ); let keccak_looked = TableWithColumns::new( Table::Keccak, - keccak_stark::ctl_data_output(), - Some(keccak_stark::ctl_filter_output()), + keccak_stark::ctl_data_outputs(), + Some(keccak_stark::ctl_filter_outputs()), ); CrossTableLookup::new(vec![keccak_sponge_looking], keccak_looked) } diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs index 6ce280d228..f1c1e0fa59 100644 --- a/evm/src/keccak/keccak_stark.rs +++ b/evm/src/keccak/keccak_stark.rs @@ -33,24 +33,24 @@ pub(crate) const NUM_ROUNDS: usize = 24; /// Number of 64-bit elements in the Keccak permutation input. pub(crate) const NUM_INPUTS: usize = 25; -pub fn ctl_data_output() -> Vec> { - let mut res: Vec<_> = Column::singles((0..2 * NUM_INPUTS).map(reg_output_limb)).collect(); +pub fn ctl_data_inputs() -> Vec> { + let mut res: Vec<_> = (0..2 * NUM_INPUTS).map(reg_input_limb).collect(); res.push(Column::single(TIMESTAMP)); res } -pub fn ctl_data_input() -> Vec> { - let mut res: Vec<_> = (0..2 * NUM_INPUTS).map(reg_input_limb).collect(); +pub fn ctl_data_outputs() -> Vec> { + let mut res: Vec<_> = Column::singles((0..2 * NUM_INPUTS).map(reg_output_limb)).collect(); res.push(Column::single(TIMESTAMP)); res } -pub fn ctl_filter_output() -> Column { - Column::single(reg_step(NUM_ROUNDS - 1)) +pub fn ctl_filter_inputs() -> Column { + Column::single(reg_step(0)) } -pub fn ctl_filter_input() -> Column { - Column::single(reg_step(0)) +pub fn ctl_filter_outputs() -> Column { + Column::single(reg_step(NUM_ROUNDS - 1)) } #[derive(Copy, Clone, Default)] @@ -63,15 +63,15 @@ impl, const D: usize> KeccakStark { /// in our lookup arguments, as those are computed after transposing to column-wise form. fn generate_trace_rows( &self, - inputs: Vec<([u64; NUM_INPUTS], usize)>, + inputs_and_timestamps: Vec<([u64; NUM_INPUTS], usize)>, min_rows: usize, ) -> Vec<[F; NUM_COLUMNS]> { - let num_rows = (inputs.len() * NUM_ROUNDS) + let num_rows = (inputs_and_timestamps.len() * NUM_ROUNDS) .max(min_rows) .next_power_of_two(); let mut rows = Vec::with_capacity(num_rows); - for input in inputs.iter() { + for input in inputs_and_timestamps.iter() { let rows_for_perm = self.generate_trace_rows_for_perm(*input); rows.extend(rows_for_perm); } diff --git a/evm/src/keccak_sponge/keccak_sponge_stark.rs b/evm/src/keccak_sponge/keccak_sponge_stark.rs index b01d700edc..e491252ba8 100644 --- a/evm/src/keccak_sponge/keccak_sponge_stark.rs +++ b/evm/src/keccak_sponge/keccak_sponge_stark.rs @@ -47,7 +47,7 @@ pub(crate) fn ctl_looked_data() -> Vec> { .collect() } -pub(crate) fn ctl_looking_keccak_input() -> Vec> { +pub(crate) fn ctl_looking_keccak_inputs() -> Vec> { let cols = KECCAK_SPONGE_COL_MAP; let mut res: Vec<_> = Column::singles( [ @@ -62,7 +62,7 @@ pub(crate) fn ctl_looking_keccak_input() -> Vec> { res } -pub(crate) fn ctl_looking_keccak_output() -> Vec> { +pub(crate) fn ctl_looking_keccak_outputs() -> Vec> { let cols = KECCAK_SPONGE_COL_MAP; // We recover the 32-bit digest limbs from their corresponding bytes, From f7e40313dd398b6758e271ba0f928ea1463905ff Mon Sep 17 00:00:00 2001 From: Linda Guiga Date: Fri, 6 Oct 2023 15:29:32 -0400 Subject: [PATCH 3/3] Minor cleanup --- evm/src/keccak/keccak_stark.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs index f1c1e0fa59..2745d03302 100644 --- a/evm/src/keccak/keccak_stark.rs +++ b/evm/src/keccak/keccak_stark.rs @@ -71,8 +71,8 @@ impl, const D: usize> KeccakStark { .next_power_of_two(); let mut rows = Vec::with_capacity(num_rows); - for input in inputs_and_timestamps.iter() { - let rows_for_perm = self.generate_trace_rows_for_perm(*input); + for input_and_timestamp in inputs_and_timestamps.iter() { + let rows_for_perm = self.generate_trace_rows_for_perm(*input_and_timestamp); rows.extend(rows_for_perm); }