Merge branch 'main' into insertion_gate

src/circuit_builder.rs

@@ -16,7 +16,7 @@ use crate::gates::gate_tree::Tree;
 use crate::gates::noop::NoopGate;
 use crate::generator::{CopyGenerator, RandomValueGenerator, WitnessGenerator};
 use crate::hash::hash_n_to_hash;
-use crate::permutation_argument::TargetPartitions;
+use crate::permutation_argument::TargetPartition;
 use crate::plonk_common::PlonkPolynomials;
 use crate::polynomial::commitment::ListPolynomialCommitment;
 use crate::polynomial::polynomial::PolynomialValues;
@@ -359,28 +359,34 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     fn sigma_vecs(&self, k_is: &[F], subgroup: &[F]) -> Vec<PolynomialValues<F>> {
         let degree = self.gate_instances.len();
         let degree_log = log2_strict(degree);
-        let mut target_partitions = TargetPartitions::new();
+        let mut target_partition = TargetPartition::new(|t| match t {
+            Target::Wire(Wire { gate, input }) => gate * self.config.num_routed_wires + input,
+            Target::PublicInput { index } => degree * self.config.num_routed_wires + index,
+            Target::VirtualTarget { index } => {
+                degree * self.config.num_routed_wires + self.public_input_index + index
+            }
+        });
 
         for gate in 0..degree {
             for input in 0..self.config.num_routed_wires {
-                target_partitions.add_partition(Target::Wire(Wire { gate, input }));
+                target_partition.add(Target::Wire(Wire { gate, input }));
             }
         }
 
         for index in 0..self.public_input_index {
-            target_partitions.add_partition(Target::PublicInput { index });
+            target_partition.add(Target::PublicInput { index });
         }
 
         for index in 0..self.virtual_target_index {
-            target_partitions.add_partition(Target::VirtualTarget { index });
+            target_partition.add(Target::VirtualTarget { index });
         }
 
         for &(a, b) in &self.copy_constraints {
-            target_partitions.merge(a, b);
+            target_partition.merge(a, b);
         }
 
-        let wire_partitions = target_partitions.to_wire_partitions();
-        wire_partitions.get_sigma_polys(degree_log, k_is, subgroup)
+        let wire_partition = target_partition.wire_partition();
+        wire_partition.get_sigma_polys(degree_log, k_is, subgroup)
     }
 
     /// Builds a "full circuit", with both prover and verifier data.

src/gadgets/insert.rs

@@ -37,8 +37,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         self.sub(one, not_equal)
     }
 
-    /// Inserts a `Target` in a vector at a non-deterministic index. This is done by rotating to the
-    /// left, inserting at 0 and then rotating to the right.
+    /// Inserts a `Target` in a vector at a non-deterministic index.
     /// Note: `index` is not range-checked.
     pub fn insert(
         &mut self,
@@ -49,9 +48,8 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let mut already_inserted = self.zero();
         let mut new_list = Vec::new();
 
-        for i in 0..v.len() {
         let one = self.one();
-
+        for i in 0..=v.len() {
             let cur_index = self.constant(F::from_canonical_usize(i));
             let insert_here = self.is_equal(cur_index, index);
 
@@ -63,11 +61,14 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             already_inserted = self.add(already_inserted, insert_here);
 
             let not_already_inserted = self.sub(one, already_inserted);
+            if i < v.len() {
                 new_item = self.scalar_mul_add_extension(not_already_inserted, v[i], new_item);
+            }
 
             new_list.push(new_item);
         }
 
+
         new_list
     }
 }
