diff --git a/src/gates/comparison.rs b/src/gates/comparison.rs new file mode 100644 index 00000000..93454ebc --- /dev/null +++ b/src/gates/comparison.rs @@ -0,0 +1,480 @@ +use std::marker::PhantomData; + +use crate::field::extension_field::target::ExtensionTarget; +use crate::field::extension_field::Extendable; +use crate::field::field_types::{Field, PrimeField, RichField}; +use crate::gates::gate::Gate; +use crate::iop::generator::{GeneratedValues, SimpleGenerator, WitnessGenerator}; +use crate::iop::target::Target; +use crate::iop::wire::Wire; +use crate::iop::witness::{PartitionWitness, Witness}; +use crate::plonk::circuit_builder::CircuitBuilder; +use crate::plonk::circuit_data::CircuitConfig; +use crate::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase}; +use crate::util::{ceil_div_usize, log2_ceil}; + +/// A gate for checking that one value is smaller than another. +#[derive(Clone, Debug)] +pub(crate) struct ComparisonGate, const D: usize> { + pub(crate) chunk_bits: usize, + pub(crate) num_copies: usize, + _phantom: PhantomData, +} + +impl, const D: usize> ComparisonGate { + pub fn new(num_copies: usize, chunk_bits: usize) -> Self { + Self { + chunk_bits, + num_copies, + _phantom: PhantomData, + } + } + + pub fn field_bits() -> usize { + log2_ceil(F::ORDER) + } + + pub fn num_chunks(&self) -> usize { + ceil_div_usize(Self::field_bits(), self.chunk_bits) + } + + pub fn new_from_config(config: CircuitConfig, chunk_bits: usize) -> Self { + let num_copies = Self::max_num_copies(config.num_routed_wires, chunk_bits); + Self::new(num_copies, chunk_bits) + } + + pub fn max_num_copies(num_routed_wires: usize, chunk_bits: usize) -> usize { + let num_chunks = ceil_div_usize(Self::field_bits(), chunk_bits); + let wires_per_copy = 4 + chunk_bits + 4 * num_chunks; + num_routed_wires / wires_per_copy + } + + pub fn wire_first_input(&self, copy: usize) -> usize { + debug_assert!(copy < self.num_copies); + copy * (4 + self.chunk_bits + 4 * self.num_chunks()) + } + + pub fn wire_second_input(&self, copy: usize) -> usize { + debug_assert!(copy < self.num_copies); + copy * (4 + self.chunk_bits + 4 * self.num_chunks()) + 1 + } + + pub fn wire_z_val(&self, copy: usize) -> usize { + copy * (4 + self.chunk_bits + 4 * self.num_chunks()) + 3 + } + + pub fn wire_z_bit(&self, copy: usize, bit_index: usize) -> usize { + debug_assert!(bit_index < self.chunk_bits + 1); + copy * (4 + self.chunk_bits + 4 * self.num_chunks()) + 4 + bit_index + } + + pub fn wire_first_chunk_val(&self, copy: usize, chunk: usize) -> usize { + debug_assert!(copy < self.num_copies); + debug_assert!(chunk < self.num_chunks()); + copy * (4 + self.chunk_bits + 4 * self.num_chunks()) + 4 + self.chunk_bits + chunk + } + + pub fn wire_second_chunk_val(&self, copy: usize, chunk: usize) -> usize { + debug_assert!(copy < self.num_copies); + debug_assert!(chunk < self.num_chunks()); + copy * (4 + self.chunk_bits + 4 * self.num_chunks()) + + 4 + + self.chunk_bits + + self.num_chunks() + + chunk + } + + pub fn wire_equality_dummy(&self, copy: usize, chunk: usize) -> usize { + debug_assert!(copy < self.num_copies); + debug_assert!(chunk < self.num_chunks()); + copy * (4 + self.chunk_bits + 4 * self.num_chunks()) + + 4 + + self.chunk_bits + + 2 * self.num_chunks() + + chunk + } + + pub fn wire_chunks_equal(&self, copy: usize, chunk: usize) -> usize { + debug_assert!(copy < self.num_copies); + debug_assert!(chunk < self.num_chunks()); + copy * (4 + self.chunk_bits + 4 * self.num_chunks()) + + 4 + + self.chunk_bits + + 3 * self.num_chunks() + + chunk + } +} + +impl, const D: usize> Gate for ComparisonGate { + fn id(&self) -> String { + format!("{:?}", self, D) + } + + fn eval_unfiltered(&self, vars: EvaluationVars) -> Vec { + let mut constraints = Vec::with_capacity(self.num_constraints()); + + for c in 0..self.num_copies { + let first_input = vars.local_wires[self.wire_first_input(c)]; + let second_input = vars.local_wires[self.wire_second_input(c)]; + + // Get chunks and assert that they match + let first_chunks: Vec = (0..self.num_chunks()) + .map(|i| vars.local_wires[self.wire_first_chunk_val(c, i)]) + .collect(); + let second_chunks: Vec = (0..self.num_chunks()) + .map(|i| vars.local_wires[self.wire_second_chunk_val(c, i)]) + .collect(); + + let chunk_base_powers = (0..self.chunk_bits) + .map(|i| F::TWO.exp_u64(i * self.chunk_bits as u64)) + .collect(); + + let first_chunks_combined = first_chunks + .iter() + .zip(chunk_base_powers.iter()) + .map(|(b, x)| b * x) + .fold(F::ZERO, |a, b| a + b); + let second_chunks_combined = second_chunks + .iter() + .zip(chunk_base_powers.iter()) + .map(|(b, x)| b * x) + .fold(F::ZERO, |a, b| a + b); + + constraints.push(first_chunks_combined - first_input); + constraints.push(second_chunks_combined - second_input); + + // Get bits to assert they match the chosen chunk. + let powers_of_two: Vec = (0..self.chunk_bits) + .map(|i| F::TWO.exp_u64(i as u64)) + .collect(); + + let mut most_significant_diff = + first_chunks[self.num_chunks() - 1] - second_chunks[self.num_chunks() - 1]; + + // Find the chosen chunk. + for i in (0..self.num_chunks()).rev() { + let difference = first_chunks[i] - second_chunks[i]; + let equality_dummy = vars.local_wires[self.wire_equality_dummy(c, i)]; + let chunks_equal = vars.local_wires[self.wires_chunks_equal(c, i)]; + + // Two constraints identifying index. + constraints.push(difference * equality_dummy - (F::Extension::ONE - chunks_equal)); + constraints.push(chunks_equal * difference); + + let this_diff = first_chunks[i] - second_chunks[i]; + most_significant_diff = chunks_equal * most_significant_diff + + (F::Extension::ONE - chunks_equal) * this_diff; + } + + constraints.push(first_bits_combined - most_significant_diff[0]); + constraints.push(second_bits_combined - most_significant_diff[1]); + + let z_bits: Vec = (0..self.chunk_size + 1) + .map(|i| vars.local_wires[self.wire_z_bit(c, i)]) + .collect(); + + let powers_of_two: Vec = (0..self.chunk_bits + 1) + .map(|i| F::TWO.exp_u64(i as u64)) + .collect(); + let z_bits_combined = z_bits + .iter() + .zip(powers_of_two.iter()) + .map(|(b, x)| b * x) + .fold(F::ZERO, |a, b| a + b); + + let two_n = F::TWO.exp_u64(self.chunk_bits); + let (x, y) = most_significant_diff; + constraints.push(z_bits_combined - (two_n + x - y)); + + constraints.push(z_bits[self.chunk_bits - 1]); + } + + constraints + } + + fn eval_unfiltered_base(&self, vars: EvaluationVarsBase) -> Vec { + todo!() + } + + fn eval_unfiltered_recursively( + &self, + builder: &mut CircuitBuilder, + vars: EvaluationTargets, + ) -> Vec> { + todo!() + } + + fn generators( + &self, + gate_index: usize, + _local_constants: &[F], + ) -> Vec>> { + (0..self.num_copies) + .map(|c| { + let g: Box> = Box::new(ComparisonGenerator:: { + gate_index, + gate: self.clone(), + copy: c, + }); + g + }) + .collect() + } + + fn num_wires(&self) -> usize { + self.wire_switch_bool(self.num_copies - 1) + 1 + } + + fn num_constants(&self) -> usize { + 0 + } + + fn degree(&self) -> usize { + 2 + } + + fn num_constraints(&self) -> usize { + 4 * self.num_copies * self.chunk_bits + } +} + +#[derive(Debug)] +struct ComparisonGenerator, const D: usize> { + gate_index: usize, + gate: ComparisonGate, + copy: usize, +} + +impl, const D: usize> SimpleGenerator + for ComparisonGenerator +{ + fn dependencies(&self) -> Vec { + let local_target = |input| Target::wire(self.gate_index, input); + + let mut deps = Vec::new(); + deps.push(local_target(self.gate.wire_first_input(self.copy))); + deps.push(local_target(self.gate.wire_second_input(self.copy))); + deps + } + + fn run_once(&self, witness: &PartitionWitness, out_buffer: &mut GeneratedValues) { + let local_wire = |input| Wire { + gate: self.gate_index, + input, + }; + + let get_local_wire = |input| witness.get_wire(local_wire(input)); + + let first_input = get_local_wire(self.gate.wire_first_input(self.copy)); + let second_input = get_local_wire(self.gate.wire_second_input(self.copy)); + + let field_bits = log2_ceil(F::ORDER); + let first_input_u64 = first_input.to_canonical_u64(); + let second_input_u64 = second_input.to_canonical_u64(); + + let first_input_bits: Vec = (0..field_bits) + .scan(first_input_u64, |acc, _| { + let tmp = *acc % 2; + *acc /= 2; + Some(F::from_canonical_u64(tmp)) + }) + .collect(); + let second_input_bits: Vec = (0..field_bits) + .scan(second_input_u64, |acc, _| { + let tmp = *acc % 2; + *acc /= 2; + Some(F::from_canonical_u64(tmp)) + }) + .collect(); + + let powers_of_two: Vec = (0..self.gate.chunk_bits) + .map(|i| F::TWO.exp_u64(i as u64)) + .collect(); + let first_input_chunks: Vec = first_input_bits + .chunks(self.gate.chunk_bits) + .map(|bits| { + bits.iter() + .zip(powers_of_two.iter()) + .map(|(b, x)| *b * *x) + .fold(F::ZERO, |a, b| a + b) + }) + .collect(); + let second_input_chunks: Vec = second_input_bits + .chunks(self.gate.chunk_bits) + .map(|bits| { + bits.iter() + .zip(powers_of_two.iter()) + .map(|(b, x)| *b * *x) + .fold(F::ZERO, |a, b| a + b) + }) + .collect(); + + let chunks_equal: Vec = (0..self.gate.num_chunks()) + .map(|i| F::from_bool(first_input_chunks[i] == second_input_chunks[i])) + .collect(); + let equality_dummies: Vec = first_input_chunks + .iter() + .zip(second_input_chunks.iter()) + .map(|(f, s)| if *f == *s { F::ONE } else { F::ONE / (*f - *s) }) + .collect(); + + let z = F::TWO.exp_u64(self.gate.chunk_bits as u64) + first_input - second_input; + let z_bits: Vec = (0..self.gate.chunk_bits + 1) + .scan(z.to_canonical_u64(), |acc, _| { + let tmp = *acc % 2; + *acc /= 2; + Some(F::from_canonical_u64(tmp)) + }) + .collect(); + + out_buffer.set_wire(local_wire(self.gate.wire_z_val(self.copy)), z); + for b in 0..self.gate.chunk_bits + 1 { + out_buffer.set_wire(local_wire(self.gate.wire_z_bit(c, b)), z_bits[b]); + } + for i in 0..self.gate.num_chunks() { + out_buffer.set_wire( + local_wire(self.gate.wire_first_chunk_val(self.copy, i)), + first_input_chunks[i], + ); + out_buffer.set_wire( + local_wire(self.gate.wire_second_chunk_val(self.copy, i)), + second_input_chunks[i], + ); + out_buffer.set_wire( + local_wire(self.gate.wire_chunks_equal(self.copy, i)), + chunks_equal[i], + ); + out_buffer.set_wire( + local_wire(self.gate.wire_equality_dummy(self.copy, i)), + equality_dummies[i], + ); + } + } +} + +#[cfg(test)] +mod tests { + use std::marker::PhantomData; + + use anyhow::Result; + + use crate::field::crandall_field::CrandallField; + use crate::field::extension_field::quartic::QuarticCrandallField; + use crate::field::field_types::Field; + use crate::gates::comparison::ComparisonGate; + use crate::gates::gate::Gate; + use crate::gates::gate_testing::{test_eval_fns, test_low_degree}; + use crate::hash::hash_types::HashOut; + use crate::plonk::circuit_data::CircuitConfig; + use crate::plonk::vars::EvaluationVars; + + #[test] + fn wire_indices() { + type CG = ComparisonGate; + let num_copies = 3; + let chunk_bits = 3; + + let gate = CG { + chunk_bits, + num_copies, + _phantom: PhantomData, + }; + + assert_eq!(gate.wire_first_input(0, 0), 0); + assert_eq!(gate.wire_first_input(0, 2), 2); + assert_eq!(gate.wire_second_input(0, 0), 3); + assert_eq!(gate.wire_second_input(0, 2), 5); + assert_eq!(gate.wire_first_output(0, 0), 6); + assert_eq!(gate.wire_second_output(0, 2), 11); + assert_eq!(gate.wire_switch_bool(0), 12); + assert_eq!(gate.wire_first_input(1, 0), 13); + assert_eq!(gate.wire_second_output(1, 2), 24); + assert_eq!(gate.wire_switch_bool(1), 25); + assert_eq!(gate.wire_first_input(2, 0), 26); + assert_eq!(gate.wire_second_output(2, 2), 37); + assert_eq!(gate.wire_switch_bool(2), 38); + } + + #[test] + fn low_degree() { + test_low_degree::(SwitchGate::<_, 4>::new_from_config( + CircuitConfig::large_config(), + 3, + )); + } + + #[test] + fn eval_fns() -> Result<()> { + test_eval_fns::(SwitchGate::<_, 4>::new_from_config( + CircuitConfig::large_config(), + 3, + )) + } + + #[test] + fn test_gate_constraint() { + type F = CrandallField; + type FF = QuarticCrandallField; + const D: usize = 4; + const CHUNK_SIZE: usize = 4; + let num_copies = 3; + + /// Returns the local wires for a switch gate given the inputs and the switch booleans. + fn get_wires( + first_inputs: Vec>, + second_inputs: Vec>, + switch_bools: Vec, + ) -> Vec { + let num_copies = first_inputs.len(); + + let mut v = Vec::new(); + for c in 0..num_copies { + let switch = switch_bools[c]; + + let mut first_input_chunk = Vec::with_capacity(CHUNK_SIZE); + let mut second_input_chunk = Vec::with_capacity(CHUNK_SIZE); + let mut first_output_chunk = Vec::with_capacity(CHUNK_SIZE); + let mut second_output_chunk = Vec::with_capacity(CHUNK_SIZE); + for e in 0..CHUNK_SIZE { + let first_input = first_inputs[c][e]; + let second_input = second_inputs[c][e]; + let first_output = if switch { second_input } else { first_input }; + let second_output = if switch { first_input } else { second_input }; + first_input_chunk.push(first_input); + second_input_chunk.push(second_input); + first_output_chunk.push(first_output); + second_output_chunk.push(second_output); + } + v.append(&mut first_input_chunk); + v.append(&mut second_input_chunk); + v.append(&mut first_output_chunk); + v.append(&mut second_output_chunk); + + v.push(F::from_bool(switch)); + } + + v.iter().map(|&x| x.into()).collect::>() + } + + let first_inputs: Vec> = (0..num_copies).map(|_| F::rand_vec(CHUNK_SIZE)).collect(); + let second_inputs: Vec> = (0..num_copies).map(|_| F::rand_vec(CHUNK_SIZE)).collect(); + let switch_bools = vec![true, false, true]; + + let gate = SwitchGate:: { + chunk_bits: CHUNK_SIZE, + num_copies, + _phantom: PhantomData, + }; + + let vars = EvaluationVars { + local_constants: &[], + local_wires: &get_wires(first_inputs, second_inputs, switch_bools), + public_inputs_hash: &HashOut::rand(), + }; + + assert!( + gate.eval_unfiltered(vars).iter().all(|x| x.is_zero()), + "Gate constraints are not satisfied." + ); + } +} diff --git a/src/gates/mod.rs b/src/gates/mod.rs index 993623c3..053ba3fa 100644 --- a/src/gates/mod.rs +++ b/src/gates/mod.rs @@ -3,6 +3,7 @@ pub mod arithmetic; pub mod base_sum; +pub mod comparison; pub mod constant; pub mod exponentiation; pub mod gate;