From 63a3090452befc3799d41b145b8a17d89f6670b4 Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Wed, 16 Mar 2022 18:09:36 +0100
Subject: [PATCH] Start selectors

---
 plonky2/src/gates/mod.rs             |  1 +
 plonky2/src/gates/selectors.rs       | 40 ++++++++++++++++++++++++++++
 plonky2/src/plonk/circuit_builder.rs | 11 +++++++-
 plonky2/src/plonk/circuit_data.rs    |  2 +-
 4 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 plonky2/src/gates/selectors.rs

diff --git a/plonky2/src/gates/mod.rs b/plonky2/src/gates/mod.rs
index 18e3e99b..2450492f 100644
--- a/plonky2/src/gates/mod.rs
+++ b/plonky2/src/gates/mod.rs
@@ -24,6 +24,7 @@ pub mod random_access;
 pub mod range_check_u32;
 pub mod reducing;
 pub mod reducing_extension;
+pub(crate) mod selectors;
 pub mod subtraction_u32;
 pub mod switch;
 pub mod util;
diff --git a/plonky2/src/gates/selectors.rs b/plonky2/src/gates/selectors.rs
new file mode 100644
index 00000000..11463512
--- /dev/null
+++ b/plonky2/src/gates/selectors.rs
@@ -0,0 +1,40 @@
+use plonky2_field::extension_field::Extendable;
+use plonky2_field::polynomial::PolynomialValues;
+
+use crate::gates::gate::{GateInstance, GateRef};
+use crate::hash::hash_types::RichField;
+
+pub(crate) fn compute_selectors<F: RichField + Extendable<D>, const D: usize>(
+    mut gates: Vec<GateRef<F, D>>,
+    instances: &[GateInstance<F, D>],
+    max_degree: usize,
+) {
+    let n = instances.len();
+    gates.sort_unstable_by_key(|g| g.0.degree());
+
+    let mut combinations = Vec::new();
+    let mut pos = 0;
+
+    while pos < gates.len() {
+        let mut i = 0;
+        while (pos + i < gates.len()) && (i + gates[pos + i].0.degree() <= max_degree + 1) {
+            i += 1;
+        }
+        combinations.push((pos, pos + i));
+        pos += i;
+    }
+
+    let num_constants_polynomials =
+        0.max(gates.iter().map(|g| g.0.num_constants()).max().unwrap() - combinations.len() - 1);
+    let mut polynomials =
+        vec![PolynomialValues::zero(n); combinations.len() + num_constants_polynomials];
+
+    let index = |id| gates.iter().position(|g| g.0.id() == id).unwrap();
+    let combination = |i| combinations.iter().position(|&(a, _)| a <= i).unwrap();
+
+    for (j, g) in instances.iter().enumerate() {
+        let i = index(g.gate_ref.0.id());
+        let comb = combination(i);
+        polynomials[comb].values[j] = i - combinations[comb].0;
+    }
+}
diff --git a/plonky2/src/plonk/circuit_builder.rs b/plonky2/src/plonk/circuit_builder.rs
index 8e2f2e10..50b5931a 100644
--- a/plonky2/src/plonk/circuit_builder.rs
+++ b/plonky2/src/plonk/circuit_builder.rs
@@ -23,6 +23,7 @@ use crate::gates::gate::{CurrentSlot, Gate, GateInstance, GateRef, PrefixedGate}
 use crate::gates::gate_tree::Tree;
 use crate::gates::noop::NoopGate;
 use crate::gates::public_input::PublicInputGate;
+use crate::gates::selectors::compute_selectors;
 use crate::hash::hash_types::{HashOutTarget, MerkleCapTarget, RichField};
 use crate::hash::merkle_proofs::MerkleProofTarget;
 use crate::iop::ext_target::ExtensionTarget;
@@ -669,7 +670,15 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             "FRI total reduction arity is too large.",
         );
 
-        let gates = self.gates.iter().cloned().collect();
+        let gates = self.gates.iter().cloned().collect::<Vec<_>>();
+        for g in &gates {
+            println!("{} {}", g.0.id(), g.0.num_constants());
+        }
+        dbg!(compute_selectors(
+            gates.clone(),
+            &self.gate_instances,
+            self.config.max_quotient_degree_factor + 1
+        ));
         let (gate_tree, max_filtered_constraint_degree, num_constants) = Tree::from_gates(gates);
         let prefixed_gates = PrefixedGate::from_tree(gate_tree);
 
diff --git a/plonky2/src/plonk/circuit_data.rs b/plonky2/src/plonk/circuit_data.rs
index 3d4ee2df..f0fa4bd5 100644
--- a/plonky2/src/plonk/circuit_data.rs
+++ b/plonky2/src/plonk/circuit_data.rs
@@ -63,7 +63,7 @@ impl CircuitConfig {
         Self {
             num_wires: 135,
             num_routed_wires: 80,
-            constant_gate_size: 5,
+            constant_gate_size: 2,
             use_base_arithmetic_gate: true,
             security_bits: 100,
             num_challenges: 2,