zk: Rework arithmetic chip to allow for more configuration and cell reuse.

example/gt.rs

@@ -1,5 +1,5 @@
 use darkfi::zk::{
-    arith_chip::{ArithmeticChip, ArithmeticChipConfig},
+    arith_chip::{ArithChip, ArithConfig, ArithInstruction},
     even_bits::{EvenBitsChip, EvenBitsConfig, EvenBitsLookup},
     greater_than::{GreaterThanChip, GreaterThanConfig, GreaterThanInstruction},
 };
@@ -20,7 +20,7 @@ struct ZkConfig {
     advices: [Column<Advice>; 3],
     evenbits_config: EvenBitsConfig,
     greaterthan_config: GreaterThanConfig,
-    arith_config: ArithmeticChipConfig,
+    arith_config: ArithConfig,
 }
 
 impl ZkConfig {
@@ -32,8 +32,8 @@ impl ZkConfig {
         GreaterThanChip::construct(self.greaterthan_config.clone())
     }
 
-    fn arith_chip(&self) -> ArithmeticChip {
-        ArithmeticChip::construct(self.arith_config.clone())
+    fn arith_chip(&self) -> ArithChip {
+        ArithChip::construct(self.arith_config.clone())
     }
 }
 
@@ -72,7 +72,7 @@ impl Circuit<pallas::Base> for ZkCircuit {
             [advices[1], advices[2]],
             primary,
         );
-        let arith_config = ArithmeticChip::configure(meta);
+        let arith_config = ArithChip::configure(meta, advices[1], advices[2], advices[0]);
 
         ZkConfig { primary, advices, evenbits_config, greaterthan_config, arith_config }
     }
