From 15709b8845e9c90d884f419e14c44420da97c2a5 Mon Sep 17 00:00:00 2001 From: mrmr1993 Date: Sun, 25 Jun 2023 17:58:23 +0100 Subject: [PATCH] Add keccak sponge --- Cargo.lock | 1 + kimchi/Cargo.toml | 2 + kimchi/src/keccak_sponge.rs | 164 ++++++++++++++++++++++++++++++++++++ kimchi/src/lib.rs | 2 + kimchi/src/plonk_sponge.rs | 70 ++++++++------- kimchi/src/tests/generic.rs | 39 +++++++++ 6 files changed, 241 insertions(+), 37 deletions(-) create mode 100644 kimchi/src/keccak_sponge.rs diff --git a/Cargo.lock b/Cargo.lock index b26d8e1c04..a0cbcd4b25 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1103,6 +1103,7 @@ dependencies = [ "strum", "strum_macros", "thiserror", + "tiny-keccak", "turshi", "wasm-bindgen", ] diff --git a/kimchi/Cargo.toml b/kimchi/Cargo.toml index a2dedda78a..b1b5d2f920 100644 --- a/kimchi/Cargo.toml +++ b/kimchi/Cargo.toml @@ -36,6 +36,7 @@ once_cell = "1.10.0" hex = "0.4" strum = "0.24.0" strum_macros = "0.24.0" +tiny-keccak = { version = "2.0.2", features = ["keccak"], optional = true } # TODO: audit this disjoint-set = "0.0.2" @@ -83,5 +84,6 @@ default = [] internal_tracing = [ "internal-tracing/enabled" ] ocaml_types = [ "ocaml", "ocaml-gen", "poly-commitment/ocaml_types", "mina-poseidon/ocaml_types", "internal-tracing/ocaml_types" ] bn254 = [ "ark-bn254" ] +keccak-sponges = [ "tiny-keccak" ] wasm_types = [ "wasm-bindgen" ] check_feature_flags = [] diff --git a/kimchi/src/keccak_sponge.rs b/kimchi/src/keccak_sponge.rs new file mode 100644 index 0000000000..e1a1d39332 --- /dev/null +++ b/kimchi/src/keccak_sponge.rs @@ -0,0 +1,164 @@ +use crate::plonk_sponge::FrSponge; +use ark_ec::{short_weierstrass_jacobian::GroupAffine, ModelParameters, SWModelParameters}; +use ark_ff::{fields::FpParameters, BigInteger, PrimeField}; +use mina_poseidon::{ + poseidon::ArithmeticSpongeParams, + sponge::{FqSponge, ScalarChallenge}, +}; +use std::marker::PhantomData; +use tiny_keccak::{Hasher, Keccak}; + +/// A sponge designed to be directly compatible with the EVM's Keccak precompile. +#[derive(Debug, Clone)] +struct Keccak256Sponge { + pending: Vec, +} + +impl Keccak256Sponge { + /// Create a new cryptographic sponge backed by keccak + fn new() -> Self { + Keccak256Sponge { pending: vec![] } + } + + /// Absorb arbitrary bytes + fn absorb_bytes(&mut self, x: &[u8]) { + self.pending.extend(x.iter().map(|x| *x)) + } + + /// Squeeze an output from the sponge + fn squeeze(&mut self, n: usize) -> Vec { + let mut final_state = Vec::with_capacity(n); + let mut counter = 0; + while counter < n { + // Create a fresh keccak instance, and hash the entire pending contents into `output`. + let mut hasher = Keccak::v256(); + hasher.update(self.pending.as_slice()); + let mut output = [0u8; 32]; + hasher.finalize(&mut output); + + // Extend the `final_state` buffer with any additional output bytes. + for i in 0..32 { + counter += 1; + if counter >= n { + break; + } + final_state.push(output[i]); + } + + // Update the pending state to contain only the current output. + self.pending = output.to_vec(); + } + final_state + } +} + +#[derive(Debug, Clone)] +pub struct Keccak256FqSponge { + sponge: Keccak256Sponge, + _base_field: PhantomData, + _g: PhantomData, + _scalar_field: PhantomData, +} + +impl< + BaseField: PrimeField, + ScalarField: PrimeField, + P: SWModelParameters + ModelParameters, + > FqSponge, ScalarField> + for Keccak256FqSponge, ScalarField> +{ + fn new(_: &'static ArithmeticSpongeParams) -> Self { + Keccak256FqSponge { + sponge: Keccak256Sponge::new(), + _base_field: PhantomData::default(), + _g: PhantomData::default(), + _scalar_field: PhantomData::default(), + } + } + + fn absorb_g(&mut self, g: &[GroupAffine

]) { + for g in g { + if g.infinity { + // absorb a fake point (0, 0) + let zero = BaseField::zero(); + self.absorb_fq(&[zero, zero]); + } else { + self.absorb_fq(&[g.x, g.y]); + } + } + } + + fn absorb_fq(&mut self, x: &[BaseField]) { + let repr_bytes: usize = (BaseField::Params::MODULUS_BITS as usize + 7) / 8; + let mut bytes: Vec = Vec::with_capacity(repr_bytes * x.len()); + for x in x { + bytes.extend(x.into_repr().to_bytes_be()); + } + self.sponge.absorb_bytes(bytes.as_slice()) + } + + fn absorb_fr(&mut self, x: &[ScalarField]) { + let repr_bytes: usize = (ScalarField::Params::MODULUS_BITS as usize + 7) / 8; + let mut bytes: Vec = Vec::with_capacity(repr_bytes * x.len()); + for x in x { + bytes.extend(x.into_repr().to_bytes_be()); + } + self.sponge.absorb_bytes(bytes.as_slice()) + } + + fn challenge(&mut self) -> ScalarField { + let repr_bytes: usize = (ScalarField::Params::MODULUS_BITS as usize as usize + 7) / 8; + ScalarField::from_be_bytes_mod_order(&self.sponge.squeeze(repr_bytes / 2)) + } + + fn challenge_fq(&mut self) -> BaseField { + let repr_bytes: usize = (BaseField::Params::MODULUS_BITS as usize + 7) / 8; + BaseField::from_be_bytes_mod_order(&self.sponge.squeeze(repr_bytes / 2)) + } + + fn digest(mut self) -> ScalarField { + let repr_bytes: usize = (ScalarField::Params::MODULUS_BITS as usize + 7) / 8; + ScalarField::from_be_bytes_mod_order(&self.sponge.squeeze(repr_bytes)) + } + + fn digest_fq(mut self) -> BaseField { + let repr_bytes: usize = (BaseField::Params::MODULUS_BITS as usize + 7) / 8; + BaseField::from_be_bytes_mod_order(&self.sponge.squeeze(repr_bytes)) + } +} + +#[derive(Debug, Clone)] +pub struct Keccak256FrSponge { + sponge: Keccak256Sponge, + _f: PhantomData, +} + +impl FrSponge for Keccak256FrSponge { + fn new(_: &'static ArithmeticSpongeParams) -> Self { + Keccak256FrSponge { + sponge: Keccak256Sponge::new(), + _f: PhantomData::default(), + } + } + + fn absorb_multiple(&mut self, x: &[F]) { + let repr_bytes: usize = (F::Params::MODULUS_BITS as usize + 7) / 8; + let mut bytes: Vec = Vec::with_capacity(repr_bytes * x.len()); + for x in x { + bytes.extend(x.into_repr().to_bytes_be()); + } + self.sponge.absorb_bytes(bytes.as_slice()) + } + + fn challenge(&mut self) -> ScalarChallenge { + let repr_bytes: usize = (F::Params::MODULUS_BITS as usize + 7) / 8; + ScalarChallenge(F::from_be_bytes_mod_order( + &self.sponge.squeeze(repr_bytes / 2), + )) + } + + fn digest(mut self) -> F { + let repr_bytes: usize = (F::Params::MODULUS_BITS as usize + 7) / 8; + F::from_be_bytes_mod_order(&self.sponge.squeeze(repr_bytes)) + } +} diff --git a/kimchi/src/lib.rs b/kimchi/src/lib.rs index 6a1e004c57..b5ebc096fd 100644 --- a/kimchi/src/lib.rs +++ b/kimchi/src/lib.rs @@ -15,6 +15,8 @@ pub mod bench; pub mod circuits; pub mod curve; pub mod error; +#[cfg(feature = "keccak-sponges")] +pub mod keccak_sponge; pub mod lagrange_basis_evaluations; pub mod linearization; pub mod oracles; diff --git a/kimchi/src/plonk_sponge.rs b/kimchi/src/plonk_sponge.rs index da54379f4d..c7982ac262 100644 --- a/kimchi/src/plonk_sponge.rs +++ b/kimchi/src/plonk_sponge.rs @@ -12,7 +12,9 @@ pub trait FrSponge { fn new(p: &'static ArithmeticSpongeParams) -> Self; /// Absorbs the field element into the sponge. - fn absorb(&mut self, x: &Fr); + fn absorb(&mut self, x: &Fr) { + self.absorb_multiple(&[*x]) + } /// Absorbs a slice of field elements into the sponge. fn absorb_multiple(&mut self, x: &[Fr]); @@ -24,41 +26,7 @@ pub trait FrSponge { fn digest(self) -> Fr; /// Absorbs the given evaluations into the sponge. - // TODO: IMO this function should be inlined in prover/verifier - fn absorb_evaluations(&mut self, e: &ProofEvaluations>>); -} - -impl FrSponge for DefaultFrSponge { - fn new(params: &'static ArithmeticSpongeParams) -> DefaultFrSponge { - DefaultFrSponge { - sponge: ArithmeticSponge::new(params), - last_squeezed: vec![], - } - } - - fn absorb(&mut self, x: &Fr) { - self.last_squeezed = vec![]; - self.sponge.absorb(&[*x]); - } - - fn absorb_multiple(&mut self, x: &[Fr]) { - self.last_squeezed = vec![]; - self.sponge.absorb(x); - } - - fn challenge(&mut self) -> ScalarChallenge { - // TODO: why involve sponge_5_wires here? - ScalarChallenge(self.squeeze(mina_poseidon::sponge::CHALLENGE_LENGTH_IN_LIMBS)) - } - - fn digest(mut self) -> Fr { - self.sponge.squeeze() - } - - // We absorb all evaluations of the same polynomial at the same time fn absorb_evaluations(&mut self, e: &ProofEvaluations>>) { - self.last_squeezed = vec![]; - let ProofEvaluations { w, z, @@ -88,8 +56,36 @@ impl FrSponge for DefaultFrSponge { } points.into_iter().for_each(|p| { - self.sponge.absorb(&p.zeta); - self.sponge.absorb(&p.zeta_omega); + self.absorb_multiple(&p.zeta); + self.absorb_multiple(&p.zeta_omega); }) } } + +impl FrSponge for DefaultFrSponge { + fn new(params: &'static ArithmeticSpongeParams) -> DefaultFrSponge { + DefaultFrSponge { + sponge: ArithmeticSponge::new(params), + last_squeezed: vec![], + } + } + + fn absorb(&mut self, x: &Fr) { + self.last_squeezed = vec![]; + self.sponge.absorb(&[*x]); + } + + fn absorb_multiple(&mut self, x: &[Fr]) { + self.last_squeezed = vec![]; + self.sponge.absorb(x); + } + + fn challenge(&mut self) -> ScalarChallenge { + // TODO: why involve sponge_5_wires here? + ScalarChallenge(self.squeeze(mina_poseidon::sponge::CHALLENGE_LENGTH_IN_LIMBS)) + } + + fn digest(mut self) -> Fr { + self.sponge.squeeze() + } +} diff --git a/kimchi/src/tests/generic.rs b/kimchi/src/tests/generic.rs index b3d923e67b..4ed03c6781 100644 --- a/kimchi/src/tests/generic.rs +++ b/kimchi/src/tests/generic.rs @@ -124,3 +124,42 @@ fn test_generic_gate_pairing() { .prove_and_verify::() .unwrap(); } + +#[cfg(feature = "bn254")] +#[cfg(feature = "keccak-sponges")] +#[test] +fn test_generic_gate_pairing_keccak() { + type Fp = ark_bn254::Fr; + type BaseSponge = + crate::keccak_sponge::Keccak256FqSponge; + type ScalarSponge = crate::keccak_sponge::Keccak256FrSponge; + + use ark_ff::UniformRand; + use ark_poly::EvaluationDomain; + + let public = vec![Fp::from(3u8); 5]; + let gates = create_circuit(0, public.len()); + + let rng = &mut rand::rngs::OsRng; + let x = Fp::rand(rng); + + // create witness + let mut witness: [Vec; COLUMNS] = array::from_fn(|_| vec![Fp::zero(); gates.len()]); + fill_in_witness(0, &mut witness, &public); + + // create and verify proof based on the witness + >, + > as Default>::default() + .gates(gates) + .witness(witness) + .public_inputs(public) + .setup_with_custom_srs(|d1| { + let mut srs = poly_commitment::pairing_proof::PairingSRS::create(x, d1.size()); + srs.full_srs.add_lagrange_basis(d1); + srs + }) + .prove_and_verify::() + .unwrap(); +}