Add generic to keccak

keccak256/src/permutation.rs

@@ -3,6 +3,7 @@
 pub(crate) mod absorb;
 pub(crate) mod base_conversion;
 pub mod circuit;
+pub(crate) mod generic;
 pub(crate) mod iota;
 pub(crate) mod mixing;
 pub(crate) mod pi;

keccak256/src/permutation/circuit.rs

@@ -3,9 +3,9 @@ use crate::{
     common::{NEXT_INPUTS_LANES, PERMUTATION, ROUND_CONSTANTS},
     keccak_arith::*,
     permutation::{
-        base_conversion::BaseConversionConfig, iota::IotaConfig, mixing::MixingConfig,
-        pi::pi_gate_permutation, rho::RhoConfig, tables::FromBase9TableConfig, theta::ThetaConfig,
-        xi::XiConfig,
+        base_conversion::BaseConversionConfig, generic::GenericConfig, iota::IotaConfig,
+        mixing::MixingConfig, pi::pi_gate_permutation, rho::RhoConfig,
+        tables::FromBase9TableConfig, theta::ThetaConfig, xi::XiConfig,
     },
 };
 use eth_types::Field;
@@ -34,7 +34,7 @@ pub struct KeccakFConfig<F: Field> {
 impl<F: Field> KeccakFConfig<F> {
     // We assume state is received in base-9.
     pub fn configure(meta: &mut ConstraintSystem<F>) -> Self {
-        let state = (0..25)
+        let state: [Column<Advice>; 25] = (0..25)
             .map(|_| {
                 let column = meta.advice_column();
                 meta.enable_equality(column);
@@ -50,11 +50,12 @@ impl<F: Field> KeccakFConfig<F> {
             meta.fixed_column(),
             meta.fixed_column(),
         ];
+        let generic = GenericConfig::configure(meta, state[0..3].try_into().unwrap(), fixed[0]);
 
         // theta
         let theta_config = ThetaConfig::configure(meta.selector(), meta, state);
         // rho
-        let rho_config = RhoConfig::configure(meta, state, fixed);
+        let rho_config = RhoConfig::configure(meta, state, fixed, &generic);
         // xi
         let xi_config = XiConfig::configure(meta.selector(), meta, state);
         let iota_config = IotaConfig::configure(meta, state[0], flag, fixed[0]);

keccak256/src/permutation/generic.rs

@@ -0,0 +1,287 @@
+use eth_types::Field;
+use halo2_proofs::{
+    circuit::{AssignedCell, Layouter},
+    plonk::{Advice, Column, ConstraintSystem, Error, Fixed, Selector},
+    poly::Rotation,
+};
+use itertools::Itertools;
+use std::marker::PhantomData;
+
+/// A versatile gate to do running sum, conditional add, and linear combination,
+/// etc.
+#[derive(Clone, Debug)]
+pub struct GenericConfig<F> {
+    q_enable: Selector,
+    io: Column<Advice>,
+    left: Column<Advice>,
+    right: Column<Advice>,
+    _marker: PhantomData<F>,
+}
+
+#[allow(dead_code)]
+impl<F: Field> GenericConfig<F> {
+    pub fn configure(
+        meta: &mut ConstraintSystem<F>,
+        advices: [Column<Advice>; 3],
+        fixed: Column<Fixed>,
+    ) -> Self {
+        let q_enable = meta.selector();
+        let [io, left, right] = advices;
+        meta.enable_equality(io);
+        meta.enable_equality(left);
+        meta.enable_equality(right);
+        meta.enable_constant(fixed);
+
+        meta.create_gate("add", |meta| {
+            let q_enable = meta.query_selector(q_enable);
+            let input = meta.query_advice(io, Rotation::cur());
+            let output = meta.query_advice(io, Rotation::next());
+            let left = meta.query_advice(left, Rotation::cur());
+            let right = meta.query_advice(right, Rotation::cur());
+            vec![q_enable * (output - input - left * right)]
+        });
+
+        Self {
+            q_enable,
+            io,
+            left,
+            right,
+            _marker: PhantomData,
+        }
+    }
+
+    fn add_generic(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        input: AssignedCell<F, F>,
+        left: Option<AssignedCell<F, F>>,
+        right: Option<AssignedCell<F, F>>,
+        value: Option<F>,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        layouter.assign_region(
+            || "add advice",
+            |mut region| {
+                let offset = 0;
+                self.q_enable.enable(&mut region, offset)?;
+                input.copy_advice(|| "input", &mut region, self.io, offset)?;
+                let left = match &left {
+                    Some(x) => {
+                        // copy x to use as a flag
+                        (*x).copy_advice(|| "left adv", &mut region, self.left, offset)?
+                    }
+                    None => {
+                        // constrain advice to 1 for a simple add.
+                        region.assign_advice_from_constant(
+                            || "left const",
+                            self.left,
+                            offset,
+                            F::one(),
+                        )?
+                    }
+                };
+
+                let right = match &right {
+                    Some(right) => {
+                        if value.is_some() {
+                            panic!("right and value can't be both some");
+                        }
+                        right.copy_advice(|| "right adv", &mut region, self.right, offset)?
+                    }
+                    None => {
+                        match value {
+                            Some(value) => region.assign_advice_from_constant(
+                                || "fixed value",
+                                self.right,
+                                offset,
+                                value,
+                            )?,
+                            None => {
+                                // constrain fixed to 1 for a simple add.
+                                region.assign_advice_from_constant(
+                                    || "fixed value",
+                                    self.right,
+                                    offset,
+                                    F::one(),
+                                )?
+                            }
+                        }
+                    }
+                };
+
+                let offset = 1;
+                region.assign_advice(
+                    || "input + x",
+                    self.io,
+                    offset,
+                    || {
+                        Ok(input.value().cloned().ok_or(Error::Synthesis)?
+                            + left.value().cloned().ok_or(Error::Synthesis)?
+                                * right.value().cloned().ok_or(Error::Synthesis)?)
+                    },
+                )
+            },
+        )
+    }
+    /// input += v * x
+    pub fn add_advice_mul_const(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        input: AssignedCell<F, F>,
+        x: AssignedCell<F, F>,
+        v: F,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        self.add_generic(layouter, input, Some(x), None, Some(v))
+    }
+    /// input -= x
+    pub fn sub_advice(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        input: AssignedCell<F, F>,
+        x: AssignedCell<F, F>,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        self.add_generic(layouter, input, Some(x), None, Some(-F::one()))
+    }
+    /// input += v
+    pub fn add_fixed(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        input: AssignedCell<F, F>,
+        value: F,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        self.add_generic(layouter, input, None, None, Some(value))
+    }
+    /// input += flag * v
+    /// No boolean check on the flag, we assume the flag is checked before
+    /// copied to here
+    pub fn conditional_add_const(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        input: AssignedCell<F, F>,
+        flag: AssignedCell<F, F>,
+        value: F,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        self.add_generic(layouter, input, Some(flag), None, Some(value))
+    }
+    /// input += flag * x
+    /// No boolean check on the flag, we assume the flag is checked before
+    /// copied to here
+    pub fn conditional_add_advice(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        input: AssignedCell<F, F>,
+        flag: AssignedCell<F, F>,
+        x: AssignedCell<F, F>,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        self.add_generic(layouter, input, Some(flag), Some(x), None)
+    }
+    fn linear_combine_generic(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        xs: Vec<AssignedCell<F, F>>,
+        ys: Option<Vec<AssignedCell<F, F>>>,
+        vs: Option<Vec<F>>,
+        outcome: Option<AssignedCell<F, F>>,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        debug_assert_eq!(
+            ys.is_some(),
+            vs.is_none(),
+            "They can't both some or both none"
+        );
+        if let Some(ref vs) = vs {
+            debug_assert_eq!(xs.len(), vs.len());
+        }
+        if let Some(ref ys) = ys {
+            debug_assert_eq!(xs.len(), ys.len());
+        }
+        layouter.assign_region(
+            || "linear combine",
+            |mut region| {
+                // | offset |        input |        x |   y |
+                // | ------ | -----------: | -------: | ------: |
+                // | 0      |            0 |       x0 |      y0 |
+                // | 1      |         x0y0 |       x1 |      y1 |
+                // | 2      |  x0y0 + x1y1 |       x2 |      y2 |
+                // | ...    |          ... |      ... |     ... |
+                // | N - 1  |              |  x_(N-1) | y_(N-1) |
+                // | N      |    (sum)     |          |         |
+                let mut acc = region.assign_advice(|| "input 0", self.io, 0, || Ok(F::zero()))?;
+                region.constrain_constant(acc.cell(), F::zero())?;
+                let mut sum = F::zero();
+                for (offset, x) in xs.iter().enumerate() {
+                    self.q_enable.enable(&mut region, offset)?;
+                    x.copy_advice(|| "x", &mut region, self.left, offset)?;
+                    let right = {
+                        match &vs {
+                            Some(vs) => region.assign_advice_from_constant(
+                                || "v",
+                                self.right,
+                                offset,
+                                vs[offset],
+                            )?,
+                            None => match &ys {
+                                Some(ys) => ys[offset].copy_advice(
+                                    || "y",
+                                    &mut region,
+                                    self.right,
+                                    offset,
+                                )?,
+                                None => {
+                                    unreachable!()
+                                }
+                            },
+                        }
+                    };
+                    acc = region.assign_advice(
+                        || "accumulation",
+                        self.io,
+                        offset + 1,
+                        || {
+                            sum += x.value().cloned().ok_or(Error::Synthesis)?
+                                * right.value().cloned().ok_or(Error::Synthesis)?;
+                            Ok(sum)
+                        },
+                    )?;
+                }
+                if let Some(outcome) = &outcome {
+                    region.constrain_equal(outcome.cell(), acc.cell())?;
+                }
+                Ok(acc)
+            },
+        )
+    }
+
+    pub fn linear_combine_consts(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        xs: Vec<AssignedCell<F, F>>,
+        vs: Vec<F>,
+        outcome: Option<AssignedCell<F, F>>,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        self.linear_combine_generic(layouter, xs, None, Some(vs), outcome)
+    }
+
+    pub fn linear_combine_advices(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        xs: Vec<AssignedCell<F, F>>,
+        ys: Vec<AssignedCell<F, F>>,
+        outcome: Option<AssignedCell<F, F>>,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        self.linear_combine_generic(layouter, xs, Some(ys), None, outcome)
+    }
+
+    pub fn running_sum(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        xs: Vec<AssignedCell<F, F>>,
+        outcome: Option<AssignedCell<F, F>>,
+    ) -> Result<AssignedCell<F, F>, Error> {
+        let len = xs.len();
+        self.linear_combine_consts(
+            layouter,
+            xs,
+            (0..len).map(|_| F::one()).collect_vec(),
+            outcome,
+        )
+    }
+}

keccak256/src/permutation/rho.rs

@@ -1,4 +1,5 @@
 use crate::permutation::{
+    generic::GenericConfig,
     rho_checks::{LaneRotateConversionConfig, OverflowCheckConfig},
     rho_helpers::{STEP2_RANGE, STEP3_RANGE},
     tables::{Base13toBase9TableConfig, RangeCheckConfig, SpecialChunkTableConfig},
@@ -26,6 +27,7 @@ impl<F: Field> RhoConfig<F> {
         meta: &mut ConstraintSystem<F>,
         state: [Column<Advice>; 25],
         fixed: [Column<Fixed>; 3],
+        generic: &GenericConfig<F>,
     ) -> Self {
         state.iter().for_each(|col| meta.enable_equality(*col));
         let base13_to_9_table = Base13toBase9TableConfig::configure(meta);
@@ -46,6 +48,7 @@ impl<F: Field> RhoConfig<F> {
             &step2_range_table,
             &step3_range_table,
             state[5..7].try_into().unwrap(),
+            generic.clone(),
         );
         Self {
             lane_config,
@@ -113,12 +116,13 @@ mod tests {
     use crate::common::*;
     use crate::gate_helpers::biguint_to_f;
     use crate::keccak_arith::*;
-    use halo2_proofs::circuit::Layouter;
-    use halo2_proofs::pairing::bn256::Fr as Fp;
-    use halo2_proofs::plonk::Selector;
-    use halo2_proofs::plonk::{Advice, Column, ConstraintSystem, Error};
-    use halo2_proofs::poly::Rotation;
-    use halo2_proofs::{circuit::SimpleFloorPlanner, dev::MockProver, plonk::Circuit};
+    use halo2_proofs::{
+        circuit::{Layouter, SimpleFloorPlanner},
+        dev::MockProver,
+        pairing::bn256::Fr as Fp,
+        plonk::{Advice, Circuit, Column, ConstraintSystem, Error, Selector},
+        poly::Rotation,
+    };
     use itertools::Itertools;
     use std::convert::TryInto;
 
@@ -156,8 +160,10 @@ mod tests {
                     meta.fixed_column(),
                     meta.fixed_column(),
                 ];
+                let generic =
+                    GenericConfig::configure(meta, state[0..3].try_into().unwrap(), fixed[0]);
 
-                let rho_config = RhoConfig::configure(meta, state, fixed);
+                let rho_config = RhoConfig::configure(meta, state, fixed, &generic);
 
                 let q_enable = meta.selector();
                 meta.create_gate("Check states", |meta| {