From bef17692002c75e7a615f763eef67f920f5556c2 Mon Sep 17 00:00:00 2001 From: Supragya Raj Date: Sun, 24 Mar 2024 16:31:11 +0530 Subject: [PATCH 1/2] add: basic custom gate --- Cargo.toml | 3 +- zk-plonky2-permutation-circuit/Cargo.toml | 9 ++ .../src/custom_gates/mod.rs | 1 + .../src/custom_gates/switch.rs | 97 +++++++++++++++++++ zk-plonky2-permutation-circuit/src/lib.rs | 89 +++++++++++++++++ 5 files changed, 198 insertions(+), 1 deletion(-) create mode 100644 zk-plonky2-permutation-circuit/Cargo.toml create mode 100644 zk-plonky2-permutation-circuit/src/custom_gates/mod.rs create mode 100644 zk-plonky2-permutation-circuit/src/custom_gates/switch.rs create mode 100644 zk-plonky2-permutation-circuit/src/lib.rs diff --git a/Cargo.toml b/Cargo.toml index 6297a23..00342b8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,8 @@ [workspace] members = [ "polynomial", - "univariate-polynomial-iop-zerotest" + "univariate-polynomial-iop-zerotest", + "zk-plonky2-permutation-circuit" ] resolver = "2" diff --git a/zk-plonky2-permutation-circuit/Cargo.toml b/zk-plonky2-permutation-circuit/Cargo.toml new file mode 100644 index 0000000..09bf4b9 --- /dev/null +++ b/zk-plonky2-permutation-circuit/Cargo.toml @@ -0,0 +1,9 @@ +[package] +edition = "2021" +name = "zk-plonky2-permutation-circuit" +version = "0.1.0" + +[dependencies] +anyhow = "1.0.81" +plonky2 = "0.2.0" +plonky2_field = "0.2.0" diff --git a/zk-plonky2-permutation-circuit/src/custom_gates/mod.rs b/zk-plonky2-permutation-circuit/src/custom_gates/mod.rs new file mode 100644 index 0000000..5a2a8f4 --- /dev/null +++ b/zk-plonky2-permutation-circuit/src/custom_gates/mod.rs @@ -0,0 +1 @@ +pub mod switch; diff --git a/zk-plonky2-permutation-circuit/src/custom_gates/switch.rs b/zk-plonky2-permutation-circuit/src/custom_gates/switch.rs new file mode 100644 index 0000000..b58cd8c --- /dev/null +++ b/zk-plonky2-permutation-circuit/src/custom_gates/switch.rs @@ -0,0 +1,97 @@ +use std::marker::PhantomData; + +use plonky2::{gates::gate::Gate, hash::hash_types::RichField, iop::generator::WitnessGeneratorRef, util::serialization::Write}; +use plonky2_field::extension::Extendable; + + +/// A gate for conditionally swapping input values based on a boolean. +#[derive(Copy, Clone, Debug)] +pub struct SwitchGate, const D: usize> { + _phantom: PhantomData +} + +impl, const D: usize> Gate for SwitchGate { + fn id(&self) -> String { + format!("{self:?}") + } + + fn serialize(&self, dst: &mut Vec, common_data: &plonky2::plonk::circuit_data::CommonCircuitData) -> plonky2::util::serialization::IoResult<()> { + dst.write_bool(false) // TODO: remove + } + + fn deserialize(src: &mut plonky2::util::serialization::Buffer, common_data: &plonky2::plonk::circuit_data::CommonCircuitData) -> plonky2::util::serialization::IoResult + where + Self: Sized { + Ok(Self) + } + + fn eval_unfiltered(&self, vars: plonky2::plonk::vars::EvaluationVars) -> Vec<::Extension> { + let mut constraints = Vec::with_capacity(4); + let switch_bool = vars.local_wires[0]; + let not_switch = F::Extension::ONE - switch_bool; + + let input_1 = vars.local_wires[1]; + let input_2 = vars.local_wires[2]; + let output_1 = vars.local_wires[3]; + let output_2 = vars.local_wires[4]; + + constraints.push(not_switch * (output_1 - input_1)); + constraints.push(not_switch * (output_2 - input_2)); + + constraints.push(switch_bool * (output_2 - input_1)); + constraints.push(switch_bool * (output_1 - input_2)); + constraints + } + + fn eval_unfiltered_circuit( + &self, + builder: &mut plonky2::plonk::circuit_builder::CircuitBuilder, + vars: plonky2::plonk::vars::EvaluationTargets, + ) -> Vec> { + let mut constraints = Vec::with_capacity(4); + + let one = builder.one_extension(); + + let switch_bool = vars.local_wires[0]; + let not_switch = builder.sub_extension(one, switch_bool); + + let input_1 = vars.local_wires[1]; + let input_2 = vars.local_wires[2]; + let output_1 = vars.local_wires[3]; + let output_2 = vars.local_wires[4]; + + constraints.push(builder.mul_extension(not_switch, builder.sub_extension(input_1, output_1))); + constraints.push(builder.mul_extension(not_switch, builder.sub_extension(input_2, output_2))); + + constraints.push(builder.mul_extension(switch_bool, builder.sub_extension(input_1, output_2))); + constraints.push(builder.mul_extension(switch_bool, builder.sub_extension(input_2, output_1))); + + constraints + } + + fn generators(&self, row: usize, local_constants: &[F]) -> Vec> { + let g: Box> = Box::new(SwitchGenerator::{ + row, + gate: *self, + copy: c, + }); + g + } + + + fn num_wires(&self) -> usize { + 5 + } + + fn num_constants(&self) -> usize { + 0 + } + + fn degree(&self) -> usize { + 2 + } + + fn num_constraints(&self) -> usize { + 4 + } +} diff --git a/zk-plonky2-permutation-circuit/src/lib.rs b/zk-plonky2-permutation-circuit/src/lib.rs new file mode 100644 index 0000000..bfc99a3 --- /dev/null +++ b/zk-plonky2-permutation-circuit/src/lib.rs @@ -0,0 +1,89 @@ +mod custom_gates; + +use plonky2::hash::hash_types::RichField; +use plonky2_field::extension::Extendable; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use plonky2::iop::target::Target; + +// Inspired by https://github.com/0xPolygonZero/plonky2-waksman/blob/main/src/permutation.rs + +/// Assert that two set of targets are permutation of each other +pub fn assert_permutation_circuit, const D: usize>( + builder: &mut CircuitBuilder, + a: Vec, + b: Vec, +) { + assert_eq!( + a.len(), + b.len(), + "Permutation must have same number of inputs and outputs" + ); + + match a.len() { + // Two empty lists are permutations of one another, trivially. + 0 => (), + // Two singleton lists are permutations of one another as long as their items are equal. + 1 => { + builder.connect(a[0], b[0]); + } + 2 => assert_permutation_2x2_circuit( + builder, + a[0], + a[1], + b[0], + b[1], + ), + // For larger lists, we recursively use two smaller permutation networks. + _ => assert_permutation_helper_circuit(builder, a, b), + } +} + +/// Assert that [a1, a2] is a permutation of [b1, b2]. +fn assert_permutation_2x2_circuit, const D: usize>( + builder: &mut CircuitBuilder, + a1: Target, + a2: Target, + b1: Target, + b2: Target, +) { + let (_switch, gate_out1, gate_out2) = create_switch_circuit(builder, a1, a2); + for e in 0..chunk_size { + builder.connect(b1[e], gate_out1[e]); + builder.connect(b2[e], gate_out2[e]); + } +} + + +/// Given two input wire chunks, add a new switch to the circuit (by adding one copy to a switch +/// gate). Returns the wire for the switch boolean, and the two output wire chunks. +fn create_switch_circuit, const D: usize>( + builder: &mut CircuitBuilder, + a1: Target, + a2: Target, +) -> (Target, Target, Target) { + + let chunk_size = a1.len(); + + let gate = SwitchGate::new_from_config(&builder.config, chunk_size); + let params = vec![F::from_canonical_usize(chunk_size)]; + let (row, next_copy) = builder.find_slot(gate, ¶ms, &[]); + + let mut c = Vec::new(); + let mut d = Vec::new(); + for e in 0..chunk_size { + builder.connect( + a1[e], + Target::wire(row, gate.wire_first_input(next_copy, e)), + ); + builder.connect( + a2[e], + Target::wire(row, gate.wire_second_input(next_copy, e)), + ); + c.push(Target::wire(row, gate.wire_first_output(next_copy, e))); + d.push(Target::wire(row, gate.wire_second_output(next_copy, e))); + } + + let switch = Target::wire(row, gate.wire_switch_bool(next_copy)); + + (switch, c, d) +} From c659c89b7dcfe8d5a1c52d71d36e1c2b1ee925f0 Mon Sep 17 00:00:00 2001 From: Supragya Raj Date: Sun, 24 Mar 2024 19:31:03 +0530 Subject: [PATCH 2/2] add: basic switch circuit --- zk-plonky2-permutation-circuit/Cargo.toml | 1 + .../src/custom_gates/switch.rs | 235 +++++++++++++++--- zk-plonky2-permutation-circuit/src/lib.rs | 59 ++--- 3 files changed, 222 insertions(+), 73 deletions(-) diff --git a/zk-plonky2-permutation-circuit/Cargo.toml b/zk-plonky2-permutation-circuit/Cargo.toml index 09bf4b9..3f30d95 100644 --- a/zk-plonky2-permutation-circuit/Cargo.toml +++ b/zk-plonky2-permutation-circuit/Cargo.toml @@ -7,3 +7,4 @@ version = "0.1.0" anyhow = "1.0.81" plonky2 = "0.2.0" plonky2_field = "0.2.0" +array_tool = "1.0.3" diff --git a/zk-plonky2-permutation-circuit/src/custom_gates/switch.rs b/zk-plonky2-permutation-circuit/src/custom_gates/switch.rs index b58cd8c..08a538f 100644 --- a/zk-plonky2-permutation-circuit/src/custom_gates/switch.rs +++ b/zk-plonky2-permutation-circuit/src/custom_gates/switch.rs @@ -1,39 +1,76 @@ +use array_tool::vec::Union; use std::marker::PhantomData; -use plonky2::{gates::gate::Gate, hash::hash_types::RichField, iop::generator::WitnessGeneratorRef, util::serialization::Write}; +use plonky2::{ + gates::gate::Gate, + hash::hash_types::RichField, + iop::{generator::{GeneratedValues, WitnessGenerator, WitnessGeneratorRef}, target::Target, wire::Wire, witness::{Witness, WitnessWrite, PartitionWitness}}, + plonk::circuit_data::CommonCircuitData, + util::serialization::{Buffer, IoResult, Read, Write}, +}; +use plonky2_field::types::Field; use plonky2_field::extension::Extendable; - /// A gate for conditionally swapping input values based on a boolean. #[derive(Copy, Clone, Debug)] -pub struct SwitchGate, const D: usize> { - _phantom: PhantomData +pub struct SwitchGate, const D: usize> { + _phantom: PhantomData, +} + +impl, const D: usize> SwitchGate { + pub fn new() -> Self { + Self{_phantom: PhantomData} + } + + pub fn wire_switch_bool() -> usize { + 0 + } + + pub fn wire_input_1() -> usize { + 1 + } + + pub fn wire_input_2() -> usize { + 2 + } + + pub fn wire_output_1() -> usize { + 3 + } + + pub fn wire_output_2() -> usize { + 4 + } } -impl, const D: usize> Gate for SwitchGate { +impl, const D: usize> Gate for SwitchGate { fn id(&self) -> String { format!("{self:?}") } - fn serialize(&self, dst: &mut Vec, common_data: &plonky2::plonk::circuit_data::CommonCircuitData) -> plonky2::util::serialization::IoResult<()> { + fn serialize(&self, dst: &mut Vec, common_data: &CommonCircuitData) -> IoResult<()> { dst.write_bool(false) // TODO: remove } - fn deserialize(src: &mut plonky2::util::serialization::Buffer, common_data: &plonky2::plonk::circuit_data::CommonCircuitData) -> plonky2::util::serialization::IoResult - where - Self: Sized { - Ok(Self) + fn deserialize(src: &mut Buffer, common_data: &CommonCircuitData) -> IoResult + where + Self: Sized, + { + Ok(Self{_phantom: PhantomData}) } - fn eval_unfiltered(&self, vars: plonky2::plonk::vars::EvaluationVars) -> Vec<::Extension> { + fn eval_unfiltered( + &self, + vars: plonky2::plonk::vars::EvaluationVars, + ) -> Vec<>::Extension> { let mut constraints = Vec::with_capacity(4); - let switch_bool = vars.local_wires[0]; + let switch_bool = vars.local_wires[Self::wire_switch_bool()]; let not_switch = F::Extension::ONE - switch_bool; - let input_1 = vars.local_wires[1]; - let input_2 = vars.local_wires[2]; - let output_1 = vars.local_wires[3]; - let output_2 = vars.local_wires[4]; + let input_1 = vars.local_wires[Self::wire_input_1()]; + let input_2 = vars.local_wires[Self::wire_input_2()]; + let output_1 = vars.local_wires[Self::wire_output_1()]; + let output_2 = vars.local_wires[Self::wire_output_2()]; constraints.push(not_switch * (output_1 - input_1)); constraints.push(not_switch * (output_2 - input_2)); @@ -44,41 +81,44 @@ impl, const D: usize> Gate for SwitchGate, - vars: plonky2::plonk::vars::EvaluationTargets, - ) -> Vec> { + &self, + builder: &mut plonky2::plonk::circuit_builder::CircuitBuilder, + vars: plonky2::plonk::vars::EvaluationTargets, + ) -> Vec> { let mut constraints = Vec::with_capacity(4); let one = builder.one_extension(); - let switch_bool = vars.local_wires[0]; + let switch_bool = vars.local_wires[Self::wire_switch_bool()]; let not_switch = builder.sub_extension(one, switch_bool); - let input_1 = vars.local_wires[1]; - let input_2 = vars.local_wires[2]; - let output_1 = vars.local_wires[3]; - let output_2 = vars.local_wires[4]; + let input_1 = vars.local_wires[Self::wire_input_1()]; + let input_2 = vars.local_wires[Self::wire_input_2()]; + let output_1 = vars.local_wires[Self::wire_output_1()]; + let output_2 = vars.local_wires[Self::wire_output_2()]; - constraints.push(builder.mul_extension(not_switch, builder.sub_extension(input_1, output_1))); - constraints.push(builder.mul_extension(not_switch, builder.sub_extension(input_2, output_2))); + constraints + .push(builder.mul_extension(not_switch, builder.sub_extension(input_1, output_1))); + constraints + .push(builder.mul_extension(not_switch, builder.sub_extension(input_2, output_2))); - constraints.push(builder.mul_extension(switch_bool, builder.sub_extension(input_1, output_2))); - constraints.push(builder.mul_extension(switch_bool, builder.sub_extension(input_2, output_1))); + constraints + .push(builder.mul_extension(switch_bool, builder.sub_extension(input_1, output_2))); + constraints + .push(builder.mul_extension(switch_bool, builder.sub_extension(input_2, output_1))); constraints } fn generators(&self, row: usize, local_constants: &[F]) -> Vec> { - let g: Box> = Box::new(SwitchGenerator::{ + // unimplemented!() + let g = Box::new(SwitchGenerator:: { row, gate: *self, - copy: c, }); - g + vec![g] } - fn num_wires(&self) -> usize { 5 } @@ -95,3 +135,132 @@ impl, const D: usize> Gate for SwitchGate, const D: usize> { + row: usize, + gate: SwitchGate, +} + +impl, const D: usize> SwitchGenerator { + /// List of wire targets for inputs and outputs + fn dependencies_inputs_outputs(&self) -> Vec { + let local_target = |column| Target::wire(self.row, column); + + let mut deps = Vec::new(); + + deps.push(local_target(SwitchGate::wire_first_input())); + deps.push(local_target(SwitchGate::wire_second_input())); + deps.push(local_target(SwitchGate::wire_first_output())); + deps.push(local_target(SwitchGate::wire_second_output())); + + deps + } + + /// List of wire targets for inputs and switching boolean + fn dependencies_switch_inputs(&self) -> Vec { + let local_target = |column| Target::wire(self.row, column); + + let mut deps = Vec::new(); + + deps.push(local_target(SwitchGate::wire_first_input())); + deps.push(local_target(SwitchGate::wire_second_input())); + deps.push(local_target(SwitchGate::wire_switch_bool())); + + deps + } + + /// Run when all input and output wires are present + fn set_switch_wire(&self, witness: &PartitionWitness, out_buffer: &mut GeneratedValues) { + let get_local_wire = |column| { + witness.get_wire(Wire { + row: self.row, + column, + }) + }; + let switch_bool_wire = Wire { + row: self.row, + column: SwitchGate::wire_switch_bool(), + }; + + let mut input_1 = get_local_wire(SwitchGate::wire_input_1()); + let mut input_2 = get_local_wire(SwitchGate::wire_input_2()); + let mut output_1 = get_local_wire(SwitchGate::wire_output_1()); + let mut output_2 = get_local_wire(SwitchGate::wire_output_2()); + + if input_1 == output_1 && input_2 == output_2 { + out_buffer.set_wire(switch_bool_wire, F::ZERO); + } else if input_1 == output_2 && input_2 == output_1 { + out_buffer.set_wire(switch_bool_wire, F::ONE); + } else { + panic!("No permutation from given inputs to given outputs"); + } + } + + /// Run when only inputs and switching boolean is available + fn set_output_wires(&self, witness: &PartitionWitness, out_buffer: &mut GeneratedValues) { + let get_local_wire = |column| { + witness.get_wire(Wire { + row: self.row, + column, + }) + }; + let switch_bool_wire = Wire { + row: self.row, + column: SwitchGate::wire_switch_bool(), + }; + + let mut input_1 = get_local_wire(SwitchGate::wire_input_1()); + let mut input_2 = get_local_wire(SwitchGate::wire_input_2()); + let mut output_1 = get_local_wire(SwitchGate::wire_output_1()); + let mut output_2 = get_local_wire(SwitchGate::wire_output_2()); + + let (expected_output_1, expected_output_2) = if switch_bool_wire == F::ZERO { + (input_1, input_2) + } else if switch_bool_wire == F::ONE { + (input_2, input_1) + } else { + panic!("Invalid switch bool value"); + }; + + out_buffer.set_wire(output_1, expected_output_1); + out_buffer.set_wire(output_2, expected_output_2); + } +} + +impl, const D: usize> WitnessGenerator for SwitchGenerator { + fn id(&self) -> String { + format!("{self:?}") + } + + fn serialize(&self, dst: &mut Vec, common_data: &CommonCircuitData) -> IoResult<()> { + dst.write_usize(self.row) + } + + fn deserialize(src: &mut Buffer, common_data: &CommonCircuitData) -> IoResult + where + Self: Sized { + Ok(Self{row: src.read_bool().unwrap(), gate: SwitchGate::new()}) + } + /// Register the different columns to watch + fn watch_list(&self) -> Vec { + self.dependencies_inputs_outputs() + .union(self.dependencies_switch_inputs()) + } + + /// Figure out which columns change and set the remaining + /// Can work in two modes: + /// 1. If input and switch wires are pre-populated + /// 2. If input and output wires are pre-populated + fn run(&self, witness: &PartitionWitness, out_buffer: &mut GeneratedValues) -> bool { + if witness.contains_all(&self.dependencies_switch_inputs()) { + self.set_output_wires(witness, out_buffer); + true + } else if witness.contains_all(&self.dependencies_inputs_outputs()) { + self.set_switch_wire(witness, out_buffer); + true + } else { + false + } + } +} diff --git a/zk-plonky2-permutation-circuit/src/lib.rs b/zk-plonky2-permutation-circuit/src/lib.rs index bfc99a3..bacdca9 100644 --- a/zk-plonky2-permutation-circuit/src/lib.rs +++ b/zk-plonky2-permutation-circuit/src/lib.rs @@ -1,9 +1,10 @@ mod custom_gates; +use custom_gates::switch::SwitchGate; use plonky2::hash::hash_types::RichField; -use plonky2_field::extension::Extendable; -use plonky2::plonk::circuit_builder::CircuitBuilder; use plonky2::iop::target::Target; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use plonky2_field::extension::Extendable; // Inspired by https://github.com/0xPolygonZero/plonky2-waksman/blob/main/src/permutation.rs @@ -18,7 +19,7 @@ pub fn assert_permutation_circuit, const D: usize>( b.len(), "Permutation must have same number of inputs and outputs" ); - + match a.len() { // Two empty lists are permutations of one another, trivially. 0 => (), @@ -26,15 +27,9 @@ pub fn assert_permutation_circuit, const D: usize>( 1 => { builder.connect(a[0], b[0]); } - 2 => assert_permutation_2x2_circuit( - builder, - a[0], - a[1], - b[0], - b[1], - ), + 2 => assert_permutation_2x2_circuit(builder, a[0], a[1], b[0], b[1]), // For larger lists, we recursively use two smaller permutation networks. - _ => assert_permutation_helper_circuit(builder, a, b), + _ => unimplemented!(), // assert_permutation_helper_circuit(builder, a, b), } } @@ -46,14 +41,12 @@ fn assert_permutation_2x2_circuit, const D: usize>( b1: Target, b2: Target, ) { - let (_switch, gate_out1, gate_out2) = create_switch_circuit(builder, a1, a2); - for e in 0..chunk_size { - builder.connect(b1[e], gate_out1[e]); - builder.connect(b2[e], gate_out2[e]); - } + let (_switch, out_1, out_2) = create_switch_circuit(builder, a1, a2); + // Add constraints + builder.connect(b1, out_1); + builder.connect(b2, out_2); } - /// Given two input wire chunks, add a new switch to the circuit (by adding one copy to a switch /// gate). Returns the wire for the switch boolean, and the two output wire chunks. fn create_switch_circuit, const D: usize>( @@ -61,29 +54,15 @@ fn create_switch_circuit, const D: usize>( a1: Target, a2: Target, ) -> (Target, Target, Target) { + let gate = SwitchGate::new(); + let (row, _next_copy) = builder.find_slot(gate, &vec![], &[]); - let chunk_size = a1.len(); - - let gate = SwitchGate::new_from_config(&builder.config, chunk_size); - let params = vec![F::from_canonical_usize(chunk_size)]; - let (row, next_copy) = builder.find_slot(gate, ¶ms, &[]); - - let mut c = Vec::new(); - let mut d = Vec::new(); - for e in 0..chunk_size { - builder.connect( - a1[e], - Target::wire(row, gate.wire_first_input(next_copy, e)), - ); - builder.connect( - a2[e], - Target::wire(row, gate.wire_second_input(next_copy, e)), - ); - c.push(Target::wire(row, gate.wire_first_output(next_copy, e))); - d.push(Target::wire(row, gate.wire_second_output(next_copy, e))); - } - - let switch = Target::wire(row, gate.wire_switch_bool(next_copy)); + builder.connect(a1, Target::wire(row, SwitchGate::wire_input_1())); + builder.connect(a2, Target::wire(row, SwitchGate::wire_input_2())); - (switch, c, d) + ( + Target::wire(row, SwitchGate::wire_switch_bool()), + Target::wire(row, SwitchGate::wire_output_1()), + Target::wire(row, SwitchGate::wire_output_2()), + ) }