From 07b71a961398ec12876b7d8be084dc7ccae8b84f Mon Sep 17 00:00:00 2001 From: Nicholas Ward Date: Thu, 20 Jan 2022 16:10:21 -0800 Subject: [PATCH] ECDSA merge --- plonky2/src/gadgets/biguint.rs | 6 +- plonky2/src/gadgets/multiple_comparison.rs | 13 +- plonky2/src/gates/add_many_u32.rs | 437 +++++++++++++++++++++ plonky2/src/gates/mod.rs | 1 + 4 files changed, 448 insertions(+), 9 deletions(-) create mode 100644 plonky2/src/gates/add_many_u32.rs diff --git a/plonky2/src/gadgets/biguint.rs b/plonky2/src/gadgets/biguint.rs index 6f9081b7..f647a30f 100644 --- a/plonky2/src/gadgets/biguint.rs +++ b/plonky2/src/gadgets/biguint.rs @@ -83,7 +83,7 @@ impl, const D: usize> CircuitBuilder { pub fn cmp_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BoolTarget { let (a, b) = self.pad_biguints(a, b); - self.list_le_binary::<30>(a.limbs, b.limbs) + self.list_le_30(a.limbs, b.limbs) } pub fn add_virtual_biguint_target(&mut self, num_limbs: usize) -> BigUintTarget { @@ -138,6 +138,8 @@ impl, const D: usize> CircuitBuilder { } pub fn mul_biguint(&mut self, a: &BigUintTarget, b: &BigUintTarget) -> BigUintTarget { + let before = self.num_gates(); + let total_limbs = a.limbs.len() + b.limbs.len(); let mut to_add = vec![vec![]; total_limbs]; @@ -159,6 +161,8 @@ impl, const D: usize> CircuitBuilder { } combined_limbs.push(carry); + println!("NUMBER OF GATES: {}", self.num_gates() - before); + BigUintTarget { limbs: combined_limbs, } diff --git a/plonky2/src/gadgets/multiple_comparison.rs b/plonky2/src/gadgets/multiple_comparison.rs index 222a6858..bca337fe 100644 --- a/plonky2/src/gadgets/multiple_comparison.rs +++ b/plonky2/src/gadgets/multiple_comparison.rs @@ -59,14 +59,11 @@ impl, const D: usize> CircuitBuilder { } /// Helper function for comparing, specifically, lists of `U32Target`s. - pub fn list_le_binary( - &mut self, - a: Vec>, - b: Vec>, - ) -> BoolTarget { - let a_targets = a.iter().map(|&t| t.0).collect(); - let b_targets = b.iter().map(|&t| t.0).collect(); - self.list_le(a_targets, b_targets, BITS) + pub fn list_le_30(&mut self, a: Vec>, b: Vec>) -> BoolTarget { + let a_targets: Vec = a.iter().map(|&t| t.0).collect(); + let b_targets: Vec = b.iter().map(|&t| t.0).collect(); + + self.list_le(a_targets, b_targets, 30) } } diff --git a/plonky2/src/gates/add_many_u32.rs b/plonky2/src/gates/add_many_u32.rs new file mode 100644 index 00000000..7ad8523d --- /dev/null +++ b/plonky2/src/gates/add_many_u32.rs @@ -0,0 +1,437 @@ +use std::marker::PhantomData; + +use itertools::unfold; + +use crate::field::extension_field::target::ExtensionTarget; +use crate::field::extension_field::Extendable; +use crate::field::field_types::{Field, 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}; + +const LOG2_MAX_NUM_ADDENDS: usize = 6; +const MAX_NUM_ADDENDS: usize = 1 << LOG2_MAX_NUM_ADDENDS; + +/// A gate to perform addition on `num_addends` different 32-bit values, plus a small carry +#[derive(Copy, Clone, Debug)] +pub struct U32AddManyGate, const D: usize> { + pub num_addends: usize, + pub num_ops: usize, + _phantom: PhantomData, +} + +impl, const D: usize> U32AddManyGate { + pub fn new_from_config(num_addends: usize, config: &CircuitConfig) -> Self { + Self { + num_addends, + num_ops: Self::num_ops(num_addends, config), + _phantom: PhantomData, + } + } + + pub(crate) fn num_ops(num_addends: usize, config: &CircuitConfig) -> usize { + debug_assert!(num_addends < MAX_NUM_ADDENDS); + let wires_per_op = (num_addends + 3) + Self::num_limbs(); + let routed_wires_per_op = 5; + (config.num_wires / wires_per_op).min(config.num_routed_wires / routed_wires_per_op) + } + + pub fn wire_ith_op_jth_addend(&self, i: usize, j: usize) -> usize { + debug_assert!(i < self.num_ops); + debug_assert!(i < self.num_addends); + (self.num_addends + 3) * i + j + } + pub fn wire_ith_carry(&self, i: usize) -> usize { + debug_assert!(i < self.num_ops); + (self.num_addends + 3) * i + self.num_addends + } + + pub fn wire_ith_output_low_half(&self, i: usize) -> usize { + debug_assert!(i < self.num_ops); + (self.num_addends + 3) * i + self.num_addends + 1 + } + pub fn wire_ith_output_high_half(&self, i: usize) -> usize { + debug_assert!(i < self.num_ops); + (self.num_addends + 3) * i + self.num_addends + 2 + } + + pub fn limb_bits() -> usize { + 2 + } + pub fn num_limbs() -> usize { + 32 / Self::limb_bits() + } + + pub fn wire_ith_output_jth_limb(&self, i: usize, j: usize) -> usize { + debug_assert!(i < self.num_ops); + debug_assert!(j < Self::num_limbs()); + (self.num_addends + 3) * self.num_ops + Self::num_limbs() * i + j + } +} + +impl, const D: usize> Gate for U32AddManyGate { + fn id(&self) -> String { + format!("{:?}", self) + } + + fn eval_unfiltered(&self, vars: EvaluationVars) -> Vec { + let mut constraints = Vec::with_capacity(self.num_constraints()); + for i in 0..self.num_ops { + let addends: Vec = (0..self.num_addends) + .map(|j| vars.local_wires[self.wire_ith_op_jth_addend(i, j)]) + .collect(); + let borrow = vars.local_wires[self.wire_ith_carry(i)]; + + let computed_output = addends.iter().fold(F::Extension::ZERO, |x, &y| x + y) + borrow; + + let output_low = vars.local_wires[self.wire_ith_output_low_half(i)]; + let output_high = vars.local_wires[self.wire_ith_output_high_half(i)]; + + let base = F::Extension::from_canonical_u64(1 << 32u64); + let combined_output = output_high * base + output_low; + + constraints.push(combined_output - computed_output); + + let mut combined_low_limbs = F::Extension::ZERO; + let base = F::Extension::from_canonical_u64(1u64 << Self::limb_bits()); + for j in (0..Self::num_limbs()).rev() { + let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)]; + let max_limb = 1 << Self::limb_bits(); + let product = (0..max_limb) + .map(|x| this_limb - F::Extension::from_canonical_usize(x)) + .product(); + constraints.push(product); + + combined_low_limbs = base * combined_low_limbs + this_limb; + } + constraints.push(combined_low_limbs - output_low); + + let max_overflow = self.num_addends; + let product = (0..max_overflow) + .map(|x| output_high - F::Extension::from_canonical_usize(x)) + .product(); + constraints.push(product); + } + + constraints + } + + fn eval_unfiltered_base(&self, vars: EvaluationVarsBase) -> Vec { + let mut constraints = Vec::with_capacity(self.num_constraints()); + for i in 0..self.num_ops { + let addends: Vec = (0..self.num_addends) + .map(|j| vars.local_wires[self.wire_ith_op_jth_addend(i, j)]) + .collect(); + let borrow = vars.local_wires[self.wire_ith_carry(i)]; + + let computed_output = addends.iter().fold(F::ZERO, |x, &y| x + y) + borrow; + + let output_low = vars.local_wires[self.wire_ith_output_low_half(i)]; + let output_high = vars.local_wires[self.wire_ith_output_high_half(i)]; + + let base = F::from_canonical_u64(1 << 32u64); + let combined_output = output_high * base + output_low; + + constraints.push(combined_output - computed_output); + + let mut combined_low_limbs = F::ZERO; + let base = F::from_canonical_u64(1u64 << Self::limb_bits()); + for j in (0..Self::num_limbs()).rev() { + let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)]; + let max_limb = 1 << Self::limb_bits(); + let product = (0..max_limb) + .map(|x| this_limb - F::from_canonical_usize(x)) + .product(); + constraints.push(product); + + combined_low_limbs = base * combined_low_limbs + this_limb; + } + constraints.push(combined_low_limbs - output_low); + + let max_overflow = self.num_addends; + let product = (0..max_overflow) + .map(|x| output_high - F::from_canonical_usize(x)) + .product(); + constraints.push(product); + } + + constraints + } + + fn eval_unfiltered_recursively( + &self, + builder: &mut CircuitBuilder, + vars: EvaluationTargets, + ) -> Vec> { + let mut constraints = Vec::with_capacity(self.num_constraints()); + + for i in 0..self.num_ops { + let addends: Vec> = (0..self.num_addends) + .map(|j| vars.local_wires[self.wire_ith_op_jth_addend(i, j)]) + .collect(); + let borrow = vars.local_wires[self.wire_ith_carry(i)]; + + let mut computed_output = borrow; + for addend in addends { + computed_output = builder.add_extension(computed_output, addend); + } + + let output_low = vars.local_wires[self.wire_ith_output_low_half(i)]; + let output_high = vars.local_wires[self.wire_ith_output_high_half(i)]; + + let base: F::Extension = F::from_canonical_u64(1 << 32u64).into(); + let base_target = builder.constant_extension(base); + let combined_output = builder.mul_add_extension(output_high, base_target, output_low); + + constraints.push(builder.sub_extension(combined_output, computed_output)); + + let mut combined_low_limbs = builder.zero_extension(); + let base = builder + .constant_extension(F::Extension::from_canonical_u64(1u64 << Self::limb_bits())); + for j in (0..Self::num_limbs()).rev() { + let this_limb = vars.local_wires[self.wire_ith_output_jth_limb(i, j)]; + let max_limb = 1 << Self::limb_bits(); + + let mut product = builder.one_extension(); + for x in 0..max_limb { + let x_target = + builder.constant_extension(F::Extension::from_canonical_usize(x)); + let diff = builder.sub_extension(this_limb, x_target); + product = builder.mul_extension(product, diff); + } + constraints.push(product); + + combined_low_limbs = builder.mul_add_extension(base, combined_low_limbs, this_limb); + } + constraints.push(builder.sub_extension(combined_low_limbs, output_low)); + + let max_overflow = self.num_addends; + let mut product = builder.one_extension(); + for x in 0..max_overflow { + let x_target = builder.constant_extension(F::Extension::from_canonical_usize(x)); + let diff = builder.sub_extension(output_high, x_target); + product = builder.mul_extension(product, diff); + } + constraints.push(product); + } + + constraints + } + + fn generators( + &self, + gate_index: usize, + _local_constants: &[F], + ) -> Vec>> { + (0..self.num_ops) + .map(|i| { + let g: Box> = Box::new( + U32AddManyGenerator { + gate: *self, + gate_index, + i, + _phantom: PhantomData, + } + .adapter(), + ); + g + }) + .collect::>() + } + + fn num_wires(&self) -> usize { + (self.num_addends + 3) * self.num_ops + Self::num_limbs() * self.num_ops + } + + fn num_constants(&self) -> usize { + 0 + } + + fn degree(&self) -> usize { + 1 << Self::limb_bits() + } + + fn num_constraints(&self) -> usize { + self.num_ops * (3 + Self::num_limbs()) + } +} + +#[derive(Clone, Debug)] +struct U32AddManyGenerator, const D: usize> { + gate: U32AddManyGate, + gate_index: usize, + i: usize, + _phantom: PhantomData, +} + +impl, const D: usize> SimpleGenerator + for U32AddManyGenerator +{ + fn dependencies(&self) -> Vec { + let local_target = |input| Target::wire(self.gate_index, input); + + (0..self.gate.num_addends) + .map(|j| local_target(self.gate.wire_ith_op_jth_addend(self.i, j))) + .chain([local_target(self.gate.wire_ith_carry(self.i))]) + .collect() + } + + 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 addends: Vec<_> = (0..self.gate.num_addends).map(|j| get_local_wire(self.gate.wire_ith_output_jth_limb(self.i, j))).collect(); + let carry = get_local_wire(self.gate.wire_ith_carry(self.i)); + + let output = addends.iter().fold(F::ZERO, |x, &y| x + y) + carry; + let mut output_u64 = output.to_canonical_u64(); + + let output_high_u64 = output_u64 >> 32; + let output_low_u64 = output_u64 & ((1 << 32) - 1); + + let output_high = F::from_canonical_u64(output_high_u64); + let output_low = F::from_canonical_u64(output_low_u64); + + let output_high_wire = local_wire(self.gate.wire_ith_output_high_half(self.i)); + let output_low_wire = local_wire(self.gate.wire_ith_output_low_half(self.i)); + + out_buffer.set_wire(output_high_wire, output_high); + out_buffer.set_wire(output_low_wire, output_low); + + let num_limbs = U32AddManyGate::::num_limbs(); + let limb_base = 1 << U32AddManyGate::::limb_bits(); + let output_limbs_u64 = unfold((), move |_| { + let ret = output_u64 % limb_base; + output_u64 /= limb_base; + Some(ret) + }) + .take(num_limbs); + let output_limbs_f = output_limbs_u64.map(F::from_canonical_u64); + + for (j, output_limb) in output_limbs_f.enumerate() { + let wire = local_wire(self.gate.wire_ith_output_jth_limb(self.i, j)); + out_buffer.set_wire(wire, output_limb); + } + } +} + +#[cfg(test)] +mod tests { + use std::marker::PhantomData; + + use anyhow::Result; + use rand::Rng; + + use crate::field::extension_field::quartic::QuarticExtension; + use crate::field::field_types::Field; + use crate::field::goldilocks_field::GoldilocksField; + use crate::gates::add_many_u32::U32AddManyGate; + 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::vars::EvaluationVars; + + #[test] + fn low_degree() { + test_low_degree::(U32AddManyGate:: { + num_addends: 4, + num_ops: 3, + _phantom: PhantomData, + }) + } + + #[test] + fn eval_fns() -> Result<()> { + test_eval_fns::(U32AddManyGate:: { + num_addends: 4, + num_ops: 3, + _phantom: PhantomData, + }) + } + + #[test] + fn test_gate_constraint() { + type F = GoldilocksField; + type FF = QuarticExtension; + const D: usize = 4; + const NUM_ADDENDS: usize = 4; + const NUM_U32_ADD_MANY_OPS: usize = 3; + + fn get_wires( + addends: Vec>, + carries: Vec, + ) -> Vec { + let mut v0 = Vec::new(); + let mut v1 = Vec::new(); + + let limb_bits = U32AddManyGate::::limb_bits(); + let num_limbs = U32AddManyGate::::num_limbs(); + let limb_base = 1 << limb_bits; + for op in 0..NUM_U32_ADD_MANY_OPS { + let adds = &addends[op]; + let ca = carries[op]; + + let mut output = adds.iter().sum::() + ca; + let output_low = output & ((1 << 32) - 1); + let output_high = output >> 32; + + let mut output_limbs = Vec::with_capacity(num_limbs); + for _i in 0..num_limbs { + output_limbs.push(output % limb_base); + output /= limb_base; + } + let mut output_limbs_f: Vec<_> = output_limbs + .into_iter() + .map(F::from_canonical_u64) + .collect(); + + for a in adds { + v0.push(F::from_canonical_u64(*a)); + } + v0.push(F::from_canonical_u64(ca)); + v0.push(F::from_canonical_u64(output_low)); + v0.push(F::from_canonical_u64(output_high)); + v1.append(&mut output_limbs_f); + } + + v0.iter() + .chain(v1.iter()) + .map(|&x| x.into()) + .collect::>() + } + + let mut rng = rand::thread_rng(); + let addends: Vec> = (0..NUM_U32_ADD_MANY_OPS) + .map(|_| (0..NUM_ADDENDS).map(|_| rng.gen::() as u64).collect()) + .collect(); + let carries: Vec<_> = (0..NUM_U32_ADD_MANY_OPS) + .map(|_| rng.gen::() as u64) + .collect(); + + let gate = U32AddManyGate:: { + num_addends: NUM_ADDENDS, + num_ops: NUM_U32_ADD_MANY_OPS, + _phantom: PhantomData, + }; + + let vars = EvaluationVars { + local_constants: &[], + local_wires: &get_wires(addends, carries), + public_inputs_hash: &HashOut::rand(), + }; + + assert!( + gate.eval_unfiltered(vars).iter().all(|x| x.is_zero()), + "Gate constraints are not satisfied." + ); + } +} diff --git a/plonky2/src/gates/mod.rs b/plonky2/src/gates/mod.rs index f63080e5..a7591648 100644 --- a/plonky2/src/gates/mod.rs +++ b/plonky2/src/gates/mod.rs @@ -1,6 +1,7 @@ // Gates have `new` methods that return `GateRef`s. #![allow(clippy::new_ret_no_self)] +pub mod add_many_u32; pub mod arithmetic_base; pub mod arithmetic_extension; pub mod arithmetic_u32;