Merge branch 'main' into fix_reduction_strategy

# Conflicts:
#	starky/src/prover.rs

field/src/polynomial/mod.rs

@@ -5,6 +5,7 @@ use std::iter::Sum;
 use std::ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign};
 
 use anyhow::{ensure, Result};
+use itertools::Itertools;
 use plonky2_util::log2_strict;
 use serde::{Deserialize, Serialize};
 
@@ -26,8 +27,12 @@ impl<F: Field> PolynomialValues<F> {
         PolynomialValues { values }
     }
 
+    pub fn constant(value: F, len: usize) -> Self {
+        Self::new(vec![value; len])
+    }
+
     pub fn zero(len: usize) -> Self {
-        Self::new(vec![F::ZERO; len])
+        Self::constant(F::ZERO, len)
     }
 
     /// Returns the polynomial whole value is one at the given index, and zero elsewhere.
@@ -83,6 +88,14 @@ impl<F: Field> PolynomialValues<F> {
     pub fn degree_plus_one(&self) -> usize {
         self.clone().ifft().degree_plus_one()
     }
+
+    /// Adds `rhs * rhs_weight` to `self`. Assumes `self.len() == rhs.len()`.
+    pub fn add_assign_scaled(&mut self, rhs: &Self, rhs_weight: F) {
+        self.values
+            .iter_mut()
+            .zip_eq(&rhs.values)
+            .for_each(|(self_v, rhs_v)| *self_v += *rhs_v * rhs_weight)
+    }
 }
 
 impl<F: Field> From<Vec<F>> for PolynomialValues<F> {

starky/src/get_challenges.rs

@@ -11,14 +11,14 @@ use plonky2::plonk::circuit_builder::CircuitBuilder;
 use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 
 use crate::config::StarkConfig;
-use crate::proof::{
-    StarkOpeningSet, StarkOpeningSetTarget, StarkProof, StarkProofChallenges,
-    StarkProofChallengesTarget, StarkProofTarget, StarkProofWithPublicInputs,
-    StarkProofWithPublicInputsTarget,
-};
+use crate::permutation::get_n_permutation_challenge_sets;
+use crate::proof::*;
+use crate::stark::Stark;
 
-fn get_challenges<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>(
+fn get_challenges<F, C, S, const D: usize>(
+    stark: &S,
     trace_cap: &MerkleCap<F, C::Hasher>,
+    permutation_zs_cap: Option<&MerkleCap<F, C::Hasher>>,
     quotient_polys_cap: &MerkleCap<F, C::Hasher>,
     openings: &StarkOpeningSet<F, D>,
     commit_phase_merkle_caps: &[MerkleCap<F, C::Hasher>],
@@ -26,12 +26,33 @@ fn get_challenges<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, cons
     pow_witness: F,
     config: &StarkConfig,
     degree_bits: usize,
-) -> Result<StarkProofChallenges<F, D>> {
+) -> Result<StarkProofChallenges<F, D>>
+where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+{
     let num_challenges = config.num_challenges;
 
     let mut challenger = Challenger::<F, C::Hasher>::new();
 
     challenger.observe_cap(trace_cap);
+
+    let permutation_challenge_sets = if stark.uses_permutation_args() {
+        get_n_permutation_challenge_sets(
+            &mut challenger,
+            num_challenges,
+            stark.permutation_batch_size(),
+        )
+    } else {
+        vec![]
+    };
+    if stark.uses_permutation_args() {
+        let cap =
+            permutation_zs_cap.ok_or_else(|| anyhow::Error::msg("expected permutation_zs_cap"));
+        challenger.observe_cap(cap?);
+    }
+
     let stark_alphas = challenger.get_n_challenges(num_challenges);
 
     challenger.observe_cap(quotient_polys_cap);
@@ -40,6 +61,7 @@ fn get_challenges<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, cons
     challenger.observe_openings(&openings.to_fri_openings());
 
     Ok(StarkProofChallenges {
+        permutation_challenge_sets,
         stark_alphas,
         stark_zeta,
         fri_challenges: challenger.fri_challenges::<C, D>(
@@ -52,28 +74,33 @@ fn get_challenges<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, cons
     })
 }
 
-impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
-    StarkProofWithPublicInputs<F, C, D>
+impl<F, C, const D: usize> StarkProofWithPublicInputs<F, C, D>
+where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
 {
-    pub(crate) fn fri_query_indices(
+    pub(crate) fn fri_query_indices<S: Stark<F, D>>(
         &self,
+        stark: &S,
         config: &StarkConfig,
         degree_bits: usize,
     ) -> anyhow::Result<Vec<usize>> {
         Ok(self
-            .get_challenges(config, degree_bits)?
+            .get_challenges(stark, config, degree_bits)?
             .fri_challenges
             .fri_query_indices)
     }
 
     /// Computes all Fiat-Shamir challenges used in the STARK proof.
-    pub(crate) fn get_challenges(
+    pub(crate) fn get_challenges<S: Stark<F, D>>(
         &self,
+        stark: &S,
         config: &StarkConfig,
         degree_bits: usize,
     ) -> Result<StarkProofChallenges<F, D>> {
         let StarkProof {
             trace_cap,
+            permutation_zs_cap,
             quotient_polys_cap,
             openings,
             opening_proof:
@@ -85,8 +112,10 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
                 },
         } = &self.proof;
 
-        get_challenges::<F, C, D>(
+        get_challenges::<F, C, S, D>(
+            stark,
             trace_cap,
+            permutation_zs_cap.as_ref(),
             quotient_polys_cap,
             openings,
             commit_phase_merkle_caps,
@@ -106,6 +135,7 @@ pub(crate) fn get_challenges_target<
 >(
     builder: &mut CircuitBuilder<F, D>,
     trace_cap: &MerkleCapTarget,
+    permutation_zs_cap: Option<&MerkleCapTarget>,
     quotient_polys_cap: &MerkleCapTarget,
     openings: &StarkOpeningSetTarget<D>,
     commit_phase_merkle_caps: &[MerkleCapTarget],
@@ -152,6 +182,7 @@ impl<const D: usize> StarkProofWithPublicInputsTarget<D> {
     {
         let StarkProofTarget {
             trace_cap,
+            permutation_zs_cap,
             quotient_polys_cap,
             openings,
             opening_proof:
@@ -166,6 +197,7 @@ impl<const D: usize> StarkProofWithPublicInputsTarget<D> {
         get_challenges_target::<F, C, D>(
             builder,
             trace_cap,
+            permutation_zs_cap.as_ref(),
             quotient_polys_cap,
             openings,
             commit_phase_merkle_caps,

starky/src/lib.rs

@@ -8,6 +8,7 @@
 pub mod config;
 pub mod constraint_consumer;
 mod get_challenges;
+pub mod permutation;
 pub mod proof;
 pub mod prover;
 pub mod recursive_verifier;

starky/src/permutation.rs

@@ -0,0 +1,149 @@
+//! Permutation arguments.
+
+use itertools::Itertools;
+use plonky2::field::batch_util::batch_multiply_inplace;
+use plonky2::field::extension_field::Extendable;
+use plonky2::field::field_types::Field;
+use plonky2::field::polynomial::PolynomialValues;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::challenger::Challenger;
+use plonky2::plonk::config::{GenericConfig, Hasher};
+use rayon::prelude::*;
+
+use crate::config::StarkConfig;
+use crate::stark::Stark;
+
+/// A pair of lists of columns, `lhs` and `rhs`, that should be permutations of one another.
+/// In particular, there should exist some permutation `pi` such that for any `i`,
+/// `trace[lhs[i]] = pi(trace[rhs[i]])`. Here `trace` denotes the trace in column-major form, so
+/// `trace[col]` is a column vector.
+pub struct PermutationPair {
+    /// Each entry contains two column indices, representing two columns which should be
+    /// permutations of one another.
+    pub column_pairs: Vec<(usize, usize)>,
+}
+
+/// A single instance of a permutation check protocol.
+pub(crate) struct PermutationInstance<'a, F: Field> {
+    pub(crate) pair: &'a PermutationPair,
+    pub(crate) challenge: PermutationChallenge<F>,
+}
+
+/// Randomness for a single instance of a permutation check protocol.
+#[derive(Copy, Clone)]
+pub(crate) struct PermutationChallenge<F: Field> {
+    /// Randomness used to combine multiple columns into one.
+    pub(crate) beta: F,
+    /// Random offset that's added to the beta-reduced column values.
+    pub(crate) gamma: F,
+}
+
+/// Like `PermutationChallenge`, but with `num_challenges` copies to boost soundness.
+pub(crate) struct PermutationChallengeSet<F: Field> {
+    pub(crate) challenges: Vec<PermutationChallenge<F>>,
+}
+
+/// Compute all Z polynomials (for permutation arguments).
+pub(crate) fn compute_permutation_z_polys<F, C, S, const D: usize>(
+    stark: &S,
+    config: &StarkConfig,
+    challenger: &mut Challenger<F, C::Hasher>,
+    trace_poly_values: &[PolynomialValues<F>],
+) -> Vec<PolynomialValues<F>>
+where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+{
+    let permutation_pairs = stark.permutation_pairs();
+    let permutation_challenge_sets = get_n_permutation_challenge_sets(
+        challenger,
+        config.num_challenges,
+        stark.permutation_batch_size(),
+    );
+
+    // Get a list of instances of our batch-permutation argument. These are permutation arguments
+    // where the same `Z(x)` polynomial is used to check more than one permutation.
+    // Before batching, each permutation pair leads to `num_challenges` permutation arguments, so we
+    // start with the cartesian product of `permutation_pairs` and `0..num_challenges`. Then we
+    // chunk these arguments based on our batch size.
+    let permutation_instances = permutation_pairs
+        .iter()
+        .cartesian_product(0..config.num_challenges)
+        .chunks(stark.permutation_batch_size())
+        .into_iter()
+        .flat_map(|batch| {
+            batch.enumerate().map(|(i, (pair, chal))| {
+                let challenge = permutation_challenge_sets[i].challenges[chal];
+                PermutationInstance { pair, challenge }
+            })
+        })
+        .collect_vec();
+
+    permutation_instances
+        .into_par_iter()
+        .map(|instance| compute_permutation_z_poly(instance, trace_poly_values))
+        .collect()
+}
+
+/// Compute a single Z polynomial.
+// TODO: Change this to handle a batch of `PermutationInstance`s.
+fn compute_permutation_z_poly<F: Field>(
+    instance: PermutationInstance<F>,
+    trace_poly_values: &[PolynomialValues<F>],
+) -> PolynomialValues<F> {
+    let PermutationInstance { pair, challenge } = instance;
+    let PermutationPair { column_pairs } = pair;
+    let PermutationChallenge { beta, gamma } = challenge;
+
+    let degree = trace_poly_values[0].len();
+    let mut reduced_lhs = PolynomialValues::constant(gamma, degree);
+    let mut reduced_rhs = PolynomialValues::constant(gamma, degree);
+
+    for ((lhs, rhs), weight) in column_pairs.iter().zip(beta.powers()) {
+        reduced_lhs.add_assign_scaled(&trace_poly_values[*lhs], weight);
+        reduced_rhs.add_assign_scaled(&trace_poly_values[*rhs], weight);
+    }
+
+    // Compute the quotients.
+    let reduced_rhs_inverses = F::batch_multiplicative_inverse(&reduced_rhs.values);
+    let mut quotients = reduced_lhs.values;
+    batch_multiply_inplace(&mut quotients, &reduced_rhs_inverses);
+
+    // Compute Z, which contains partial products of the quotients.
+    let mut partial_products = Vec::with_capacity(degree);
+    let mut acc = F::ONE;
+    for q in quotients {
+        partial_products.push(acc);
+        acc *= q;
+    }
+    PolynomialValues::new(partial_products)
+}
+
+fn get_permutation_challenge<F: RichField, H: Hasher<F>>(
+    challenger: &mut Challenger<F, H>,
+) -> PermutationChallenge<F> {
+    let beta = challenger.get_challenge();
+    let gamma = challenger.get_challenge();
+    PermutationChallenge { beta, gamma }
+}
+
+fn get_permutation_challenge_set<F: RichField, H: Hasher<F>>(
+    challenger: &mut Challenger<F, H>,
+    num_challenges: usize,
+) -> PermutationChallengeSet<F> {
+    let challenges = (0..num_challenges)
+        .map(|_| get_permutation_challenge(challenger))
+        .collect();
+    PermutationChallengeSet { challenges }
+}
+
+pub(crate) fn get_n_permutation_challenge_sets<F: RichField, H: Hasher<F>>(
+    challenger: &mut Challenger<F, H>,
+    num_challenges: usize,
+    num_sets: usize,
+) -> Vec<PermutationChallengeSet<F>> {
+    (0..num_sets)
+        .map(|_| get_permutation_challenge_set(challenger, num_challenges))
+        .collect()
+}

starky/src/proof.rs

@@ -1,3 +1,4 @@
+use itertools::Itertools;
 use plonky2::field::extension_field::{Extendable, FieldExtension};
 use plonky2::fri::oracle::PolynomialBatch;
 use plonky2::fri::proof::{
@@ -14,11 +15,14 @@ use plonky2::plonk::config::GenericConfig;
 use rayon::prelude::*;
 
 use crate::config::StarkConfig;
+use crate::permutation::PermutationChallengeSet;
 
 #[derive(Debug, Clone)]
 pub struct StarkProof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> {
     /// Merkle cap of LDEs of trace values.
     pub trace_cap: MerkleCap<F, C::Hasher>,
+    /// Merkle cap of LDEs of permutation Z values.
+    pub permutation_zs_cap: Option<MerkleCap<F, C::Hasher>>,
     /// Merkle cap of LDEs of trace values.
     pub quotient_polys_cap: MerkleCap<F, C::Hasher>,
     /// Purported values of each polynomial at the challenge point.
@@ -40,6 +44,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> S
 
 pub struct StarkProofTarget<const D: usize> {
     pub trace_cap: MerkleCapTarget,
+    pub permutation_zs_cap: Option<MerkleCapTarget>,
     pub quotient_polys_cap: MerkleCapTarget,
     pub openings: StarkOpeningSetTarget<D>,
     pub opening_proof: FriProofTarget<D>,
@@ -95,6 +100,9 @@ pub struct CompressedStarkProofWithPublicInputs<
 }
 
 pub(crate) struct StarkProofChallenges<F: RichField + Extendable<D>, const D: usize> {
+    /// Randomness used in any permutation arguments.
+    pub permutation_challenge_sets: Vec<PermutationChallengeSet<F>>,
+
     /// Random values used to combine STARK constraints.
     pub stark_alphas: Vec<F>,
 
@@ -115,8 +123,8 @@ pub(crate) struct StarkProofChallengesTarget<const D: usize> {
 pub struct StarkOpeningSet<F: RichField + Extendable<D>, const D: usize> {
     pub local_values: Vec<F::Extension>,
     pub next_values: Vec<F::Extension>,
-    pub permutation_zs: Vec<F::Extension>,
-    pub permutation_zs_right: Vec<F::Extension>,
+    pub permutation_zs: Option<Vec<F::Extension>>,
+    pub permutation_zs_right: Option<Vec<F::Extension>>,
     pub quotient_polys: Vec<F::Extension>,
 }
 
@@ -125,6 +133,7 @@ impl<F: RichField + Extendable<D>, const D: usize> StarkOpeningSet<F, D> {
         zeta: F::Extension,
         g: F,
         trace_commitment: &PolynomialBatch<F, C, D>,
+        permutation_zs_commitment: Option<&PolynomialBatch<F, C, D>>,
         quotient_commitment: &PolynomialBatch<F, C, D>,
     ) -> Self {
         let eval_commitment = |z: F::Extension, c: &PolynomialBatch<F, C, D>| {
@@ -133,30 +142,33 @@ impl<F: RichField + Extendable<D>, const D: usize> StarkOpeningSet<F, D> {
                 .map(|p| p.to_extension().eval(z))
                 .collect::<Vec<_>>()
         };
+        let zeta_right = zeta.scalar_mul(g);
         Self {
             local_values: eval_commitment(zeta, trace_commitment),
-            next_values: eval_commitment(zeta.scalar_mul(g), trace_commitment),
-            permutation_zs: vec![/*TODO*/],
-            permutation_zs_right: vec![/*TODO*/],
+            next_values: eval_commitment(zeta_right, trace_commitment),
+            permutation_zs: permutation_zs_commitment.map(|c| eval_commitment(zeta, c)),
+            permutation_zs_right: permutation_zs_commitment.map(|c| eval_commitment(zeta_right, c)),
             quotient_polys: eval_commitment(zeta, quotient_commitment),
         }
     }
 
     pub(crate) fn to_fri_openings(&self) -> FriOpenings<F, D> {
         let zeta_batch = FriOpeningBatch {
-            values: [
-                self.local_values.as_slice(),
-                self.quotient_polys.as_slice(),
-                self.permutation_zs.as_slice(),
-            ]
-            .concat(),
+            values: self
+                .local_values
+                .iter()
+                .chain(self.permutation_zs.iter().flatten())
+                .chain(&self.quotient_polys)
+                .copied()
+                .collect_vec(),
         };
         let zeta_right_batch = FriOpeningBatch {
-            values: [
-                self.next_values.as_slice(),
-                self.permutation_zs_right.as_slice(),
-            ]
-            .concat(),
+            values: self
+                .next_values
+                .iter()
+                .chain(self.permutation_zs_right.iter().flatten())
+                .copied()
+                .collect_vec(),
         };
         FriOpenings {
             batches: vec![zeta_batch, zeta_right_batch],

starky/src/prover.rs

@@ -1,3 +1,5 @@
+use std::iter::once;
+
 use anyhow::{ensure, Result};
 use itertools::Itertools;
 use plonky2::field::extension_field::Extendable;
@@ -16,6 +18,7 @@ use rayon::prelude::*;
 
 use crate::config::StarkConfig;
 use crate::constraint_consumer::ConstraintConsumer;
+use crate::permutation::compute_permutation_z_polys;
 use crate::proof::{StarkOpeningSet, StarkProof, StarkProofWithPublicInputs};
 use crate::stark::Stark;
 use crate::vars::StarkEvaluationVars;
@@ -43,7 +46,7 @@ where
         "FRI total reduction arity is too large.",
     );
 
-    let trace_vecs = trace.into_iter().map(|row| row.to_vec()).collect_vec();
+    let trace_vecs = trace.iter().map(|row| row.to_vec()).collect_vec();
     let trace_col_major: Vec<Vec<F>> = transpose(&trace_vecs);
 
     let trace_poly_values: Vec<PolynomialValues<F>> = timed!(
@@ -61,7 +64,9 @@ where
         timing,
         "compute trace commitment",
         PolynomialBatch::<F, C, D>::from_values(
-            trace_poly_values,
+            // TODO: Cloning this isn't great; consider having `from_values` accept a reference,
+            // or having `compute_permutation_z_polys` read trace values from the `PolynomialBatch`.
+            trace_poly_values.clone(),
             rate_bits,
             false,
             cap_height,
@@ -74,6 +79,36 @@ where
     let mut challenger = Challenger::new();
     challenger.observe_cap(&trace_cap);
 
+    // Permutation arguments.
+    let permutation_zs_commitment = if stark.uses_permutation_args() {
+        let permutation_z_polys = compute_permutation_z_polys::<F, C, S, D>(
+            &stark,
+            config,
+            &mut challenger,
+            &trace_poly_values,
+        );
+        timed!(
+            timing,
+            "compute permutation Z commitments",
+            Some(PolynomialBatch::from_values(
+                permutation_z_polys,
+                rate_bits,
+                false,
+                config.fri_config.cap_height,
+                timing,
+                None,
+            ))
+        )
+    } else {
+        None
+    };
+    let permutation_zs_cap = permutation_zs_commitment
+        .as_ref()
+        .map(|commit| commit.merkle_tree.cap.clone());
+    for cap in &permutation_zs_cap {
+        challenger.observe_cap(cap);
+    }
+
     let alphas = challenger.get_n_challenges(config.num_challenges);
     let quotient_polys = compute_quotient_polys::<F, C, S, D>(
         &stark,
@@ -117,18 +152,26 @@ where
         zeta.exp_power_of_2(degree_bits) != F::Extension::ONE,
         "Opening point is in the subgroup."
     );
-    let openings = StarkOpeningSet::new(zeta, g, &trace_commitment, &quotient_commitment);
+    let openings = StarkOpeningSet::new(
+        zeta,
+        g,
+        &trace_commitment,
+        permutation_zs_commitment.as_ref(),
+        &quotient_commitment,
+    );
     challenger.observe_openings(&openings.to_fri_openings());
 
-    // TODO: Add permutation checks
-    let initial_merkle_trees = &[&trace_commitment, &quotient_commitment];
+    let initial_merkle_trees = once(&trace_commitment)
+        .chain(permutation_zs_commitment.as_ref())
+        .chain(once(&quotient_commitment))
+        .collect_vec();
 
     let opening_proof = timed!(
         timing,
         "compute openings proof",
         PolynomialBatch::prove_openings(
-            &stark.fri_instance(zeta, g, config.num_challenges),
-            initial_merkle_trees,
+            &stark.fri_instance(zeta, g, config),
+            &initial_merkle_trees,
             &mut challenger,
             &fri_params,
             timing,
@@ -136,6 +179,7 @@ where
     );
     let proof = StarkProof {
         trace_cap,
+        permutation_zs_cap,
         quotient_polys_cap,
         openings,
         opening_proof,
@@ -212,6 +256,7 @@ where
                 public_inputs: &public_inputs,
             };
             stark.eval_packed_base(vars, &mut consumer);
+            // TODO: Add in constraints for permutation arguments.
             // TODO: Fix this once we use a genuine `PackedField`.
             let mut constraints_evals = consumer.accumulators();
             // We divide the constraints evaluations by `Z_H(x)`.

starky/src/recursive_verifier.rs

@@ -1,3 +1,5 @@
+use std::iter::once;
+
 use itertools::Itertools;
 use plonky2::field::extension_field::Extendable;
 use plonky2::field::field_types::Field;
@@ -103,6 +105,7 @@ fn recursively_verify_stark_proof_with_challenges<
         l_last,
     );
     stark.eval_ext_recursively(builder, vars, &mut consumer);
+    // TODO: Add in constraints for permutation arguments.
     let vanishing_polys_zeta = consumer.accumulators();
 
     // Check each polynomial identity, of the form `vanishing(x) = Z_H(x) quotient(x)`, at zeta.
@@ -117,20 +120,22 @@ fn recursively_verify_stark_proof_with_challenges<
         builder.connect_extension(vanishing_polys_zeta[i], computed_vanishing_poly);
     }
 
-    // TODO: Permutation polynomials.
-    let merkle_caps = &[proof.trace_cap, proof.quotient_polys_cap];
+    let merkle_caps = once(proof.trace_cap)
+        .chain(proof.permutation_zs_cap)
+        .chain(once(proof.quotient_polys_cap))
+        .collect_vec();
 
     let fri_instance = stark.fri_instance_target(
         builder,
         challenges.stark_zeta,
         F::primitive_root_of_unity(degree_bits),
-        inner_config.num_challenges,
+        inner_config,
     );
     builder.verify_fri_proof::<C>(
         &fri_instance,
         &proof.openings.to_fri_openings(),
         &challenges.fri_challenges,
-        merkle_caps,
+        &merkle_caps,
         &proof.opening_proof,
         &inner_config.fri_params(degree_bits),
     );
@@ -188,8 +193,15 @@ pub fn add_virtual_stark_proof<F: RichField + Extendable<D>, S: Stark<F, D>, con
         stark.quotient_degree_factor() * config.num_challenges,
     ];
 
+    let permutation_zs_cap = if stark.uses_permutation_args() {
+        Some(builder.add_virtual_cap(cap_height))
+    } else {
+        None
+    };
+
     StarkProofTarget {
         trace_cap: builder.add_virtual_cap(cap_height),
+        permutation_zs_cap,
         quotient_polys_cap: builder.add_virtual_cap(cap_height),
         openings: add_stark_opening_set::<F, S, D>(builder, stark, config),
         opening_proof: builder.add_virtual_fri_proof(num_leaves_per_oracle, &fri_params),

starky/src/stark.rs

@@ -7,8 +7,11 @@ use plonky2::fri::structure::{
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::ext_target::ExtensionTarget;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2_util::ceil_div_usize;
 
+use crate::config::StarkConfig;
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+use crate::permutation::PermutationPair;
 use crate::vars::StarkEvaluationTargets;
 use crate::vars::StarkEvaluationVars;
 
@@ -75,56 +78,126 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
     }
 
     /// Computes the FRI instance used to prove this Stark.
-    // TODO: Permutation polynomials.
     fn fri_instance(
         &self,
         zeta: F::Extension,
         g: F,
-        num_challenges: usize,
+        config: &StarkConfig,
     ) -> FriInstanceInfo<F, D> {
         let no_blinding_oracle = FriOracleInfo { blinding: false };
-        let trace_info = FriPolynomialInfo::from_range(0, 0..Self::COLUMNS);
-        let quotient_info =
-            FriPolynomialInfo::from_range(1, 0..self.quotient_degree_factor() * num_challenges);
+        let mut oracle_indices = 0..;
+
+        let trace_info =
+            FriPolynomialInfo::from_range(oracle_indices.next().unwrap(), 0..Self::COLUMNS);
+
+        let permutation_zs_info = if self.uses_permutation_args() {
+            FriPolynomialInfo::from_range(
+                oracle_indices.next().unwrap(),
+                0..self.num_permutation_batches(config),
+            )
+        } else {
+            vec![]
+        };
+
+        let quotient_info = FriPolynomialInfo::from_range(
+            oracle_indices.next().unwrap(),
+            0..self.quotient_degree_factor() * config.num_challenges,
+        );
+
         let zeta_batch = FriBatchInfo {
             point: zeta,
-            polynomials: [trace_info.clone(), quotient_info].concat(),
+            polynomials: [
+                trace_info.clone(),
+                permutation_zs_info.clone(),
+                quotient_info,
+            ]
+            .concat(),
         };
-        let zeta_right_batch = FriBatchInfo::<F, D> {
+        let zeta_right_batch = FriBatchInfo {
             point: zeta.scalar_mul(g),
-            polynomials: trace_info,
+            polynomials: [trace_info, permutation_zs_info].concat(),
         };
         FriInstanceInfo {
-            oracles: vec![no_blinding_oracle; 3],
+            oracles: vec![no_blinding_oracle; oracle_indices.next().unwrap()],
             batches: vec![zeta_batch, zeta_right_batch],
         }
     }
 
     /// Computes the FRI instance used to prove this Stark.
-    // TODO: Permutation polynomials.
     fn fri_instance_target(
         &self,
         builder: &mut CircuitBuilder<F, D>,
         zeta: ExtensionTarget<D>,
         g: F,
-        num_challenges: usize,
+        config: &StarkConfig,
     ) -> FriInstanceInfoTarget<D> {
         let no_blinding_oracle = FriOracleInfo { blinding: false };
-        let trace_info = FriPolynomialInfo::from_range(0, 0..Self::COLUMNS);
-        let quotient_info =
-            FriPolynomialInfo::from_range(1, 0..self.quotient_degree_factor() * num_challenges);
+        let mut oracle_indices = 0..;
+
+        let trace_info =
+            FriPolynomialInfo::from_range(oracle_indices.next().unwrap(), 0..Self::COLUMNS);
+
+        let permutation_zs_info = if self.uses_permutation_args() {
+            FriPolynomialInfo::from_range(
+                oracle_indices.next().unwrap(),
+                0..self.num_permutation_batches(config),
+            )
+        } else {
+            vec![]
+        };
+
+        let quotient_info = FriPolynomialInfo::from_range(
+            oracle_indices.next().unwrap(),
+            0..self.quotient_degree_factor() * config.num_challenges,
+        );
+
         let zeta_batch = FriBatchInfoTarget {
             point: zeta,
-            polynomials: [trace_info.clone(), quotient_info].concat(),
+            polynomials: [
+                trace_info.clone(),
+                permutation_zs_info.clone(),
+                quotient_info,
+            ]
+            .concat(),
         };
         let zeta_right = builder.mul_const_extension(g, zeta);
         let zeta_right_batch = FriBatchInfoTarget {
             point: zeta_right,
-            polynomials: trace_info,
+            polynomials: [trace_info, permutation_zs_info].concat(),
         };
         FriInstanceInfoTarget {
-            oracles: vec![no_blinding_oracle; 3],
+            oracles: vec![no_blinding_oracle; oracle_indices.next().unwrap()],
             batches: vec![zeta_batch, zeta_right_batch],
         }
     }
+
+    /// Pairs of lists of columns that should be permutations of one another. A permutation argument
+    /// will be used for each such pair. Empty by default.
+    fn permutation_pairs(&self) -> Vec<PermutationPair> {
+        vec![]
+    }
+
+    fn uses_permutation_args(&self) -> bool {
+        !self.permutation_pairs().is_empty()
+    }
+
+    /// The number of permutation argument instances that can be combined into a single constraint.
+    fn permutation_batch_size(&self) -> usize {
+        // The permutation argument constraints look like
+        //     Z(x) \prod(...) = Z(g x) \prod(...)
+        // where each product has a number of terms equal to the batch size. So our batch size
+        // should be one less than our constraint degree, which happens to be our quotient degree.
+        self.quotient_degree_factor()
+    }
+
+    fn num_permutation_instances(&self, config: &StarkConfig) -> usize {
+        self.permutation_pairs().len() * config.num_challenges
+    }
+
+    fn num_permutation_batches(&self, config: &StarkConfig) -> usize {
+        ceil_div_usize(
+            self.num_permutation_instances(config),
+            self.permutation_batch_size(),
+        )
+    }
 }

starky/src/verifier.rs

@@ -1,4 +1,7 @@
+use std::iter::once;
+
 use anyhow::{ensure, Result};
+use itertools::Itertools;
 use plonky2::field::extension_field::{Extendable, FieldExtension};
 use plonky2::field::field_types::Field;
 use plonky2::fri::verifier::verify_fri_proof;
@@ -29,7 +32,7 @@ where
 {
     ensure!(proof_with_pis.public_inputs.len() == S::PUBLIC_INPUTS);
     let degree_bits = proof_with_pis.proof.recover_degree_bits(config);
-    let challenges = proof_with_pis.get_challenges(config, degree_bits)?;
+    let challenges = proof_with_pis.get_challenges(&stark, config, degree_bits)?;
     verify_stark_proof_with_challenges(stark, proof_with_pis, challenges, degree_bits, config)
 }
 
@@ -86,6 +89,7 @@ where
         l_last,
     );
     stark.eval_ext(vars, &mut consumer);
+    // TODO: Add in constraints for permutation arguments.
     let vanishing_polys_zeta = consumer.accumulators();
 
     // Check each polynomial identity, of the form `vanishing(x) = Z_H(x) quotient(x)`, at zeta.
@@ -104,18 +108,20 @@ where
         ensure!(vanishing_polys_zeta[i] == z_h_zeta * reduce_with_powers(chunk, zeta_pow_deg));
     }
 
-    // TODO: Permutation polynomials.
-    let merkle_caps = &[proof.trace_cap, proof.quotient_polys_cap];
+    let merkle_caps = once(proof.trace_cap)
+        .chain(proof.permutation_zs_cap)
+        .chain(once(proof.quotient_polys_cap))
+        .collect_vec();
 
     verify_fri_proof::<F, C, D>(
         &stark.fri_instance(
             challenges.stark_zeta,
             F::primitive_root_of_unity(degree_bits),
-            config.num_challenges,
+            config,
         ),
         &proof.openings.to_fri_openings(),
         &challenges.fri_challenges,
-        merkle_caps,
+        &merkle_caps,
         &proof.opening_proof,
         &config.fri_params(degree_bits),
     )?;

system_zero/src/system_zero.rs

@@ -5,6 +5,7 @@ use plonky2::field::packed_field::PackedField;
 use plonky2::hash::hash_types::RichField;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+use starky::permutation::PermutationPair;
 use starky::stark::Stark;
 use starky::vars::StarkEvaluationTargets;
 use starky::vars::StarkEvaluationVars;
@@ -103,6 +104,12 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for SystemZero<F,
     fn constraint_degree(&self) -> usize {
         3
     }
+
+    fn permutation_pairs(&self) -> Vec<PermutationPair> {
+        // TODO: Add permutation pairs for memory.
+        // TODO: Add permutation pairs for range checks.
+        vec![]
+    }
 }
 
 #[cfg(test)]