From 4b0ea0be155b17a2d0e8364d9ce8a6d7a5184dc3 Mon Sep 17 00:00:00 2001 From: therealyingtong Date: Thu, 3 Jun 2021 09:42:36 +0800 Subject: [PATCH] Add conditional swap chip --- src/circuit/gadget/utilities.rs | 1 + src/circuit/gadget/utilities/cond_swap.rs | 356 ++++++++++++++++++++++ 2 files changed, 357 insertions(+) create mode 100644 src/circuit/gadget/utilities/cond_swap.rs diff --git a/src/circuit/gadget/utilities.rs b/src/circuit/gadget/utilities.rs index a7644726..a7794868 100644 --- a/src/circuit/gadget/utilities.rs +++ b/src/circuit/gadget/utilities.rs @@ -4,6 +4,7 @@ use halo2::{ }; use pasta_curves::arithmetic::FieldExt; +mod cond_swap; mod plonk; /// A variable representing a number. diff --git a/src/circuit/gadget/utilities/cond_swap.rs b/src/circuit/gadget/utilities/cond_swap.rs new file mode 100644 index 00000000..566c6658 --- /dev/null +++ b/src/circuit/gadget/utilities/cond_swap.rs @@ -0,0 +1,356 @@ +use super::{copy, CellValue, UtilitiesInstructions}; +use halo2::{ + circuit::{Cell, Chip, Layouter}, + plonk::{Advice, Column, ConstraintSystem, Error, Expression, Permutation, Selector}, + poly::Rotation, +}; +use pasta_curves::arithmetic::FieldExt; +use std::marker::PhantomData; + +pub trait CondSwapInstructions: UtilitiesInstructions { + /// Variable representing cell with a certain value in the circuit. + type Var; + + /// Variable representing a `swap` boolean flag. + type Swap: From<>::Var>; + + /// Given an input pair (x,y) and a `swap` boolean flag, return + /// (y,x) if `swap` is set, else (x,y) if `swap` is not set. + fn swap( + &self, + layouter: impl Layouter, + pair: ( + >::Var, + >::Var, + ), + swap: Self::Swap, + ) -> Result< + ( + >::Var, + >::Var, + ), + Error, + >; +} + +/// A chip implementing a conditional swap. +#[derive(Clone, Debug)] +pub struct CondSwapChip { + config: CondSwapConfig, + _marker: PhantomData, +} + +impl Chip for CondSwapChip { + type Config = CondSwapConfig; + type Loaded = (); + + fn config(&self) -> &Self::Config { + &self.config + } + + fn loaded(&self) -> &Self::Loaded { + &() + } +} + +#[derive(Clone, Debug)] +pub struct CondSwapConfig { + q_swap: Selector, + x: Column, + y: Column, + x_swapped: Column, + y_swapped: Column, + swap: Column, + perm: Permutation, +} + +/// A variable representing a `swap` boolean flag. +#[derive(Copy, Clone)] +pub struct Swap { + cell: Cell, + value: Option, +} + +impl From> for Swap { + fn from(var: CellValue) -> Self { + let value = var.value.map(|value| { + let zero = value == F::zero(); + let one = value == F::one(); + if zero { + false + } else if one { + true + } else { + panic!("Value must be boolean.") + } + }); + Swap { + cell: var.cell, + value, + } + } +} + +impl UtilitiesInstructions for CondSwapChip { + type Var = CellValue; +} + +impl CondSwapInstructions for CondSwapChip { + type Var = CellValue; + type Swap = Swap; + + fn swap( + &self, + mut layouter: impl Layouter, + pair: ( + >::Var, + >::Var, + ), + swap: Self::Swap, + ) -> Result< + ( + >::Var, + >::Var, + ), + Error, + > { + let config = self.config(); + + layouter.assign_region( + || "swap", + |mut region| { + // Enable `q_swap` selector + config.q_swap.enable(&mut region, 0)?; + + // Copy in `x` value + let x = copy(&mut region, || "copy x", config.x, 0, &pair.0, &config.perm)?; + + // Copy in `y` value + let y = copy(&mut region, || "copy y", config.y, 0, &pair.1, &config.perm)?; + + // Copy in `swap` value + let swap_val = swap.value; + let swap_cell = region.assign_advice( + || "swap", + config.swap, + 0, + || { + swap_val + .map(|swap| F::from_u64(swap as u64)) + .ok_or(Error::SynthesisError) + }, + )?; + region.constrain_equal(&config.perm, swap_cell, swap.cell)?; + + // Conditionally swap x + let x_swapped = { + let x_swapped = + x.value + .zip(y.value) + .zip(swap_val) + .map(|((x, y), swap)| if swap { y } else { x }); + let x_swapped_cell = region.assign_advice( + || "x_swapped", + config.x_swapped, + 0, + || x_swapped.ok_or(Error::SynthesisError), + )?; + CellValue { + cell: x_swapped_cell, + value: x_swapped, + } + }; + + // Conditionally swap y + let y_swapped = { + let y_swapped = + x.value + .zip(y.value) + .zip(swap_val) + .map(|((x, y), swap)| if swap { x } else { y }); + let y_swapped_cell = region.assign_advice( + || "y_swapped", + config.y_swapped, + 0, + || y_swapped.ok_or(Error::SynthesisError), + )?; + CellValue { + cell: y_swapped_cell, + value: y_swapped, + } + }; + + // Return swapped pair + Ok((x_swapped, y_swapped)) + }, + ) + } +} + +impl CondSwapChip { + /// Configures this chip for use in a circuit. + pub fn configure( + meta: &mut ConstraintSystem, + advices: [Column; 5], + perm: Permutation, + ) -> CondSwapConfig { + let q_swap = meta.selector(); + + let config = CondSwapConfig { + q_swap, + x: advices[0], + y: advices[1], + x_swapped: advices[2], + y_swapped: advices[3], + swap: advices[4], + perm, + }; + + // TODO: optimise shape of gate for Merkle path validation + + meta.create_gate("x' = y ⋅ swap + x ⋅ (1-swap)", |meta| { + let q_swap = meta.query_selector(q_swap, Rotation::cur()); + + let x = meta.query_advice(config.x, Rotation::cur()); + let y = meta.query_advice(config.y, Rotation::cur()); + let x_swapped = meta.query_advice(config.x_swapped, Rotation::cur()); + let y_swapped = meta.query_advice(config.y_swapped, Rotation::cur()); + let swap = meta.query_advice(config.swap, Rotation::cur()); + + let one = Expression::Constant(F::one()); + + // x_swapped - y ⋅ swap - x ⋅ (1-swap) = 0 + // This checks that `x_swapped` is equal to `y` when `swap` is set, + // but remains as `x` when `swap` is not set. + let x_check = + x_swapped - y.clone() * swap.clone() - x.clone() * (one.clone() - swap.clone()); + + // y_swapped - x ⋅ swap - y ⋅ (1-swap) = 0 + // This checks that `y_swapped` is equal to `x` when `swap` is set, + // but remains as `y` when `swap` is not set. + let y_check = y_swapped - x * swap.clone() - y * (one.clone() - swap.clone()); + + // Check `swap` is boolean. + let bool_check = swap.clone() * (one - swap); + + [x_check, y_check, bool_check] + .iter() + .map(|poly| q_swap.clone() * poly.clone()) + .collect() + }); + + config + } + + pub fn construct(config: CondSwapConfig) -> Self { + CondSwapChip { + config, + _marker: PhantomData, + } + } +} + +#[cfg(test)] +mod tests { + use super::super::UtilitiesInstructions; + use super::{CondSwapChip, CondSwapConfig, CondSwapInstructions}; + use halo2::{ + circuit::{layouter::SingleChipLayouter, Layouter}, + dev::MockProver, + plonk::{Any, Assignment, Circuit, Column, ConstraintSystem, Error}, + }; + use pasta_curves::{arithmetic::FieldExt, pallas::Base}; + + #[test] + fn cond_swap() { + struct MyCircuit { + x: Option, + y: Option, + swap: Option, + } + + impl Circuit for MyCircuit { + type Config = CondSwapConfig; + + fn configure(meta: &mut ConstraintSystem) -> Self::Config { + let advices = [ + meta.advice_column(), + meta.advice_column(), + meta.advice_column(), + meta.advice_column(), + meta.advice_column(), + ]; + + let perm = meta.permutation( + &advices + .iter() + .map(|advice| (*advice).into()) + .collect::>>(), + ); + + CondSwapChip::::configure(meta, advices, perm) + } + + fn synthesize( + &self, + cs: &mut impl Assignment, + config: Self::Config, + ) -> Result<(), Error> { + let mut layouter = SingleChipLayouter::new(cs)?; + let chip = CondSwapChip::::construct(config.clone()); + + // Load the pair and the swap flag into the circuit. + let x = chip.load_private(layouter.namespace(|| "x"), config.x, self.x)?; + let y = chip.load_private(layouter.namespace(|| "y"), config.y, self.y)?; + let swap = + chip.load_private(layouter.namespace(|| "swap"), config.swap, self.swap)?; + + // Return the swapped pair. + let swapped_pair = + chip.swap(layouter.namespace(|| "swap"), (x, y).into(), swap.into())?; + + if let Some(swap) = self.swap { + if swap == F::one() { + // Check that `x` and `y` have been swapped + assert_eq!(swapped_pair.0.value.unwrap(), y.value.unwrap()); + assert_eq!(swapped_pair.1.value.unwrap(), x.value.unwrap()); + } else { + // Check that `x` and `y` have not been swapped + assert_eq!(swapped_pair.0.value.unwrap(), x.value.unwrap()); + assert_eq!(swapped_pair.1.value.unwrap(), y.value.unwrap()); + } + } + + Ok(()) + } + } + + // Test swap case + { + let circuit: MyCircuit = MyCircuit { + x: Some(Base::rand()), + y: Some(Base::rand()), + swap: Some(Base::one()), + }; + let prover = match MockProver::::run(3, &circuit, vec![]) { + Ok(prover) => prover, + Err(e) => panic!("{:?}", e), + }; + assert_eq!(prover.verify(), Ok(())); + } + + // Test non-swap case + { + let circuit: MyCircuit = MyCircuit { + x: Some(Base::rand()), + y: Some(Base::rand()), + swap: Some(Base::zero()), + }; + let prover = match MockProver::::run(3, &circuit, vec![]) { + Ok(prover) => prover, + Err(e) => panic!("{:?}", e), + }; + assert_eq!(prover.verify(), Ok(())); + } + } +}