Fix: under-constrained idx_to_indicator (#17)

*fix(BUG): `GateChip::idx_to_indicator` still had soundness bug where at index
`idx` the value could be 0 or 1 (instead of only 1)

* feat: add some function documentation

* test(idx_to_indicator): add comprehensive tests
* both positive and negative tests

halo2-base/src/gates/flex_gate.rs

@@ -473,48 +473,50 @@ pub trait GateInstructions<F: ScalarField> {
         indicator.split_off((1 << k) - 2)
     }
 
-    // returns vec with vec.len() == len such that:
-    //     vec[i] == 1{i == idx}
+    /// Returns a vector `indicator` of length `len` where
+    /// ```
+    /// indicator[i] == i == idx ? 1 : 0
+    /// ```
+    /// If `idx >= len` then `indicator` is all zeros.
     fn idx_to_indicator(
         &self,
         ctx: &mut Context<F>,
         idx: impl Into<QuantumCell<F>>,
         len: usize,
     ) -> Vec<AssignedValue<F>> {
         let mut idx = idx.into();
-        let mut ind = Vec::with_capacity(len);
-        let idx_val = idx.value().get_lower_32() as usize;
-        for i in 0..len {
-            // check ind[i] * (i - idx) == 0
-            let ind_val = F::from(idx_val == i);
-            let val = if idx_val == i { *idx.value() } else { F::zero() };
-            ctx.assign_region_smart(
-                [
-                    Constant(F::zero()),
-                    Witness(ind_val),
-                    idx,
-                    Witness(val),
-                    Constant(-F::from(i as u64)),
-                    Witness(ind_val),
-                    Constant(F::zero()),
-                ],
-                [0, 3],
-                [(1, 5)],
-                [],
-            );
-            // need to use assigned idx after i > 0 so equality constraint holds
-            if i == 0 {
-                idx = Existing(ctx.get(-5));
-            }
-            let ind_cell = ctx.get(-2);
-            self.assert_bit(ctx, ind_cell);
-            ind.push(ind_cell);
-        }
-        ind
+        (0..len)
+            .map(|i| {
+                // need to use assigned idx after i > 0 so equality constraint holds
+                if i == 0 {
+                    // unroll `is_zero` to make sure if `idx == Witness(_)` it is replaced by `Existing(_)` in later iterations
+                    let x = idx.value();
+                    let (is_zero, inv) = if x.is_zero_vartime() {
+                        (F::one(), Assigned::Trivial(F::one()))
+                    } else {
+                        (F::zero(), Assigned::Rational(F::one(), *x))
+                    };
+                    let cells = [
+                        Witness(is_zero),
+                        idx,
+                        WitnessFraction(inv),
+                        Constant(F::one()),
+                        Constant(F::zero()),
+                        idx,
+                        Witness(is_zero),
+                        Constant(F::zero()),
+                    ];
+                    ctx.assign_region_smart(cells, [0, 4], [(0, 6), (1, 5)], []); // note the two `idx` need to be constrained equal: (1, 5)
+                    idx = Existing(ctx.get(-3)); // replacing `idx` with Existing cell so future loop iterations constrain equality of all `idx`s
+                    ctx.get(-2)
+                } else {
+                    self.is_equal(ctx, idx, Constant(self.get_field_element(i as u64)))
+                }
+            })
+            .collect()
     }
 
-    // performs inner product on a, indicator
-    // `indicator` values are all boolean
+    /// Performs inner product on `<a, indicator>`. Assumes that `a` and `indicator` are of the same length.
     /// Assumes for witness generation that only one element of `indicator` has non-zero value and that value is `F::one()`.
     fn select_by_indicator<Q>(
         &self,
@@ -540,6 +542,8 @@ pub trait GateInstructions<F: ScalarField> {
         ctx.assign_region_last(cells, (0..len).map(|i| 3 * i as isize))
     }
 
+    /// Given `cells` and `idx`, returns `cells[idx]` if `idx < cells.len()`.
+    /// If `idx >= cells.len()` then returns `F::zero()`.
     fn select_from_idx<Q>(
         &self,
         ctx: &mut Context<F>,

halo2-base/src/gates/tests/idx_to_indicator.rs

@@ -0,0 +1,107 @@
+use crate::halo2_proofs::{
+    plonk::keygen_pk,
+    plonk::{keygen_vk, Assigned},
+    poly::kzg::commitment::ParamsKZG,
+};
+
+use itertools::Itertools;
+use rand::{thread_rng, Rng};
+
+use super::*;
+use crate::QuantumCell::Witness;
+
+// soundness checks for `idx_to_indicator` function
+fn test_idx_to_indicator_gen(k: u32, len: usize) {
+    // first create proving and verifying key
+    let mut builder = GateThreadBuilder::keygen();
+    let gate = GateChip::default();
+    let dummy_idx = Witness(Fr::zero());
+    let indicator = gate.idx_to_indicator(builder.main(0), dummy_idx, len);
+    // get the offsets of the indicator cells for later 'pranking'
+    let ind_offsets = indicator.iter().map(|ind| ind.cell.unwrap().offset).collect::<Vec<_>>();
+    // set env vars
+    builder.config(k as usize, Some(9));
+    let circuit = GateCircuitBuilder::keygen(builder);
+
+    let params = ParamsKZG::setup(k, OsRng);
+    // generate proving key
+    let vk = keygen_vk(&params, &circuit).unwrap();
+    let pk = keygen_pk(&params, vk, &circuit).unwrap();
+    let vk = pk.get_vk(); // pk consumed vk
+
+    // now create different proofs to test the soundness of the circuit
+
+    let gen_pf = |idx: usize, ind_witnesses: &[Fr]| {
+        let mut builder = GateThreadBuilder::prover();
+        let gate = GateChip::default();
+        let idx = Witness(Fr::from(idx as u64));
+        gate.idx_to_indicator(builder.main(0), idx, len);
+        // prank the indicator cells
+        for (offset, witness) in ind_offsets.iter().zip_eq(ind_witnesses) {
+            builder.main(0).advice[*offset] = Assigned::Trivial(*witness);
+        }
+        let circuit = GateCircuitBuilder::prover(builder, vec![vec![]]); // no break points
+        gen_proof(&params, &pk, circuit)
+    };
+
+    // expected answer
+    for idx in 0..len {
+        let mut ind_witnesses = vec![Fr::zero(); len];
+        ind_witnesses[idx] = Fr::one();
+        let pf = gen_pf(idx, &ind_witnesses);
+        check_proof(&params, vk, &pf, true);
+    }
+
+    let mut rng = thread_rng();
+    // bad cases
+    for idx in 0..len {
+        let mut ind_witnesses = vec![Fr::zero(); len];
+        // all zeros is bad!
+        let pf = gen_pf(idx, &ind_witnesses);
+        check_proof(&params, vk, &pf, false);
+
+        // ind[idx] != 1 is bad!
+        for _ in 0..100usize {
+            ind_witnesses.fill(Fr::zero());
+            ind_witnesses[idx] = Fr::random(OsRng);
+            if ind_witnesses[idx] == Fr::one() {
+                continue;
+            }
+            let pf = gen_pf(idx, &ind_witnesses);
+            check_proof(&params, vk, &pf, false);
+        }
+
+        if len < 2 {
+            continue;
+        }
+        // nonzeros where there should be zeros is bad!
+        for _ in 0..100usize {
+            ind_witnesses.fill(Fr::zero());
+            ind_witnesses[idx] = Fr::one();
+            let num_nonzeros = rng.gen_range(1..len);
+            let mut count = 0usize;
+            for _ in 0..num_nonzeros {
+                let index = rng.gen_range(0..len);
+                if index == idx {
+                    continue;
+                }
+                ind_witnesses[index] = Fr::random(&mut rng);
+                count += 1;
+            }
+            if count == 0usize {
+                continue;
+            }
+            let pf = gen_pf(idx, &ind_witnesses);
+            check_proof(&params, vk, &pf, false);
+        }
+    }
+}
+
+#[test]
+fn test_idx_to_indicator() {
+    test_idx_to_indicator_gen(8, 1);
+    test_idx_to_indicator_gen(8, 4);
+    test_idx_to_indicator_gen(8, 10);
+    test_idx_to_indicator_gen(8, 20);
+    test_idx_to_indicator_gen(11, 100);
+}

halo2-base/src/gates/tests.rs → halo2-base/src/gates/tests/mod.rs

@@ -1,13 +1,28 @@
 use super::builder::{GateCircuitBuilder, GateThreadBuilder, RangeCircuitBuilder};
 use super::flex_gate::{GateChip, GateInstructions};
 use super::range::{RangeChip, RangeInstructions};
-use crate::halo2_proofs::{dev::MockProver, halo2curves::bn256::Fr};
+use crate::halo2_proofs::{
+    dev::MockProver,
+    halo2curves::bn256::{Bn256, Fr, G1Affine},
+    plonk::{create_proof, ProvingKey},
+    plonk::{verify_proof, VerifyingKey},
+    poly::commitment::ParamsProver,
+    poly::kzg::{
+        commitment::KZGCommitmentScheme, commitment::ParamsKZG, multiopen::ProverSHPLONK,
+        multiopen::VerifierSHPLONK, strategy::SingleStrategy,
+    },
+    transcript::{
+        Blake2bRead, Blake2bWrite, Challenge255, TranscriptReadBuffer, TranscriptWriterBuffer,
+    },
+};
 use crate::utils::{BigPrimeField, ScalarField};
 use crate::{Context, QuantumCell::Constant};
 use ff::Field;
 use rand::rngs::OsRng;
 use rayon::prelude::*;
 
+mod idx_to_indicator;
+
 fn gate_tests<F: ScalarField>(ctx: &mut Context<F>, inputs: [F; 3]) {
     let [a, b, c]: [_; 3] = ctx.assign_witnesses(inputs).try_into().unwrap();
     let chip = GateChip::default();
@@ -166,3 +181,47 @@ fn plot_range() {
     let circuit = RangeCircuitBuilder::keygen(builder);
     halo2_proofs::dev::CircuitLayout::default().render(7, &circuit, &root).unwrap();
 }
+
+// helper functions
+
+pub fn gen_proof(
+    params: &ParamsKZG<Bn256>,
+    pk: &ProvingKey<G1Affine>,
+    circuit: GateCircuitBuilder<Fr>,
+) -> Vec<u8> {
+    let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
+    create_proof::<
+        KZGCommitmentScheme<Bn256>,
+        ProverSHPLONK<'_, Bn256>,
+        Challenge255<_>,
+        _,
+        Blake2bWrite<Vec<u8>, G1Affine, _>,
+        _,
+    >(params, pk, &[circuit], &[&[]], OsRng, &mut transcript)
+    .expect("prover should not fail");
+    transcript.finalize()
+}
+
+pub fn check_proof(
+    params: &ParamsKZG<Bn256>,
+    vk: &VerifyingKey<G1Affine>,
+    proof: &[u8],
+    expect_satisfied: bool,
+) {
+    let verifier_params = params.verifier_params();
+    let strategy = SingleStrategy::new(params);
+    let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(proof);
+    let res = verify_proof::<
+        KZGCommitmentScheme<Bn256>,
+        VerifierSHPLONK<'_, Bn256>,
+        Challenge255<G1Affine>,
+        Blake2bRead<&[u8], G1Affine, Challenge255<G1Affine>>,
+        SingleStrategy<'_, Bn256>,
+    >(verifier_params, vk, strategy, &[&[]], &mut transcript);
+
+    if expect_satisfied {
+        assert!(res.is_ok());
+    } else {
+        assert!(res.is_err());
+    }
+}