Skip to content

Commit

Permalink
Rename columns in KeccakSponge for clarity
Browse files Browse the repository at this point in the history
  • Loading branch information
Nashtare committed Sep 13, 2023
1 parent f944a08 commit b31c857
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 22 deletions.
8 changes: 5 additions & 3 deletions evm/src/keccak_sponge/columns.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,10 +52,12 @@ pub(crate) struct KeccakSpongeColumnsView<T: Copy> {
pub xored_rate_u32s: [T; KECCAK_RATE_U32S],

/// The entire state (rate + capacity) of the sponge, encoded as 32-bit chunks, after the
/// permutation is applied.
pub updated_state_u32s: [T; KECCAK_WIDTH_U32S],
/// permutation is applied, minus the first limbs where the digest is extracted from.
pub partial_updated_state_u32s: [T; KECCAK_WIDTH_U32S],

pub updated_state_bytes: [T; KECCAK_DIGEST_BYTES],
/// The first part of the state of the sponge, seen as bytes, after the permutation is applied.
/// This also represents the output digest of the Keccak sponge during the squeezing phase.
pub updated_digest_state_bytes: [T; KECCAK_DIGEST_BYTES],
}

// `u8` is guaranteed to have a `size_of` of 1.
Expand Down
46 changes: 27 additions & 19 deletions evm/src/keccak_sponge/keccak_sponge_stark.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ pub(crate) fn ctl_looked_data<F: Field>() -> Vec<Column<F>> {
let mut outputs = Vec::with_capacity(8);
for i in (0..8).rev() {
let cur_col = Column::linear_combination(
cols.updated_state_bytes[i * 4..(i + 1) * 4]
cols.updated_digest_state_bytes[i * 4..(i + 1) * 4]
.iter()
.enumerate()
.map(|(j, &c)| (c, F::from_canonical_u64(1 << (24 - 8 * j)))),
Expand All @@ -53,7 +53,7 @@ pub(crate) fn ctl_looking_keccak<F: Field>() -> Vec<Column<F>> {
[
cols.xored_rate_u32s.as_slice(),
&cols.original_capacity_u32s,
&cols.updated_state_u32s,
&cols.partial_updated_state_u32s,
]
.concat(),
)
Expand Down Expand Up @@ -239,7 +239,9 @@ impl<F: RichField + Extendable<D>, const D: usize> KeccakSpongeStark<F, D> {
block.try_into().unwrap(),
);

sponge_state = row.updated_state_u32s.map(|f| f.to_canonical_u64() as u32);
sponge_state = row
.partial_updated_state_u32s
.map(|f| f.to_canonical_u64() as u32);

rows.push(row.into());
already_absorbed_bytes += KECCAK_RATE_BYTES;
Expand Down Expand Up @@ -357,20 +359,20 @@ impl<F: RichField + Extendable<D>, const D: usize> KeccakSpongeStark<F, D> {
row.xored_rate_u32s = xored_rate_u32s.map(F::from_canonical_u32);

keccakf_u32s(&mut sponge_state);
row.updated_state_u32s = sponge_state.map(F::from_canonical_u32);
row.partial_updated_state_u32s = sponge_state.map(F::from_canonical_u32);
let is_final_block = row.is_final_input_len.iter().copied().sum::<F>() == F::ONE;
if is_final_block {
for (l, &elt) in row.updated_state_u32s[..8].iter().enumerate() {
for (l, &elt) in row.partial_updated_state_u32s[..8].iter().enumerate() {
let mut cur_elt = elt;
(0..4).for_each(|i| {
row.updated_state_bytes[l * 4 + i] =
row.updated_digest_state_bytes[l * 4 + i] =
F::from_canonical_u32((cur_elt.to_canonical_u64() & 0xFF) as u32);
cur_elt = F::from_canonical_u64(cur_elt.to_canonical_u64() >> 8);
});

let mut s = row.updated_state_bytes[l * 4].to_canonical_u64();
let mut s = row.updated_digest_state_bytes[l * 4].to_canonical_u64();
for i in 1..4 {
s += row.updated_state_bytes[l * 4 + i].to_canonical_u64() << (8 * i);
s += row.updated_digest_state_bytes[l * 4 + i].to_canonical_u64() << (8 * i);
}
assert_eq!(elt, F::from_canonical_u64(s), "not equal");
}
Expand Down Expand Up @@ -446,14 +448,14 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakSpongeS

// If this is a full-input block, the next row's "before" should match our "after" state.
for (&current_after, &next_before) in local_values
.updated_state_u32s
.partial_updated_state_u32s
.iter()
.zip(next_values.original_rate_u32s.iter())
{
yield_constr.constraint_transition(is_full_input_block * (next_before - current_after));
}
for (&current_after, &next_before) in local_values
.updated_state_u32s
.partial_updated_state_u32s
.iter()
.skip(KECCAK_RATE_U32S)
.zip(next_values.original_capacity_u32s.iter())
Expand Down Expand Up @@ -483,10 +485,13 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakSpongeS
}

// Adding constraints for byte columns.
for (l, &elt) in local_values.updated_state_u32s[..8].iter().enumerate() {
let mut s = local_values.updated_state_bytes[l * 4];
for (l, &elt) in local_values.partial_updated_state_u32s[..8]
.iter()
.enumerate()
{
let mut s = local_values.updated_digest_state_bytes[l * 4];
for i in 1..4 {
s += local_values.updated_state_bytes[l * 4 + i]
s += local_values.updated_digest_state_bytes[l * 4 + i]
* P::from(FE::from_canonical_usize(1 << (8 * i)));
}
yield_constr.constraint(is_final_block * (s - elt));
Expand Down Expand Up @@ -567,7 +572,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakSpongeS

// If this is a full-input block, the next row's "before" should match our "after" state.
for (&current_after, &next_before) in local_values
.updated_state_u32s
.partial_updated_state_u32s
.iter()
.zip(next_values.original_rate_u32s.iter())
{
Expand All @@ -576,7 +581,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakSpongeS
yield_constr.constraint_transition(builder, constraint);
}
for (&current_after, &next_before) in local_values
.updated_state_u32s
.partial_updated_state_u32s
.iter()
.skip(KECCAK_RATE_U32S)
.zip(next_values.original_capacity_u32s.iter())
Expand Down Expand Up @@ -617,12 +622,15 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakSpongeS
}

// Adding constraints for byte columns.
for (l, &elt) in local_values.updated_state_u32s[..8].iter().enumerate() {
let mut s = local_values.updated_state_bytes[l * 4];
for (l, &elt) in local_values.partial_updated_state_u32s[..8]
.iter()
.enumerate()
{
let mut s = local_values.updated_digest_state_bytes[l * 4];
for i in 1..4 {
s = builder.mul_const_add_extension(
F::from_canonical_usize(1 << (8 * i)),
local_values.updated_state_bytes[l * 4 + i],
local_values.updated_digest_state_bytes[l * 4 + i],
s,
);
}
Expand Down Expand Up @@ -698,7 +706,7 @@ mod tests {
let rows = stark.generate_rows_for_op(op);
assert_eq!(rows.len(), 1);
let last_row: &KeccakSpongeColumnsView<F> = rows.last().unwrap().borrow();
let output = last_row.updated_state_u32s[..8]
let output = last_row.partial_updated_state_u32s[..8]
.iter()
.flat_map(|x| (x.to_canonical_u64() as u32).to_le_bytes())
.collect_vec();
Expand Down

0 comments on commit b31c857

Please sign in to comment.