diff --git a/src/circuit_builder.rs b/src/circuit_builder.rs index d21279d2..2be48e25 100644 --- a/src/circuit_builder.rs +++ b/src/circuit_builder.rs @@ -15,6 +15,7 @@ use crate::gates::noop::NoopGate; use crate::generator::{CopyGenerator, WitnessGenerator}; use crate::hash::hash_n_to_hash; use crate::merkle_tree::MerkleTree; +use crate::permutation_argument::TargetPartitions; use crate::polynomial::polynomial::PolynomialValues; use crate::target::Target; use crate::util::{log2_strict, transpose, transpose_poly_values}; @@ -29,9 +30,14 @@ pub struct CircuitBuilder { /// The concrete placement of each gate. gate_instances: Vec>, + /// The next available index for a public input. + public_input_index: usize, + /// The next available index for a VirtualAdviceTarget. virtual_target_index: usize, + copy_constraints: Vec<(Target, Target)>, + /// Generators used to generate the witness. generators: Vec>>, @@ -45,13 +51,25 @@ impl CircuitBuilder { config, gates: HashSet::new(), gate_instances: Vec::new(), + public_input_index: 0, virtual_target_index: 0, + copy_constraints: Vec::new(), generators: Vec::new(), constants_to_targets: HashMap::new(), targets_to_constants: HashMap::new(), } } + pub fn add_public_input(&mut self) -> Target { + let index = self.public_input_index; + self.public_input_index += 1; + Target::PublicInput { index } + } + + pub fn add_public_inputs(&mut self, n: usize) -> Vec { + (0..n).map(|_i| self.add_public_input()).collect() + } + /// Adds a new "virtual" advice target. This is not an actual wire in the witness, but just a /// target that help facilitate witness generation. In particular, a generator can assign a /// values to a virtual target, which can then be copied to other (virtual or concrete) targets @@ -127,7 +145,7 @@ impl CircuitBuilder { y.is_routable(self.config), "Tried to route a wire that isn't routable" ); - // TODO: Add to copy_constraints. + self.copy_constraints.push((x, y)); } pub fn add_generators(&mut self, generators: Vec>>) { @@ -218,9 +236,27 @@ impl CircuitBuilder { .collect() } - fn sigma_vecs(&self) -> Vec> { - vec![PolynomialValues::zero(self.gate_instances.len()); self.config.num_routed_wires] - // TODO + fn sigma_vecs(&self, k_is: &[F]) -> Vec> { + let degree = self.gate_instances.len(); + let degree_log = log2_strict(degree); + let mut target_partitions = TargetPartitions::new(); + + for gate in 0..degree { + for input in 0..self.config.num_routed_wires { + target_partitions.add_partition(Target::Wire(Wire { gate, input })); + } + } + + for index in 0..self.public_input_index { + target_partitions.add_partition(Target::PublicInput { index }) + } + + for &(a, b) in &self.copy_constraints { + target_partitions.merge(a, b); + } + + let wire_partitions = target_partitions.to_wire_partitions(); + wire_partitions.get_sigma_polys(degree_log, k_is) } /// Builds a "full circuit", with both prover and verifier data. @@ -239,7 +275,8 @@ impl CircuitBuilder { let constant_ldes_t = transpose_poly_values(constant_ldes); let constants_tree = MerkleTree::new(constant_ldes_t, true); - let sigma_vecs = self.sigma_vecs(); + let k_is = get_unique_coset_shifts(degree, self.config.num_routed_wires); + let sigma_vecs = self.sigma_vecs(&k_is); let sigma_ldes = PolynomialValues::lde_multiple(sigma_vecs, self.config.rate_bits); let sigma_ldes_t = transpose_poly_values(sigma_ldes); let sigmas_tree = MerkleTree::new(sigma_ldes_t, true); @@ -270,7 +307,6 @@ impl CircuitBuilder { .expect("No gates?"); let degree_bits = log2_strict(degree); - let k_is = get_unique_coset_shifts(degree, self.config.num_routed_wires); // TODO: This should also include an encoding of gate constraints. let circuit_digest_parts = [constants_root.elements, sigmas_root.elements]; diff --git a/src/lib.rs b/src/lib.rs index 18c1555a..09f90f81 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,6 +9,7 @@ pub mod gmimc; pub mod hash; pub mod merkle_proofs; mod merkle_tree; +mod permutation_argument; pub mod plonk_challenger; pub mod plonk_common; pub mod polynomial; diff --git a/src/permutation_argument.rs b/src/permutation_argument.rs new file mode 100644 index 00000000..af74aed2 --- /dev/null +++ b/src/permutation_argument.rs @@ -0,0 +1,146 @@ +use std::collections::HashMap; + +use rayon::prelude::*; + +use crate::field::field::Field; +use crate::polynomial::polynomial::PolynomialValues; +use crate::target::Target; +use crate::wire::Wire; + +#[derive(Debug, Clone)] +pub struct TargetPartitions { + partitions: Vec>, + indices: HashMap, +} + +impl Default for TargetPartitions { + fn default() -> Self { + TargetPartitions::new() + } +} + +impl TargetPartitions { + pub fn new() -> Self { + Self { + partitions: Vec::new(), + indices: HashMap::new(), + } + } + + pub fn get_partition(&self, target: Target) -> &[Target] { + &self.partitions[self.indices[&target]] + } + + /// Add a new partition with a single member. + pub fn add_partition(&mut self, target: Target) { + let index = self.partitions.len(); + self.partitions.push(vec![target]); + self.indices.insert(target, index); + } + + /// Merge the two partitions containing the two given targets. Does nothing if the targets are + /// already members of the same partition. + pub fn merge(&mut self, a: Target, b: Target) { + let a_index = self.indices[&a]; + let b_index = self.indices[&b]; + if a_index != b_index { + // Merge a's partition into b's partition, leaving a's partition empty. + // We have to clone because Rust's borrow checker doesn't know that + // self.partitions[b_index] and self.partitions[b_index] are disjoint. + let mut a_partition = self.partitions[a_index].clone(); + let b_partition = &mut self.partitions[b_index]; + for a_sibling in &a_partition { + *self.indices.get_mut(a_sibling).unwrap() = b_index; + } + b_partition.append(&mut a_partition); + } + } + + pub fn to_wire_partitions(&self) -> WirePartitions { + // Here we just drop all CircuitInputs, leaving all GateInputs. + let mut partitions = Vec::new(); + let mut indices = HashMap::new(); + + for old_partition in &self.partitions { + let mut new_partition = Vec::new(); + for target in old_partition { + if let Target::Wire(w) = *target { + new_partition.push(w); + } + } + partitions.push(new_partition); + } + + for (&target, &index) in &self.indices { + if let Target::Wire(gi) = target { + indices.insert(gi, index); + } + } + + WirePartitions { + partitions, + indices, + } + } +} + +pub struct WirePartitions { + partitions: Vec>, + indices: HashMap, +} + +impl WirePartitions { + /// Find a wire's "neighbor" in the context of Plonk's "extended copy constraints" check. In + /// other words, find the next wire in the given wire's partition. If the given wire is last in + /// its partition, this will loop around. If the given wire has a partition all to itself, it + /// is considered its own neighbor. + fn get_neighbor(&self, wire: Wire) -> Wire { + let partition = &self.partitions[self.indices[&wire]]; + let n = partition.len(); + for i in 0..n { + if partition[i] == wire { + let neighbor_index = (i + 1) % n; + return partition[neighbor_index]; + } + } + panic!("Wire not found in the expected partition") + } + + pub(crate) fn get_sigma_polys( + &self, + degree_log: usize, + k_is: &[F], + ) -> Vec> { + let degree = 1 << degree_log; + let subgroup_generator = F::primitive_root_of_unity(degree_log); + let sigma = self.get_sigma_map(degree); + + sigma + .chunks(degree) + .map(|chunk| { + let values = chunk + .par_iter() + .map(|&x| k_is[x / degree] * subgroup_generator.exp_usize(x % degree)) + .collect::>(); + PolynomialValues::new(values) + }) + .collect() + } + + /// Generates sigma in the context of Plonk, which is a map from `[kn]` to `[kn]`, where `k` is + /// the number of routed wires and `n` is the number of gates. + fn get_sigma_map(&self, degree: usize) -> Vec { + debug_assert_eq!(self.indices.len() % degree, 0); + let num_routed_wires = self.indices.len() / degree; + + let mut sigma = Vec::new(); + for input in 0..num_routed_wires { + for gate in 0..degree { + let wire = Wire { gate, input }; + let neighbor = self.get_neighbor(wire); + sigma.push(neighbor.input * degree + neighbor.gate); + } + } + sigma + } +}