Add keccak sponge

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 = []

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<u8>,
+}
+
+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<u8> {
+        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<BaseField, G, ScalarField> {
+    sponge: Keccak256Sponge,
+    _base_field: PhantomData<BaseField>,
+    _g: PhantomData<G>,
+    _scalar_field: PhantomData<ScalarField>,
+}
+
+impl<
+        BaseField: PrimeField,
+        ScalarField: PrimeField,
+        P: SWModelParameters + ModelParameters<ScalarField = ScalarField, BaseField = BaseField>,
+    > FqSponge<BaseField, GroupAffine<P>, ScalarField>
+    for Keccak256FqSponge<BaseField, GroupAffine<P>, ScalarField>
+{
+    fn new(_: &'static ArithmeticSpongeParams<BaseField>) -> Self {
+        Keccak256FqSponge {
+            sponge: Keccak256Sponge::new(),
+            _base_field: PhantomData::default(),
+            _g: PhantomData::default(),
+            _scalar_field: PhantomData::default(),
+        }
+    }
+
+    fn absorb_g(&mut self, g: &[GroupAffine<P>]) {
+        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<u8> = 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<u8> = 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<F> {
+    sponge: Keccak256Sponge,
+    _f: PhantomData<F>,
+}
+
+impl<F: PrimeField> FrSponge<F> for Keccak256FrSponge<F> {
+    fn new(_: &'static ArithmeticSpongeParams<F>) -> 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<u8> = 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<F> {
+        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))
+    }
+}

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;

kimchi/src/plonk_sponge.rs

@@ -12,7 +12,9 @@ pub trait FrSponge<Fr: Field> {
     fn new(p: &'static ArithmeticSpongeParams<Fr>) -> 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<Fr: Field> {
     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<PointEvaluations<Vec<Fr>>>);
-}
-
-impl<Fr: PrimeField> FrSponge<Fr> for DefaultFrSponge<Fr, SC> {
-    fn new(params: &'static ArithmeticSpongeParams<Fr>) -> DefaultFrSponge<Fr, SC> {
-        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<Fr> {
-        // 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<PointEvaluations<Vec<Fr>>>) {
-        self.last_squeezed = vec![];
-
         let ProofEvaluations {
             public: _, // Must be absorbed first manually for now, to handle Mina annoyances
             w,
@@ -153,8 +121,36 @@ impl<Fr: PrimeField> FrSponge<Fr> for DefaultFrSponge<Fr, SC> {
         }
 
         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<Fr: PrimeField> FrSponge<Fr> for DefaultFrSponge<Fr, SC> {
+    fn new(params: &'static ArithmeticSpongeParams<Fr>) -> DefaultFrSponge<Fr, SC> {
+        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<Fr> {
+        // 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()
+    }
+}

kimchi/src/tests/generic.rs

@@ -123,3 +123,42 @@ fn test_generic_gate_pairing() {
     .prove_and_verify::<BaseSponge, ScalarSponge>()
     .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<ark_bn254::Fq, ark_bn254::G1Affine, Fp>;
+    type ScalarSponge = crate::keccak_sponge::Keccak256FrSponge<Fp>;
+
+    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<Fp>; COLUMNS] = array::from_fn(|_| vec![Fp::zero(); gates.len()]);
+    fill_in_witness(0, &mut witness, &public);
+
+    // create and verify proof based on the witness
+    <TestFramework<
+        _,
+        poly_commitment::pairing_proof::PairingProof<ark_ec::bn::Bn<ark_bn254::Parameters>>,
+    > 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::<BaseSponge, ScalarSponge>()
+    .unwrap();
+}