Merge pull request #21 from mir-protocol/perm_arg

Port over some code for the permutation argument

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<F: Field> {
     /// The concrete placement of each gate.
     gate_instances: Vec<GateInstance<F>>,
 
+    /// 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<Box<dyn WitnessGenerator<F>>>,
 
@@ -45,13 +51,25 @@ impl<F: Field> CircuitBuilder<F> {
             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<Target> {
+        (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<F: Field> CircuitBuilder<F> {
             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<Box<dyn WitnessGenerator<F>>>) {
@@ -218,9 +236,27 @@ impl<F: Field> CircuitBuilder<F> {
             .collect()
     }
 
-    fn sigma_vecs(&self) -> Vec<PolynomialValues<F>> {
-        vec![PolynomialValues::zero(self.gate_instances.len()); self.config.num_routed_wires]
-        // TODO
+    fn sigma_vecs(&self, k_is: &[F]) -> Vec<PolynomialValues<F>> {
+        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<F: Field> CircuitBuilder<F> {
         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<F: Field> CircuitBuilder<F> {
             .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];

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;

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<Vec<Target>>,
+    indices: HashMap<Target, usize>,
+}
+
+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<Vec<Wire>>,
+    indices: HashMap<Wire, usize>,
+}
+
+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<F: Field>(
+        &self,
+        degree_log: usize,
+        k_is: &[F],
+    ) -> Vec<PolynomialValues<F>> {
+        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::<Vec<_>>();
+                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<usize> {
+        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
+    }
+}