@@ -106,6 +107,8 @@ mod tests {
             let inserted = real_insert(i, elem, &v);
             let purported_inserted = builder.insert(it, elem, v.clone());
 
+            assert_eq!(inserted.len(), purported_inserted.len());
+
             for (x, y) in inserted.into_iter().zip(purported_inserted) {
                 builder.route_extension(x, y);
             }

src/gates/gmimc.rs

@@ -323,7 +323,7 @@ mod tests {
     use crate::gates::gmimc::{GMiMCGate, W};
     use crate::generator::generate_partial_witness;
     use crate::gmimc::gmimc_permute_naive;
-    use crate::permutation_argument::TargetPartitions;
+    use crate::permutation_argument::TargetPartition;
     use crate::target::Target;
     use crate::wire::Wire;
     use crate::witness::PartialWitness;

src/permutation_argument.rs

@@ -1,4 +1,6 @@
 use std::collections::HashMap;
+use std::fmt::Debug;
+use std::hash::Hash;
 
 use rayon::prelude::*;
 
@@ -7,85 +9,111 @@ use crate::polynomial::polynomial::PolynomialValues;
 use crate::target::Target;
 use crate::wire::Wire;
 
+/// Node in the Disjoint Set Forest.
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub struct ForestNode<T: Debug + Copy + Eq + PartialEq> {
+    t: T,
+    parent: usize,
+    size: usize,
+    index: usize,
+}
+
+/// Disjoint Set Forest data-structure following https://en.wikipedia.org/wiki/Disjoint-set_data_structure.
 #[derive(Debug, Clone)]
-pub struct TargetPartitions {
-    partitions: Vec<Vec<Target>>,
-    indices: HashMap<Target, usize>,
+pub struct TargetPartition<T: Debug + Copy + Eq + PartialEq + Hash, F: Fn(T) -> usize> {
+    forest: Vec<ForestNode<T>>,
+    /// Function to compute a node's index in the forest.
+    indices: F,
 }
 
-impl Default for TargetPartitions {
-    fn default() -> Self {
-        TargetPartitions::new()
-    }
-}
-
-impl TargetPartitions {
-    pub fn new() -> Self {
+impl<T: Debug + Copy + Eq + PartialEq + Hash, F: Fn(T) -> usize> TargetPartition<T, F> {
+    pub fn new(f: F) -> Self {
         Self {
-            partitions: Vec::new(),
-            indices: HashMap::new(),
+            forest: Vec::new(),
+            indices: f,
         }
     }
-
-    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);
+    pub fn add(&mut self, t: T) {
+        let index = self.forest.len();
+        debug_assert_eq!((self.indices)(t), index);
+        self.forest.push(ForestNode {
+            t,
+            parent: index,
+            size: 1,
+            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);
+    /// Path compression method, see https://en.wikipedia.org/wiki/Disjoint-set_data_structure#Finding_set_representatives.
+    pub fn find(&mut self, mut x: ForestNode<T>) -> ForestNode<T> {
+        if x.parent != x.index {
+            let root = self.find(self.forest[x.parent]);
+            self.forest[x.index].parent = root.index;
+            root
+        } else {
+            x
         }
     }
 
-    pub fn to_wire_partitions(&self) -> WirePartitions {
-        // Here we keep just the Wire targets, filtering out everything else.
-        let mut partitions = Vec::new();
+    /// Merge two sets.
+    pub fn merge(&mut self, tx: T, ty: T) {
+        let mut x = self.forest[(self.indices)(tx)];
+        let mut y = self.forest[(self.indices)(ty)];
+
+        x = self.find(x);
+        y = self.find(y);
+
+        if x == y {
+            return;
+        }
+
+        if x.size >= y.size {
+            y.parent = x.index;
+            x.size += y.size;
+        } else {
+            x.parent = y.index;
+            y.size += x.size;
+        }
+
+        self.forest[x.index] = x;
+        self.forest[y.index] = y;
+    }
+}
+impl<F: Fn(Target) -> usize> TargetPartition<Target, F> {
+    pub fn wire_partition(&mut self) -> WirePartitions {
+        let mut partition = HashMap::<_, Vec<_>>::new();
+        let nodes = self.forest.clone();
+        for x in nodes {
+            let v = partition.entry(self.find(x).t).or_default();
+            v.push(x.t);
+        }
+
         let mut indices = HashMap::new();
+        // // Here we keep just the Wire targets, filtering out everything else.
+        let partition = partition
+            .into_values()
+            .map(|v| {
+                v.into_iter()
+                    .filter_map(|t| match t {
+                        Target::Wire(w) => Some(w),
+                        _ => None,
+                    })
+                    .collect::<Vec<_>>()
+            })
+            .collect::<Vec<_>>();
+        partition.iter().enumerate().for_each(|(i, v)| {
+            v.iter().for_each(|t| {
+                indices.insert(*t, i);
+            });
+        });
 
-        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,
-        }
+        WirePartitions { partition, indices }
     }
 }
 
 pub struct WirePartitions {
-    partitions: Vec<Vec<Wire>>,
+    partition: Vec<Vec<Wire>>,
     indices: HashMap<Wire, usize>,
 }
 
@@ -95,7 +123,7 @@ impl WirePartitions {
     /// 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 partition = &self.partition[self.indices[&wire]];
         let n = partition.len();
         for i in 0..n {
             if partition[i] == wire {

src/plonk_challenger.rs

@@ -321,7 +321,7 @@ mod tests {
     use crate::field::crandall_field::CrandallField;
     use crate::field::field::Field;
     use crate::generator::generate_partial_witness;
-    use crate::permutation_argument::TargetPartitions;
+    use crate::permutation_argument::TargetPartition;
     use crate::plonk_challenger::{Challenger, RecursiveChallenger};
     use crate::target::Target;
     use crate::witness::PartialWitness;