Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Propagate witness of each Keccak step #1732

Merged
merged 5 commits into from
Jan 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 15 additions & 15 deletions optimism/src/keccak/column.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,21 +53,21 @@ pub enum KeccakColumn {

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct KeccakColumns<T> {
pub(crate) hash_index: T,
pub(crate) step_index: T,
pub(crate) flag_round: T, // Coeff Round = 0 | 1 .. 24
pub(crate) flag_absorb: T, // Coeff Absorb = 0 | 1
pub(crate) flag_squeeze: T, // Coeff Squeeze = 0 | 1
pub(crate) flag_root: T, // Coeff Root = 0 | 1
pub(crate) flag_pad: T, // Coeff Pad = 0 | 1
pub(crate) flag_length: T, // Coeff Length 0 | 1 .. 136
pub(crate) two_to_pad: T, // 2^PadLength
pub(crate) inverse_round: T, // Round^-1
pub(crate) flags_bytes: [T; RATE_IN_BYTES], // 136 boolean values
pub(crate) pad_suffix: [T; 5], // 5 values with padding suffix
pub(crate) round_constants: [T; QUARTERS], // Round constants
pub(crate) curr: [T; ZKVM_KECCAK_COLS_CURR], // Curr[0..1965)
pub(crate) next: [T; ZKVM_KECCAK_COLS_NEXT], // Next[0..100)
pub hash_index: T,
pub step_index: T,
pub flag_round: T, // Coeff Round = 0 | 1 .. 24
pub flag_absorb: T, // Coeff Absorb = 0 | 1
pub flag_squeeze: T, // Coeff Squeeze = 0 | 1
pub flag_root: T, // Coeff Root = 0 | 1
pub flag_pad: T, // Coeff Pad = 0 | 1
pub flag_length: T, // Coeff Length 0 | 1 .. 136
pub two_to_pad: T, // 2^PadLength
pub inverse_round: T, // Round^-1
pub flags_bytes: [T; RATE_IN_BYTES], // 136 boolean values
pub pad_suffix: [T; 5], // 5 values with padding suffix
pub round_constants: [T; QUARTERS], // Round constants
pub curr: [T; ZKVM_KECCAK_COLS_CURR], // Curr[0..1965)
pub next: [T; ZKVM_KECCAK_COLS_NEXT], // Next[0..100)
}

impl<T: Clone> KeccakColumns<T> {
Expand Down
169 changes: 162 additions & 7 deletions optimism/src/main.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,15 @@
use ark_bn254::FrParameters;
use ark_ec::bn::Bn;
use ark_ff::{UniformRand, Zero};
use ark_ff::{Fp256, UniformRand, Zero};
use kimchi_optimism::{
cannon::{self, Meta, Start, State},
cannon_cli,
keccak::interpreter::KeccakInterpreter,
mips::{proof, witness},
keccak::{
column::KeccakColumns,
interpreter::KeccakInterpreter,
proof::{self as keccak_proof, KeccakProofInputs},
},
mips::{proof, witness as mips_witness},
preimage_oracle::PreImageOracle,
};
use poly_commitment::pairing_proof::PairingProof;
Expand Down Expand Up @@ -70,7 +75,7 @@ pub fn main() -> ExitCode {
srs
};

let mut env = witness::Env::<ark_bn254::Fr>::create(cannon::PAGE_SIZE as usize, state, po);
let mut env = mips_witness::Env::<ark_bn254::Fr>::create(cannon::PAGE_SIZE as usize, state, po);

let mut folded_witness = proof::ProofInputs::<
ark_ec::short_weierstrass_jacobian::GroupAffine<ark_bn254::g1::Parameters>,
Expand All @@ -94,6 +99,54 @@ pub fn main() -> ExitCode {
error: Vec::with_capacity(domain_size),
};

// The keccak environment is extracted inside the loop

let mut keccak_folded_witness = KeccakProofInputs::<
ark_ec::short_weierstrass_jacobian::GroupAffine<ark_bn254::g1::Parameters>,
>::default();

let keccak_reset_pre_folding_witness =
|keccak_columns: &mut KeccakColumns<Vec<Fp256<FrParameters>>>| {
// Resize without deallocating
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is not resizing I think.

keccak_columns.hash_index.clear();
keccak_columns.step_index.clear();
keccak_columns.flag_round.clear();
keccak_columns.flag_absorb.clear();
keccak_columns.flag_squeeze.clear();
keccak_columns.flag_root.clear();
keccak_columns.flag_pad.clear();
keccak_columns.flag_length.clear();
keccak_columns.two_to_pad.clear();
keccak_columns.inverse_round.clear();
keccak_columns.flags_bytes.iter_mut().for_each(Vec::clear);
keccak_columns.pad_suffix.iter_mut().for_each(Vec::clear);
keccak_columns
.round_constants
.iter_mut()
.for_each(Vec::clear);
keccak_columns.curr.iter_mut().for_each(Vec::clear);
keccak_columns.next.iter_mut().for_each(Vec::clear);
};

let mut keccak_current_pre_folding_witness: KeccakColumns<Vec<Fp256<FrParameters>>> =
KeccakColumns {
hash_index: Vec::with_capacity(domain_size),
step_index: Vec::with_capacity(domain_size),
flag_round: Vec::with_capacity(domain_size),
flag_absorb: Vec::with_capacity(domain_size),
flag_squeeze: Vec::with_capacity(domain_size),
flag_root: Vec::with_capacity(domain_size),
flag_pad: Vec::with_capacity(domain_size),
flag_length: Vec::with_capacity(domain_size),
two_to_pad: Vec::with_capacity(domain_size),
inverse_round: Vec::with_capacity(domain_size),
flags_bytes: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
pad_suffix: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
round_constants: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the round constants are always the same no?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, they are fixed. It is a pity we use witness columns to store them. In Kimchi they were passed on as gate coefficients, but with folding we cannot really do that without making instances differ from each other.

curr: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
next: std::array::from_fn(|_| Vec::with_capacity(domain_size)),
};

while !env.halt {
env.step(&configuration, &meta, &start);

Expand All @@ -103,7 +156,88 @@ pub fn main() -> ExitCode {
keccak_env.step();
}

// TODO: update the witness with the Keccak step columns before resetting the environment
// Update the witness with the Keccak step columns before resetting the environment
// TODO: simplify the contents of the KeccakColumns or create an iterator for it
keccak_current_pre_folding_witness
.hash_index
.push(keccak_env.keccak_witness.hash_index);
keccak_current_pre_folding_witness
.step_index
.push(keccak_env.keccak_witness.step_index);
keccak_current_pre_folding_witness
.flag_round
.push(keccak_env.keccak_witness.flag_round);
keccak_current_pre_folding_witness
.flag_absorb
.push(keccak_env.keccak_witness.flag_absorb);
keccak_current_pre_folding_witness
.flag_squeeze
.push(keccak_env.keccak_witness.flag_squeeze);
keccak_current_pre_folding_witness
.flag_root
.push(keccak_env.keccak_witness.flag_root);
keccak_current_pre_folding_witness
.flag_pad
.push(keccak_env.keccak_witness.flag_pad);
keccak_current_pre_folding_witness
.flag_length
.push(keccak_env.keccak_witness.flag_length);
keccak_current_pre_folding_witness
.two_to_pad
.push(keccak_env.keccak_witness.two_to_pad);
keccak_current_pre_folding_witness
.inverse_round
.push(keccak_env.keccak_witness.inverse_round);
for (env_wit, pre_fold_wit) in keccak_env
.keccak_witness
.flags_bytes
.iter()
.zip(keccak_current_pre_folding_witness.flags_bytes.iter_mut())
{
pre_fold_wit.push(*env_wit);
}
for (env_wit, pre_fold_wit) in keccak_env
.keccak_witness
.pad_suffix
.iter()
.zip(keccak_current_pre_folding_witness.pad_suffix.iter_mut())
{
pre_fold_wit.push(*env_wit);
}
for (env_wit, pre_fold_wit) in keccak_env.keccak_witness.round_constants.iter().zip(
keccak_current_pre_folding_witness
.round_constants
.iter_mut(),
) {
pre_fold_wit.push(*env_wit);
}
for (env_wit, pre_fold_wit) in keccak_env
.keccak_witness
.curr
.iter()
.zip(keccak_current_pre_folding_witness.curr.iter_mut())
{
pre_fold_wit.push(*env_wit);
}
for (env_wit, pre_fold_wit) in keccak_env
.keccak_witness
.next
.iter()
.zip(keccak_current_pre_folding_witness.next.iter_mut())
{
pre_fold_wit.push(*env_wit);
}

if keccak_current_pre_folding_witness.step_index.len() == DOMAIN_SIZE {
keccak_proof::fold::<_, OpeningProof, BaseSponge, ScalarSponge>(
domain,
&srs,
&mut keccak_folded_witness,
&keccak_current_pre_folding_witness,
);
keccak_reset_pre_folding_witness(&mut keccak_current_pre_folding_witness);
}

// TODO: create READ lookup tables

// When the Keccak interpreter is finished, we can reset the environment
Expand Down Expand Up @@ -154,15 +288,36 @@ pub fn main() -> ExitCode {
}

{
// MIPS
let proof =
proof::prove::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, folded_witness);
println!("Generated a proof:\n{:?}", proof);
let verifies =
proof::verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
if verifies {
println!("The proof verifies")
println!("The MIPS proof verifies")
} else {
println!("The MIPS proof doesn't verify")
}
}

{
// KECCAK
let keccak_proof = keccak_proof::prove::<_, OpeningProof, BaseSponge, ScalarSponge>(
domain,
&srs,
keccak_folded_witness,
);
println!("Generated a proof:\n{:?}", keccak_proof);
let verifies = keccak_proof::verify::<_, OpeningProof, BaseSponge, ScalarSponge>(
domain,
&srs,
&keccak_proof,
);
if verifies {
println!("The KECCAK proof verifies")
} else {
println!("The proof doesn't verify")
println!("The KECCAK proof doesn't verify")
}
}

Expand Down