feat: add handcrafted shuffle argument with ZK disabled as example

halo2_proofs/examples/shuffle.rs

@@ -0,0 +1,356 @@
+use ff::BatchInvert;
+use halo2_proofs::{
+    arithmetic::FieldExt,
+    circuit::{floor_planner::V1, Layouter},
+    dev::{metadata, FailureLocation, MockProver, VerifyFailure},
+    plonk::*,
+    poly::{
+        commitment::ParamsProver,
+        kzg::{
+            commitment::{KZGCommitmentScheme, ParamsKZG},
+            multiopen::{ProverSHPLONK, VerifierSHPLONK},
+            strategy::BatchVerifier,
+        },
+        Rotation, VerificationStrategy,
+    },
+    transcript::{
+        Blake2bRead, Blake2bWrite, Challenge255, TranscriptReadBuffer, TranscriptWriterBuffer,
+    },
+};
+use halo2curves::{bn256::Bn256, pairing::MultiMillerLoop};
+use rand_core::{OsRng, RngCore};
+use std::{fmt::Debug, iter};
+
+const THETA: &'static [u8; 64] = &[
+    18, 203, 39, 115, 47, 209, 21, 212, 155, 233, 255, 31, 27, 230, 235, 78, 138, 247, 124, 4, 60,
+    17, 247, 171, 251, 15, 64, 156, 16, 241, 58, 37, 23, 54, 167, 93, 115, 132, 10, 22, 103, 117,
+    38, 45, 166, 217, 106, 78, 77, 77, 137, 159, 192, 225, 125, 166, 139, 67, 44, 237, 210, 198,
+    165, 125,
+];
+const GAMMA: &'static [u8; 64] = &[
+    39, 196, 169, 65, 0, 78, 115, 65, 233, 180, 221, 250, 75, 73, 223, 218, 91, 79, 30, 167, 215,
+    33, 193, 222, 97, 242, 232, 21, 18, 131, 214, 133, 25, 217, 37, 71, 209, 188, 122, 126, 150,
+    91, 19, 222, 151, 141, 253, 253, 116, 14, 58, 59, 164, 28, 40, 27, 181, 144, 47, 146, 159, 126,
+    138, 49,
+];
+
+fn rand_2d_array<F: FieldExt, R: RngCore, const W: usize, const H: usize>(
+    rng: &mut R,
+) -> [[F; H]; W] {
+    [(); W].map(|_| [(); H].map(|_| F::random(&mut *rng)))
+}
+
+fn shuffled<F: FieldExt, R: RngCore, const W: usize, const H: usize>(
+    original: [[F; H]; W],
+    rng: &mut R,
+) -> [[F; H]; W] {
+    let mut shuffled = original;
+
+    for row in (1..H).rev() {
+        let rand_row = (rng.next_u32() as usize) % row;
+        for column in shuffled.iter_mut() {
+            column.swap(row, rand_row);
+        }
+    }
+
+    shuffled
+}
+
+#[derive(Clone)]
+struct MyConfig<F: FieldExt, const W: usize> {
+    q_first: Selector,
+    original: [Column<Advice>; W],
+    shuffled: [Column<Advice>; W],
+    theta: F,
+    gamma: F,
+    z: Column<Advice>,
+}
+
+impl<F: FieldExt, const W: usize> MyConfig<F, W> {
+    fn configure(meta: &mut ConstraintSystem<F, false>) -> Self {
+        // Phase 0
+        let q_first = meta.selector();
+        // Phase 1
+        let original = [(); W].map(|_| meta.advice_column());
+        let shuffled = [(); W].map(|_| meta.advice_column());
+        let [theta, gamma] = [F::from_bytes_wide(THETA), F::from_bytes_wide(GAMMA)];
+        // Phase 2
+        let z = meta.advice_column();
+
+        meta.create_gate("z should start and end with 1", |meta| {
+            let q_first = meta.query_selector(q_first);
+            let z = meta.query_advice(z, Rotation::cur());
+            let one = Expression::Constant(F::one());
+
+            vec![q_first * (one - z)]
+        });
+
+        meta.create_gate("z should have valid transition", |meta| {
+            let original = original.map(|advice| meta.query_advice(advice, Rotation::cur()));
+            let shuffled = shuffled.map(|advice| meta.query_advice(advice, Rotation::cur()));
+            let [theta, gamma] = [theta, gamma].map(|challenge| Expression::Constant(challenge));
+            let [z, z_w] =
+                [Rotation::cur(), Rotation::next()].map(|rotation| meta.query_advice(z, rotation));
+
+            // compress
+            let original = original
+                .iter()
+                .cloned()
+                .reduce(|acc, a| acc * theta.clone() + a)
+                .unwrap();
+            let shuffled = shuffled
+                .iter()
+                .cloned()
+                .reduce(|acc, a| acc * theta.clone() + a)
+                .unwrap();
+
+            vec![z * (original + gamma.clone()) - z_w * (shuffled + gamma)]
+        });
+
+        Self {
+            q_first,
+            original,
+            shuffled,
+            theta,
+            gamma,
+            z,
+        }
+    }
+}
+
+#[derive(Clone, Default)]
+struct MyCircuit<F: FieldExt, const W: usize, const H: usize> {
+    original: Option<[[F; H]; W]>,
+    shuffled: Option<[[F; H]; W]>,
+}
+
+impl<F: FieldExt, const W: usize, const H: usize> MyCircuit<F, W, H> {
+    fn rand<R: RngCore>(rng: &mut R) -> Self {
+        let original = rand_2d_array::<F, _, W, H>(rng);
+        let shuffled = shuffled(original, rng);
+
+        Self {
+            original: Some(original),
+            shuffled: Some(shuffled),
+        }
+    }
+}
+
+impl<F: FieldExt, const W: usize, const H: usize> Circuit<F, false> for MyCircuit<F, W, H> {
+    type Config = MyConfig<F, W>;
+    type FloorPlanner = V1;
+
+    fn without_witnesses(&self) -> Self {
+        Self::default()
+    }
+
+    fn configure(meta: &mut ConstraintSystem<F, false>) -> Self::Config {
+        MyConfig::configure(meta)
+    }
+
+    fn synthesize(
+        &self,
+        config: Self::Config,
+        mut layouter: impl Layouter<F>,
+    ) -> Result<(), Error> {
+        let theta = config.theta;
+        let gamma = config.gamma;
+
+        layouter.assign_region(
+            || "Shuffle original into shuffled",
+            |mut region| {
+                // Keygen (Phase 0)
+                config.q_first.enable(&mut region, 0)?;
+
+                // Phase 1
+                for (idx, &column) in config.original.iter().enumerate() {
+                    for offset in 0..H {
+                        region.assign_advice(
+                            || format!("original[{}][{}]", idx, offset),
+                            column,
+                            offset,
+                            || {
+                                self.original
+                                    .as_ref()
+                                    .map(|original| original[idx][offset])
+                                    .ok_or(Error::Synthesis)
+                            },
+                        )?;
+                    }
+                }
+                for (idx, &column) in config.shuffled.iter().enumerate() {
+                    for offset in 0..H {
+                        region.assign_advice(
+                            || format!("shuffled[{}][{}]", idx, offset),
+                            column,
+                            offset,
+                            || {
+                                self.shuffled
+                                    .as_ref()
+                                    .map(|shuffled| shuffled[idx][offset])
+                                    .ok_or(Error::Synthesis)
+                            },
+                        )?;
+                    }
+                }
+
+                // Phase 2
+                let z = if let (Some(original), Some(shuffled)) = (self.original, self.shuffled) {
+                    let mut product = vec![F::zero(); H];
+                    for (idx, product) in product.iter_mut().enumerate() {
+                        let mut compressed = F::zero();
+                        for value in shuffled.iter() {
+                            compressed *= theta;
+                            compressed += value[idx];
+                        }
+
+                        *product = compressed + gamma
+                    }
+
+                    product.iter_mut().batch_invert();
+
+                    for (idx, product) in product.iter_mut().enumerate() {
+                        let mut compressed = F::zero();
+                        for value in original.iter() {
+                            compressed *= theta;
+                            compressed += value[idx];
+                        }
+
+                        *product *= compressed + gamma
+                    }
+
+                    let z = iter::once(F::one())
+                        .chain(product)
+                        .scan(F::one(), |state, cur| {
+                            *state *= &cur;
+                            Some(*state)
+                        })
+                        .take(H + 1)
+                        .collect::<Vec<_>>();
+
+                    #[cfg(feature = "sanity-checks")]
+                    assert_eq!(F::one(), *z.last().unwrap());
+
+                    Some(z)
+                } else {
+                    None
+                };
+                for offset in 0..H {
+                    region.assign_advice(
+                        || format!("z[{}]", offset),
+                        config.z,
+                        offset,
+                        || z.as_ref().map(|z| z[offset]).ok_or(Error::Synthesis),
+                    )?;
+                }
+
+                Ok(())
+            },
+        )
+    }
+}
+
+fn test_mock_prover<F: FieldExt, const W: usize, const H: usize>(
+    k: u32,
+    circuit: MyCircuit<F, W, H>,
+    expected: Result<(), Vec<(metadata::Constraint, FailureLocation)>>,
+) {
+    let prover = MockProver::<F, false>::run(k, &circuit, vec![]).unwrap();
+    match (prover.verify(), expected) {
+        (Ok(_), Ok(_)) => {}
+        (Err(err), Err(expected)) => {
+            assert_eq!(
+                err.into_iter()
+                    .map(|failure| match failure {
+                        VerifyFailure::ConstraintNotSatisfied {
+                            constraint,
+                            location,
+                            ..
+                        } => (constraint, location),
+                        _ => panic!("MockProver::verify has result unmatching expected"),
+                    })
+                    .collect::<Vec<_>>(),
+                expected
+            )
+        }
+        (_, _) => panic!("MockProver::verify has result unmatching expected"),
+    };
+}
+
+fn test_prover<E: MultiMillerLoop + Debug, const W: usize, const H: usize>(
+    k: u32,
+    circuit: MyCircuit<E::Scalar, W, H>,
+    expected: bool,
+) {
+    let params = ParamsKZG::<E>::new(k);
+    let vk = keygen_vk::<KZGCommitmentScheme<_>, _, false>(&params, &circuit).unwrap();
+    let pk = keygen_pk::<KZGCommitmentScheme<_>, _, false>(&params, vk, &circuit).unwrap();
+
+    let proof = {
+        let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
+
+        create_proof::<_, ProverSHPLONK<_>, _, _, _, _, false>(
+            &params,
+            &pk,
+            &[circuit],
+            &[&[]],
+            OsRng,
+            &mut transcript,
+        )
+        .expect("proof generation should not fail");
+
+        transcript.finalize()
+    };
+
+    let accepted = {
+        let strategy = BatchVerifier::new(&params, OsRng);
+        let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]);
+
+        VerificationStrategy::<_, VerifierSHPLONK<_>, _>::finalize(
+            verify_proof::<_, _, _, VerifierSHPLONK<_>, _, _, false>(
+                &params.verifier_params(),
+                pk.get_vk(),
+                strategy,
+                &[&[]],
+                &mut transcript,
+            )
+            .expect("proof verification should not fail"),
+        )
+    };
+
+    assert_eq!(accepted, expected);
+}
+
+fn main() {
+    const K: u32 = 8;
+    const W: usize = 4;
+    const H: usize = 1 << K;
+
+    let circuit = &MyCircuit::<_, W, H>::rand(&mut OsRng);
+
+    {
+        test_mock_prover(K, circuit.clone(), Ok(()));
+        test_prover::<Bn256, W, H>(K, circuit.clone(), true);
+    }
+
+    #[cfg(not(feature = "sanity-checks"))]
+    {
+        use std::ops::IndexMut;
+
+        let mut circuit = circuit.clone();
+        circuit.shuffled.as_mut().unwrap().index_mut(0).swap(0, 1);
+
+        test_mock_prover(
+            K,
+            circuit.clone(),
+            Err(vec![(
+                ((1, "z should have valid transition").into(), 0, "").into(),
+                FailureLocation::InRegion {
+                    region: (0, "Shuffle original into shuffled").into(),
+                    offset: 255,
+                },
+            )]),
+        );
+        test_prover::<Bn256, W, H>(K, circuit, false);
+    }
+}