@@ -93,7 +93,7 @@ impl Circuit<pallas::Base> for ZkCircuit {
         let v = self.load_private(layouter.namespace(|| "Witness v"), config.advices[0], self.v)?;
         let f = self.load_private(layouter.namespace(|| "Witness t"), config.advices[0], self.f)?;
 
-        let t = ar_chip.mul(layouter.namespace(|| "target value"), v, f)?;
+        let t = ar_chip.mul(layouter.namespace(|| "target value"), &v, &f)?;
 
         eb_chip.decompose(layouter.namespace(|| "y range check"), y.clone())?;
         eb_chip.decompose(layouter.namespace(|| "t range check"), t.clone())?;

src/zk/arith_chip.rs

@@ -1,27 +1,51 @@
 use halo2_proofs::{
+    arithmetic::FieldExt,
     circuit::{AssignedCell, Chip, Layouter},
-    plonk::{Advice, Column, ConstraintSystem, Error, Selector},
+    plonk,
+    plonk::{Advice, Column, ConstraintSystem, Constraints, Selector},
     poly::Rotation,
 };
 use pasta_curves::pallas;
 
-type Variable = AssignedCell<pallas::Base, pallas::Base>;
+pub trait ArithInstruction<F: FieldExt>: Chip<F> {
+    fn add(
+        &self,
+        layouter: impl Layouter<F>,
+        a: &AssignedCell<F, F>,
+        b: &AssignedCell<F, F>,
+    ) -> Result<AssignedCell<F, F>, plonk::Error>;
+
+    fn sub(
+        &self,
+        layouter: impl Layouter<F>,
+        a: &AssignedCell<F, F>,
+        b: &AssignedCell<F, F>,
+    ) -> Result<AssignedCell<F, F>, plonk::Error>;
+
+    fn mul(
+        &self,
+        layouter: impl Layouter<F>,
+        a: &AssignedCell<F, F>,
+        b: &AssignedCell<F, F>,
+    ) -> Result<AssignedCell<F, F>, plonk::Error>;
+}
 
 #[derive(Clone, Debug)]
-pub struct ArithmeticChipConfig {
-    a_col: Column<Advice>,
-    b_col: Column<Advice>,
-    s_add: Selector,
-    s_mul: Selector,
-    s_sub: Selector,
+pub struct ArithConfig {
+    a: Column<Advice>,
+    b: Column<Advice>,
+    c: Column<Advice>,
+    q_add: Selector,
+    q_sub: Selector,
+    q_mul: Selector,
 }
 
-pub struct ArithmeticChip {
-    config: ArithmeticChipConfig,
+pub struct ArithChip {
+    config: ArithConfig,
 }
 
-impl Chip<pallas::Base> for ArithmeticChip {
-    type Config = ArithmeticChipConfig;
+impl Chip<pallas::Base> for ArithChip {
+    type Config = ArithConfig;
     type Loaded = ();
 
     fn config(&self) -> &Self::Config {
@@ -33,187 +57,125 @@ impl Chip<pallas::Base> for ArithmeticChip {
     }
 }
 
-impl ArithmeticChip {
-    pub fn construct(config: ArithmeticChipConfig) -> Self {
-        Self { config }
-    }
-
-    pub fn configure(cs: &mut ConstraintSystem<pallas::Base>) -> ArithmeticChipConfig {
-        let a_col = cs.advice_column();
-        let b_col = cs.advice_column();
-
-        cs.enable_equality(a_col);
-        cs.enable_equality(b_col);
-
-        let s_add = cs.selector();
-        let s_mul = cs.selector();
-        let s_sub = cs.selector();
-
-        cs.create_gate("add", |cs| {
-            let lhs = cs.query_advice(a_col, Rotation::cur());
-            let rhs = cs.query_advice(b_col, Rotation::cur());
-            let out = cs.query_advice(a_col, Rotation::next());
-            let s_add = cs.query_selector(s_add);
-
-            vec![s_add * (lhs + rhs - out)]
+impl ArithChip {
+    pub fn configure(
+        meta: &mut ConstraintSystem<pallas::Base>,
+        a: Column<Advice>,
+        b: Column<Advice>,
+        c: Column<Advice>,
+    ) -> ArithConfig {
+        let q_add = meta.selector();
+        let q_sub = meta.selector();
+        let q_mul = meta.selector();
+
+        meta.create_gate("Field element addition: c = a + b", |meta| {
+            let q_add = meta.query_selector(q_add);
+            let a = meta.query_advice(a, Rotation::cur());
+            let b = meta.query_advice(b, Rotation::cur());
+            let c = meta.query_advice(c, Rotation::cur());
+
+            Constraints::with_selector(q_add, Some(a + b - c))
         });
 
-        cs.create_gate("mul", |cs| {
-            let lhs = cs.query_advice(a_col, Rotation::cur());
-            let rhs = cs.query_advice(b_col, Rotation::cur());
-            let out = cs.query_advice(a_col, Rotation::next());
-            let s_mul = cs.query_selector(s_mul);
+        meta.create_gate("Field element substitution: c = a - b", |meta| {
+            let q_sub = meta.query_selector(q_sub);
+            let a = meta.query_advice(a, Rotation::cur());
+            let b = meta.query_advice(b, Rotation::cur());
+            let c = meta.query_advice(c, Rotation::cur());
 
-            vec![s_mul * (lhs * rhs - out)]
+            Constraints::with_selector(q_sub, Some(a - b - c))
         });
 
-        cs.create_gate("sub", |cs| {
-            let lhs = cs.query_advice(a_col, Rotation::cur());
-            let rhs = cs.query_advice(b_col, Rotation::cur());
-            let out = cs.query_advice(a_col, Rotation::next());
-            let s_sub = cs.query_selector(s_sub);
+        meta.create_gate("Field element multiplication: c = a * b", |meta| {
+            let q_mul = meta.query_selector(q_mul);
+            let a = meta.query_advice(a, Rotation::cur());
+            let b = meta.query_advice(b, Rotation::cur());
+            let c = meta.query_advice(c, Rotation::cur());
 
-            vec![s_sub * (lhs - rhs - out)]
+            Constraints::with_selector(q_mul, Some(a * b - c))
         });
 
-        ArithmeticChipConfig { a_col, b_col, s_add, s_mul, s_sub }
+        ArithConfig { a, b, c, q_add, q_sub, q_mul }
     }
 
-    pub fn add(
+    pub fn construct(config: ArithConfig) -> Self {
+        Self { config }
+    }
+}
+
+impl ArithInstruction<pallas::Base> for ArithChip {
+    fn add(
         &self,
         mut layouter: impl Layouter<pallas::Base>,
-        a: Variable,
-        b: Variable,
-    ) -> Result<Variable, Error> {
-        let mut out = None;
+        a: &AssignedCell<pallas::Base, pallas::Base>,
+        b: &AssignedCell<pallas::Base, pallas::Base>,
+    ) -> Result<AssignedCell<pallas::Base, pallas::Base>, plonk::Error> {
         layouter.assign_region(
-            || "mul",
+            || "c = a + b",
             |mut region| {
-                self.config.s_add.enable(&mut region, 0)?;
+                self.config.q_add.enable(&mut region, 0)?;
 
-                let lhs = region.assign_advice(
-                    || "lhs",
-                    self.config.a_col,
-                    0,
-                    || Ok(*a.value().ok_or(Error::Synthesis)?),
-                )?;
+                a.copy_advice(|| "copy a", &mut region, self.config.a, 0)?;
+                b.copy_advice(|| "copy b", &mut region, self.config.b, 0)?;
 
-                let rhs = region.assign_advice(
-                    || "rhs",
-                    self.config.b_col,
+                let scalar_val = a.value().zip(b.value()).map(|(a, b)| a + b);
+                region.assign_advice(
+                    || "c",
+                    self.config.c,
                     0,
-                    || Ok(*b.value().ok_or(Error::Synthesis)?),
-                )?;
-
-                region.constrain_equal(a.cell(), lhs.cell())?;
-                region.constrain_equal(b.cell(), rhs.cell())?;
-
-                let value = a.value().and_then(|a| b.value().map(|b| a + b));
-
-                let cell = region.assign_advice(
-                    || "lhs + rhs",
-                    self.config.a_col,
-                    1,
-                    || value.ok_or(Error::Synthesis),
-                )?;
-
-                out = Some(cell);
-                Ok(())
+                    || scalar_val.ok_or(plonk::Error::Synthesis),
+                )
             },
-        )?;
-
-        Ok(out.unwrap())
+        )
     }
 
-    pub fn mul(
+    fn sub(
         &self,
         mut layouter: impl Layouter<pallas::Base>,
-        a: Variable,
-        b: Variable,
-    ) -> Result<Variable, Error> {
-        let mut out = None;
-
+        a: &AssignedCell<pallas::Base, pallas::Base>,
+        b: &AssignedCell<pallas::Base, pallas::Base>,
+    ) -> Result<AssignedCell<pallas::Base, pallas::Base>, plonk::Error> {
         layouter.assign_region(
-            || "mul",
+            || "c = a - b",
             |mut region| {
-                self.config.s_mul.enable(&mut region, 0)?;
+                self.config.q_sub.enable(&mut region, 0)?;
 
-                let lhs = region.assign_advice(
-                    || "lhs",
-                    self.config.a_col,
-                    0,
-                    || Ok(*a.value().ok_or(Error::Synthesis)?),
-                )?;
+                a.copy_advice(|| "copy a", &mut region, self.config.a, 0)?;
+                b.copy_advice(|| "copy b", &mut region, self.config.b, 0)?;
 
-                let rhs = region.assign_advice(
-                    || "rhs",
-                    self.config.b_col,
+                let scalar_val = a.value().zip(b.value()).map(|(a, b)| a - b);
+                region.assign_advice(
+                    || "c",
+                    self.config.c,
                     0,
-                    || Ok(*b.value().ok_or(Error::Synthesis)?),
-                )?;
-
-                region.constrain_equal(a.cell(), lhs.cell())?;
-                region.constrain_equal(b.cell(), rhs.cell())?;
-
-                let value = a.value().and_then(|a| b.value().map(|b| a * b));
-                let cell = region.assign_advice(
-                    || "lhs * rhs",
-                    self.config.a_col,
-                    1,
-                    || value.ok_or(Error::Synthesis),
-                )?;
-
-                out = Some(cell);
-                Ok(())
+                    || scalar_val.ok_or(plonk::Error::Synthesis),
+                )
             },
-        )?;
-
-        Ok(out.unwrap())
+        )
     }
 
-    pub fn sub(
+    fn mul(
         &self,
         mut layouter: impl Layouter<pallas::Base>,
-        a: Variable,
-        b: Variable,
-    ) -> Result<Variable, Error> {
-        let mut out = None;
-
+        a: &AssignedCell<pallas::Base, pallas::Base>,
+        b: &AssignedCell<pallas::Base, pallas::Base>,
+    ) -> Result<AssignedCell<pallas::Base, pallas::Base>, plonk::Error> {
         layouter.assign_region(
-            || "sub",
+            || "c = a * b",
             |mut region| {
-                self.config.s_sub.enable(&mut region, 0)?;
+                self.config.q_mul.enable(&mut region, 0)?;
 
-                let lhs = region.assign_advice(
-                    || "lhs",
-                    self.config.a_col,
-                    0,
-                    || Ok(*a.value().ok_or(Error::Synthesis)?),
-                )?;
+                a.copy_advice(|| "copy a", &mut region, self.config.a, 0)?;
+                b.copy_advice(|| "copy b", &mut region, self.config.b, 0)?;
 
-                let rhs = region.assign_advice(
-                    || "rhs",
-                    self.config.b_col,
+                let scalar_val = a.value().zip(b.value()).map(|(a, b)| a * b);
+                region.assign_advice(
+                    || "c",
+                    self.config.c,
                     0,
-                    || Ok(*b.value().ok_or(Error::Synthesis)?),
-                )?;
-
-                region.constrain_equal(a.cell(), lhs.cell())?;
-                region.constrain_equal(b.cell(), rhs.cell())?;
-
-                let value = a.value().and_then(|a| b.value().map(|b| a - b));
-                let cell = region.assign_advice(
-                    || "lhs * rhs",
-                    self.config.a_col,
-                    1,
-                    || value.ok_or(Error::Synthesis),
-                )?;
-
-                out = Some(cell);
-                Ok(())
+                    || scalar_val.ok_or(plonk::Error::Synthesis),
+                )
             },
-        )?;
-
-        Ok(out.unwrap())
+        )
     }
 }

src/zk/circuit/lead_contract.rs

@@ -31,13 +31,12 @@ use crate::crypto::{
 };
 
 use crate::zk::{
-    arith_chip::{ArithmeticChip, ArithmeticChipConfig},
+    arith_chip::{ArithChip, ArithConfig, ArithInstruction},
     even_bits::{EvenBitsChip, EvenBitsConfig, EvenBitsLookup},
     greater_than::{GreaterThanChip, GreaterThanConfig, GreaterThanInstruction},
 };
 
 use pasta_curves::group::{ff::PrimeField, GroupEncoding};
-//use halo2_proofs::arithmetic::CurveAffine;
 
 const WORD_BITS: u32 = 24;
 
@@ -55,7 +54,7 @@ pub struct LeadConfig {
         SinsemillaConfig<OrchardHashDomains, OrchardCommitDomains, OrchardFixedBases>,
     greaterthan_config: GreaterThanConfig,
     evenbits_config: EvenBitsConfig,
-    arith_config: ArithmeticChipConfig,
+    arith_config: ArithConfig,
 }
 
 impl LeadConfig {
@@ -87,8 +86,8 @@ impl LeadConfig {
         EvenBitsChip::construct(self.evenbits_config.clone())
     }
 
-    fn arith_chip(&self) -> ArithmeticChip {
-        ArithmeticChip::construct(self.arith_config.clone())
+    fn arith_chip(&self) -> ArithChip {
+        ArithChip::construct(self.arith_config.clone())
     }
 }
 
@@ -239,7 +238,7 @@ impl Circuit<pallas::Base> for LeadContract {
             primary,
         );
         let evenbits_config = EvenBitsChip::<pallas::Base, WORD_BITS>::configure(meta);
-        let arith_config = ArithmeticChip::configure(meta);
+        let arith_config = ArithChip::configure(meta, advices[7], advices[8], advices[6]);
 
         LeadConfig {
             primary,
@@ -472,11 +471,10 @@ impl Circuit<pallas::Base> for LeadContract {
         };
          */
         let coin_val = {
-            let coin_val_pt = ar_chip.mul(layouter.namespace(|| ""), coin_pk_y, coin_pk_x)?;
+            let coin_val_pt = ar_chip.mul(layouter.namespace(|| ""), &coin_pk_y, &coin_pk_x)?;
 
-            let coin_val0 =
-                ar_chip.mul(layouter.namespace(|| ""), coin_nonce.clone(), coin_value.clone())?;
-            ar_chip.mul(layouter.namespace(|| ""), coin_val_pt, coin_val0)?
+            let coin_val0 = ar_chip.mul(layouter.namespace(|| ""), &coin_nonce, &coin_value)?;
+            ar_chip.mul(layouter.namespace(|| ""), &coin_val_pt, &coin_val0)?
         };
 
         let (com, _) = {
@@ -545,16 +543,16 @@ impl Circuit<pallas::Base> for LeadContract {
          */
         let coin2_hash0 = ar_chip.mul(
             layouter.namespace(|| ""),
-            coin_pk_commit.inner().x(),
-            coin_pk_commit.inner().y(),
+            &coin_pk_commit.inner().x(),
+            &coin_pk_commit.inner().y(),
         )?;
         let coin2_hash1 = ar_chip.mul(
             layouter.namespace(|| ""),
-            coin2_nonce.inner().x(),
-            coin2_nonce.inner().y(),
+            &coin2_nonce.inner().x(),
+            &coin2_nonce.inner().y(),
         )?;
-        let coin2_hash2 = ar_chip.mul(layouter.namespace(|| ""), coin2_hash0, coin2_hash1)?;
-        let coin2_hash = ar_chip.mul(layouter.namespace(|| ""), coin_value.clone(), coin2_hash2)?;
+        let coin2_hash2 = ar_chip.mul(layouter.namespace(|| ""), &coin2_hash0, &coin2_hash1)?;
+        let coin2_hash = ar_chip.mul(layouter.namespace(|| ""), &coin_value, &coin2_hash2)?;
 
         let (com, _) = {
             let coin_commit_v = ValueCommitV;
@@ -610,8 +608,8 @@ impl Circuit<pallas::Base> for LeadContract {
 
             let res: AssignedCell<Fp, Fp> = ar_chip.mul(
                 layouter.namespace(|| ""),
-                coin_commit_coordinates.x(),
-                coin_commit_coordinates.y(),
+                &coin_commit_coordinates.x(),
+                &coin_commit_coordinates.y(),
             )?;
             res
         };
@@ -632,9 +630,9 @@ impl Circuit<pallas::Base> for LeadContract {
         //TODO (research) this multiplication panics!
         let y_commit_exp = ar_chip.mul(
             layouter.namespace(|| ""),
-            coin_nonce,
+            &coin_nonce,
             //root_sk.clone(), //(fix)
-            one.clone(),
+            &one,
         )?;
 
         let (com, _) = {
@@ -711,8 +709,8 @@ impl Circuit<pallas::Base> for LeadContract {
             config.advices[0],
             Some(pallas::Base::one()), // note! this parameter to be tuned.
         )?;
-        let ord = ar_chip.mul(layouter.namespace(|| ""), scalar, c)?;
-        let target = ar_chip.mul(layouter.namespace(|| "calculate target"), ord, coin_value)?;
+        let ord = ar_chip.mul(layouter.namespace(|| ""), &scalar, &c)?;
+        let target = ar_chip.mul(layouter.namespace(|| "calculate target"), &ord, &coin_value)?;
 
         eb_chip.decompose(layouter.namespace(|| "target range check"), target.clone())?;
         eb_chip.decompose(layouter.namespace(|| "y_commit  range check"), y_commit_base.clone())?;

src/zk/vm.rs

@@ -25,7 +25,7 @@ use log::debug;
 use pasta_curves::{group::Curve, pallas, Fp};
 
 use super::{
-    arith_chip::{ArithmeticChip, ArithmeticChipConfig},
+    arith_chip::{ArithChip, ArithConfig, ArithInstruction},
     even_bits::{EvenBitsChip, EvenBitsConfig, EvenBitsLookup},
     greater_than::{GreaterThanChip, GreaterThanConfig, GreaterThanInstruction},
 };
@@ -50,7 +50,7 @@ pub struct VmConfig {
     sinsemilla_cfg1: SinsemillaConfig<OrchardHashDomains, OrchardCommitDomains, OrchardFixedBases>,
     _sinsemilla_cfg2: SinsemillaConfig<OrchardHashDomains, OrchardCommitDomains, OrchardFixedBases>,
     poseidon_config: PoseidonConfig<pallas::Base, 3, 2>,
-    arith_config: ArithmeticChipConfig,
+    arith_config: ArithConfig,
     evenbits_config: EvenBitsConfig,
     greaterthan_config: GreaterThanConfig,
 }
@@ -90,8 +90,8 @@ impl VmConfig {
         PoseidonChip::construct(self.poseidon_config.clone())
     }
 
-    fn arithmetic_chip(&self) -> ArithmeticChip {
-        ArithmeticChip::construct(self.arith_config.clone())
+    fn arithmetic_chip(&self) -> ArithChip {
+        ArithChip::construct(self.arith_config.clone())
     }
 
     fn evenbits_chip(&self) -> EvenBitsChip<pallas::Base, 24> {
@@ -202,7 +202,7 @@ impl Circuit<pallas::Base> for ZkCircuit {
         );
 
         // Configuration for the Arithmetic chip
-        let arith_config = ArithmeticChip::configure(meta);
+        let arith_config = ArithChip::configure(meta, advices[7], advices[8], advices[6]);
 
         // Configuration for the EvenBits chip
         let evenbits_config = EvenBitsChip::<pallas::Base, 24>::configure(meta);
@@ -548,8 +548,8 @@ impl Circuit<pallas::Base> for ZkCircuit {
                     debug!("Executing `BaseAdd{:?}` opcode", opcode.1);
                     let args = &opcode.1;
 
-                    let lhs = stack[args[0]].clone().into();
-                    let rhs = stack[args[1]].clone().into();
+                    let lhs = &stack[args[0]].clone().into();
+                    let rhs = &stack[args[1]].clone().into();
 
                     let sum = arith_chip.add(layouter.namespace(|| "BaseAdd()"), lhs, rhs)?;
 
@@ -561,8 +561,8 @@ impl Circuit<pallas::Base> for ZkCircuit {
                     debug!("Executing `BaseMul{:?}` opcode", opcode.1);
                     let args = &opcode.1;
 
-                    let lhs = stack[args[0]].clone().into();
-                    let rhs = stack[args[1]].clone().into();
+                    let lhs = &stack[args[0]].clone().into();
+                    let rhs = &stack[args[1]].clone().into();
 
                     let product = arith_chip.mul(layouter.namespace(|| "BaseMul()"), lhs, rhs)?;
 
@@ -574,8 +574,8 @@ impl Circuit<pallas::Base> for ZkCircuit {
                     debug!("Executing `BaseSub{:?}` opcode", opcode.1);
                     let args = &opcode.1;
 
-                    let lhs = stack[args[0]].clone().into();
-                    let rhs = stack[args[1]].clone().into();
+                    let lhs = &stack[args[0]].clone().into();
+                    let rhs = &stack[args[1]].clone().into();
 
                     let difference =
                         arith_chip.sub(layouter.namespace(|| "BaseSub()"), lhs, rhs)?;