Merge pull request #451 from mir-protocol/start_stark_prover

Start of STARK prover

field/src/lib.rs

@@ -24,6 +24,7 @@ pub mod packed_field;
 pub mod polynomial;
 pub mod secp256k1_base;
 pub mod secp256k1_scalar;
+pub mod zero_poly_coset;
 
 #[cfg(test)]
 mod field_testing;

field/src/zero_poly_coset.rs

@@ -0,0 +1,47 @@
+use crate::field_types::Field;
+
+/// Precomputations of the evaluation of `Z_H(X) = X^n - 1` on a coset `gK` with `H <= K`.
+pub struct ZeroPolyOnCoset<F: Field> {
+    /// `n = |H|`.
+    n: F,
+    /// `rate = |K|/|H|`.
+    rate: usize,
+    /// Holds `g^n * (w^n)^i - 1 = g^n * v^i - 1` for `i in 0..rate`, with `w` a generator of `K` and `v` a
+    /// `rate`-primitive root of unity.
+    evals: Vec<F>,
+    /// Holds the multiplicative inverses of `evals`.
+    inverses: Vec<F>,
+}
+
+impl<F: Field> ZeroPolyOnCoset<F> {
+    pub fn new(n_log: usize, rate_bits: usize) -> Self {
+        let g_pow_n = F::coset_shift().exp_power_of_2(n_log);
+        let evals = F::two_adic_subgroup(rate_bits)
+            .into_iter()
+            .map(|x| g_pow_n * x - F::ONE)
+            .collect::<Vec<_>>();
+        let inverses = F::batch_multiplicative_inverse(&evals);
+        Self {
+            n: F::from_canonical_usize(1 << n_log),
+            rate: 1 << rate_bits,
+            evals,
+            inverses,
+        }
+    }
+
+    /// Returns `Z_H(g * w^i)`.
+    pub fn eval(&self, i: usize) -> F {
+        self.evals[i % self.rate]
+    }
+
+    /// Returns `1 / Z_H(g * w^i)`.
+    pub fn eval_inverse(&self, i: usize) -> F {
+        self.inverses[i % self.rate]
+    }
+
+    /// Returns `L_1(x) = Z_H(x)/(n * (x - 1))` with `x = w^i`.
+    pub fn eval_l1(&self, i: usize, x: F) -> F {
+        // Could also precompute the inverses using Montgomery.
+        self.eval(i) * (self.n * (x - F::ONE)).inverse()
+    }
+}

plonky2/src/fri/oracle.rs

@@ -127,7 +127,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
     }
 
     /// Produces a batch opening proof.
-    pub(crate) fn prove_openings(
+    pub fn prove_openings(
         instance: &FriInstanceInfo<F, D>,
         oracles: &[&Self],
         challenger: &mut Challenger<F, C::Hasher>,

plonky2/src/plonk/plonk_common.rs

@@ -63,52 +63,6 @@ pub(crate) fn eval_zero_poly<F: Field>(n: usize, x: F) -> F {
     x.exp_u64(n as u64) - F::ONE
 }
 
-/// Precomputations of the evaluation of `Z_H(X) = X^n - 1` on a coset `gK` with `H <= K`.
-pub(crate) struct ZeroPolyOnCoset<F: Field> {
-    /// `n = |H|`.
-    n: F,
-    /// `rate = |K|/|H|`.
-    rate: usize,
-    /// Holds `g^n * (w^n)^i - 1 = g^n * v^i - 1` for `i in 0..rate`, with `w` a generator of `K` and `v` a
-    /// `rate`-primitive root of unity.
-    evals: Vec<F>,
-    /// Holds the multiplicative inverses of `evals`.
-    inverses: Vec<F>,
-}
-
-impl<F: Field> ZeroPolyOnCoset<F> {
-    pub fn new(n_log: usize, rate_bits: usize) -> Self {
-        let g_pow_n = F::coset_shift().exp_power_of_2(n_log);
-        let evals = F::two_adic_subgroup(rate_bits)
-            .into_iter()
-            .map(|x| g_pow_n * x - F::ONE)
-            .collect::<Vec<_>>();
-        let inverses = F::batch_multiplicative_inverse(&evals);
-        Self {
-            n: F::from_canonical_usize(1 << n_log),
-            rate: 1 << rate_bits,
-            evals,
-            inverses,
-        }
-    }
-
-    /// Returns `Z_H(g * w^i)`.
-    pub fn eval(&self, i: usize) -> F {
-        self.evals[i % self.rate]
-    }
-
-    /// Returns `1 / Z_H(g * w^i)`.
-    pub fn eval_inverse(&self, i: usize) -> F {
-        self.inverses[i % self.rate]
-    }
-
-    /// Returns `L_1(x) = Z_H(x)/(n * (x - 1))` with `x = w^i`.
-    pub fn eval_l1(&self, i: usize, x: F) -> F {
-        // Could also precompute the inverses using Montgomery.
-        self.eval(i) * (self.n * (x - F::ONE)).inverse()
-    }
-}
-
 /// Evaluate the Lagrange basis `L_1` with `L_1(1) = 1`, and `L_1(x) = 0` for other members of an
 /// order `n` multiplicative subgroup.
 pub(crate) fn eval_l_1<F: Field>(n: usize, x: F) -> F {

plonky2/src/plonk/prover.rs

@@ -4,6 +4,7 @@ use anyhow::ensure;
 use anyhow::Result;
 use plonky2_field::extension_field::Extendable;
 use plonky2_field::polynomial::{PolynomialCoeffs, PolynomialValues};
+use plonky2_field::zero_poly_coset::ZeroPolyOnCoset;
 use plonky2_util::log2_ceil;
 use rayon::prelude::*;
 
@@ -15,7 +16,7 @@ use crate::iop::generator::generate_partial_witness;
 use crate::iop::witness::{MatrixWitness, PartialWitness, Witness};
 use crate::plonk::circuit_data::{CommonCircuitData, ProverOnlyCircuitData};
 use crate::plonk::config::{GenericConfig, Hasher};
-use crate::plonk::plonk_common::{PlonkOracle, ZeroPolyOnCoset};
+use crate::plonk::plonk_common::PlonkOracle;
 use crate::plonk::proof::OpeningSet;
 use crate::plonk::proof::{Proof, ProofWithPublicInputs};
 use crate::plonk::vanishing_poly::eval_vanishing_poly_base_batch;

plonky2/src/plonk/vanishing_poly.rs

@@ -1,6 +1,7 @@
 use plonky2_field::batch_util::batch_add_inplace;
 use plonky2_field::extension_field::{Extendable, FieldExtension};
 use plonky2_field::field_types::Field;
+use plonky2_field::zero_poly_coset::ZeroPolyOnCoset;
 
 use crate::gates::gate::PrefixedGate;
 use crate::hash::hash_types::RichField;
@@ -10,7 +11,7 @@ use crate::plonk::circuit_builder::CircuitBuilder;
 use crate::plonk::circuit_data::CommonCircuitData;
 use crate::plonk::config::GenericConfig;
 use crate::plonk::plonk_common;
-use crate::plonk::plonk_common::{eval_l_1_recursively, ZeroPolyOnCoset};
+use crate::plonk::plonk_common::eval_l_1_recursively;
 use crate::plonk::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBaseBatch};
 use crate::util::partial_products::{check_partial_products, check_partial_products_recursively};
 use crate::util::reducing::ReducingFactorTarget;

starky/src/constraint_consumer.rs

@@ -8,26 +8,45 @@ use plonky2::iop::target::Target;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 
 pub struct ConstraintConsumer<P: PackedField> {
-    /// A random value used to combine multiple constraints into one.
+    /// Random values used to combine multiple constraints into one.
-    alpha: P::Scalar,
+    alphas: Vec<P::Scalar>,
 
-    /// A running sum of constraints that have been emitted so far, scaled by powers of alpha.
+    /// Running sums of constraints that have been emitted so far, scaled by powers of alpha.
-    constraint_acc: P,
+    constraint_accs: Vec<P>,
 
     /// The evaluation of the Lagrange basis polynomial which is nonzero at the point associated
     /// with the first trace row, and zero at other points in the subgroup.
-    lagrange_basis_first: P::Scalar,
+    lagrange_basis_first: P,
 
     /// The evaluation of the Lagrange basis polynomial which is nonzero at the point associated
     /// with the last trace row, and zero at other points in the subgroup.
-    lagrange_basis_last: P::Scalar,
+    lagrange_basis_last: P,
 }
 
 impl<P: PackedField> ConstraintConsumer<P> {
+    pub fn new(alphas: Vec<P::Scalar>, lagrange_basis_first: P, lagrange_basis_last: P) -> Self {
+        Self {
+            constraint_accs: vec![P::ZEROS; alphas.len()],
+            alphas,
+            lagrange_basis_first,
+            lagrange_basis_last,
+        }
+    }
+
+    // TODO: Do this correctly.
+    pub fn accumulators(self) -> Vec<P::Scalar> {
+        self.constraint_accs
+            .into_iter()
+            .map(|acc| acc.as_slice()[0])
+            .collect()
+    }
+
     /// Add one constraint.
     pub fn one(&mut self, constraint: P) {
-        self.constraint_acc *= self.alpha;
+        for (&alpha, acc) in self.alphas.iter().zip(&mut self.constraint_accs) {
-        self.constraint_acc += constraint;
+            *acc *= alpha;
+            *acc += constraint;
+        }
     }
 
     /// Add a series of constraints.

starky/src/fibonacci_stark.rs

@@ -0,0 +1,123 @@
+use std::marker::PhantomData;
+
+use plonky2::field::extension_field::{Extendable, FieldExtension};
+use plonky2::field::packed_field::PackedField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+
+use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+use crate::stark::Stark;
+use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
+
+/// Toy STARK system used for testing.
+/// Computes a Fibonacci sequence with state `[x0, x1]` using the state transition
+/// `x0 <- x1, x1 <- x0 + x1`.
+struct FibonacciStark<F: RichField + Extendable<D>, const D: usize> {
+    num_rows: usize,
+    _phantom: PhantomData<F>,
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> FibonacciStark<F, D> {
+    // The first public input is `x0`.
+    const PI_INDEX_X0: usize = 0;
+    // The second public input is `x1`.
+    const PI_INDEX_X1: usize = 1;
+    // The third public input is the second element of the last row, which should be equal to the
+    // `num_rows`-th Fibonacci number.
+    const PI_INDEX_RES: usize = 2;
+
+    fn new(num_rows: usize) -> Self {
+        Self {
+            num_rows,
+            _phantom: PhantomData,
+        }
+    }
+
+    /// Generate the trace using `x0, x1` as inital state values.
+    fn generate_trace(&self, x0: F, x1: F) -> Vec<[F; Self::COLUMNS]> {
+        (0..self.num_rows)
+            .scan([x0, x1], |acc, _| {
+                let tmp = *acc;
+                acc[0] = tmp[1];
+                acc[1] = tmp[0] + tmp[1];
+                Some(tmp)
+            })
+            .collect()
+    }
+}
+
+impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for FibonacciStark<F, D> {
+    const COLUMNS: usize = 2;
+    const PUBLIC_INPUTS: usize = 3;
+
+    fn eval_packed_generic<FE, P, const D2: usize>(
+        &self,
+        vars: StarkEvaluationVars<FE, P, { Self::COLUMNS }, { Self::PUBLIC_INPUTS }>,
+        yield_constr: &mut ConstraintConsumer<P>,
+    ) where
+        FE: FieldExtension<D2, BaseField = F>,
+        P: PackedField<Scalar = FE>,
+    {
+        // Check public inputs.
+        yield_constr.one_first_row(vars.local_values[0] - vars.public_inputs[Self::PI_INDEX_X0]);
+        yield_constr.one_first_row(vars.local_values[1] - vars.public_inputs[Self::PI_INDEX_X1]);
+        yield_constr.one_last_row(vars.local_values[1] - vars.public_inputs[Self::PI_INDEX_RES]);
+
+        // x0 <- x1
+        yield_constr.one(vars.next_values[0] - vars.local_values[1]);
+        // x1 <- x0 + x1
+        yield_constr.one(vars.next_values[1] - vars.local_values[0] - vars.local_values[1]);
+    }
+
+    fn eval_ext_recursively(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        vars: StarkEvaluationTargets<D, { Self::COLUMNS }, { Self::PUBLIC_INPUTS }>,
+        yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+    ) {
+        todo!()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use plonky2::field::field_types::Field;
+    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+    use plonky2::util::timing::TimingTree;
+
+    use crate::config::StarkConfig;
+    use crate::fibonacci_stark::FibonacciStark;
+    use crate::prover::prove;
+
+    fn fibonacci(n: usize, x0: usize, x1: usize) -> usize {
+        (0..n).fold((0, 1), |x, _| (x.1, x.0 + x.1)).1
+    }
+
+    #[test]
+    fn test_fibonacci_stark() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type S = FibonacciStark<F, D>;
+
+        let config = StarkConfig::standard_fast_config();
+        let num_rows = 1 << 5;
+        let public_inputs = [
+            F::ZERO,
+            F::ONE,
+            F::from_canonical_usize(fibonacci(num_rows - 1, 0, 1)),
+        ];
+        let stark = S::new(num_rows);
+        let trace = stark.generate_trace(public_inputs[0], public_inputs[1]);
+        prove::<F, C, S, D>(
+            stark,
+            config,
+            trace,
+            public_inputs,
+            &mut TimingTree::default(),
+        )?;
+
+        Ok(())
+    }
+}

starky/src/lib.rs

@@ -12,3 +12,6 @@ pub mod proof;
 pub mod prover;
 pub mod stark;
 pub mod vars;
+
+#[cfg(test)]
+pub mod fibonacci_stark;

starky/src/proof.rs

@@ -1,8 +1,10 @@
 use plonky2::field::extension_field::Extendable;
+use plonky2::fri::oracle::PolynomialBatch;
 use plonky2::fri::proof::{CompressedFriProof, FriProof};
 use plonky2::hash::hash_types::RichField;
 use plonky2::hash::merkle_tree::MerkleCap;
 use plonky2::plonk::config::GenericConfig;
+use rayon::prelude::*;
 
 pub struct StarkProof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> {
     /// Merkle cap of LDEs of trace values.
@@ -33,3 +35,25 @@ pub struct StarkOpeningSet<F: RichField + Extendable<D>, const D: usize> {
     pub permutation_zs: Vec<F::Extension>,
     pub quotient_polys: Vec<F::Extension>,
 }
+
+impl<F: RichField + Extendable<D>, const D: usize> StarkOpeningSet<F, D> {
+    pub fn new<C: GenericConfig<D, F = F>>(
+        zeta: F::Extension,
+        g: F::Extension,
+        trace_commitment: &PolynomialBatch<F, C, D>,
+        quotient_commitment: &PolynomialBatch<F, C, D>,
+    ) -> Self {
+        let eval_commitment = |z: F::Extension, c: &PolynomialBatch<F, C, D>| {
+            c.polynomials
+                .par_iter()
+                .map(|p| p.to_extension().eval(z))
+                .collect::<Vec<_>>()
+        };
+        Self {
+            local_values: eval_commitment(zeta, trace_commitment),
+            next_values: eval_commitment(zeta * g, trace_commitment),
+            permutation_zs: vec![/*TODO*/],
+            quotient_polys: eval_commitment(zeta, quotient_commitment),
+        }
+    }
+}

starky/src/prover.rs

@@ -1,8 +1,10 @@
+use anyhow::{ensure, Result};
 use itertools::Itertools;
 use plonky2::field::extension_field::Extendable;
-use plonky2::field::polynomial::PolynomialValues;
+use plonky2::field::field_types::Field;
+use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
+use plonky2::field::zero_poly_coset::ZeroPolyOnCoset;
 use plonky2::fri::oracle::PolynomialBatch;
-use plonky2::fri::prover::fri_proof;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::challenger::Challenger;
 use plonky2::plonk::config::GenericConfig;
@@ -13,22 +15,28 @@ use plonky2_util::log2_strict;
 use rayon::prelude::*;
 
 use crate::config::StarkConfig;
-use crate::proof::StarkProof;
+use crate::constraint_consumer::ConstraintConsumer;
+use crate::proof::{StarkOpeningSet, StarkProof};
 use crate::stark::Stark;
+use crate::vars::StarkEvaluationVars;
 
+// TODO: Deal with public inputs.
 pub fn prove<F, C, S, const D: usize>(
     stark: S,
     config: StarkConfig,
     trace: Vec<[F; S::COLUMNS]>,
+    public_inputs: [F; S::PUBLIC_INPUTS],
     timing: &mut TimingTree,
-) -> StarkProof<F, C, D>
+) -> Result<StarkProof<F, C, D>>
 where
     F: RichField + Extendable<D>,
     C: GenericConfig<D, F = F>,
     S: Stark<F, D>,
     [(); S::COLUMNS]:,
+    [(); S::PUBLIC_INPUTS]:,
 {
-    let degree_bits = log2_strict(trace.len());
+    let degree = trace.len();
+    let degree_bits = log2_strict(degree);
 
     let trace_vecs = trace.into_iter().map(|row| row.to_vec()).collect_vec();
     let trace_col_major: Vec<Vec<F>> = transpose(&trace_vecs);
@@ -57,27 +65,148 @@ where
         )
     );
 
-    let trace_cap = trace_commitment.merkle_tree.cap;
+    let trace_cap = trace_commitment.merkle_tree.cap.clone();
-    let openings = todo!();
-
-    let initial_merkle_trees = todo!();
-    let lde_polynomial_coeffs = todo!();
-    let lde_polynomial_values = todo!();
     let mut challenger = Challenger::new();
+    challenger.observe_cap(&trace_cap);
+
+    let alphas = challenger.get_n_challenges(config.num_challenges);
+    let quotient_polys = compute_quotient_polys::<F, C, S, D>(
+        &stark,
+        &trace_commitment,
+        public_inputs,
+        alphas,
+        degree_bits,
+        rate_bits,
+    );
+    let all_quotient_chunks = quotient_polys
+        .into_par_iter()
+        .flat_map(|mut quotient_poly| {
+            quotient_poly.trim();
+            quotient_poly
+                .pad(degree << rate_bits)
+                .expect("Quotient has failed, the vanishing polynomial is not divisible by `Z_H");
+            // Split quotient into degree-n chunks.
+            quotient_poly.chunks(degree)
+        })
+        .collect();
+    let quotient_commitment = timed!(
+        timing,
+        "compute quotient commitment",
+        PolynomialBatch::from_coeffs(
+            all_quotient_chunks,
+            rate_bits,
+            false,
+            config.fri_config.cap_height,
+            timing,
+            None,
+        )
+    );
+    challenger.observe_cap(&quotient_commitment.merkle_tree.cap);
+
+    let zeta = challenger.get_extension_challenge::<D>();
+    // To avoid leaking witness data, we want to ensure that our opening locations, `zeta` and
+    // `g * zeta`, are not in our subgroup `H`. It suffices to check `zeta` only, since
+    // `(g * zeta)^n = zeta^n`, where `n` is the order of `g`.
+    let g = F::Extension::primitive_root_of_unity(degree_bits);
+    ensure!(
+        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);
+
+    // TODO: Add permuation checks
+    let initial_merkle_trees = &[&trace_commitment, &quotient_commitment];
     let fri_params = config.fri_params(degree_bits);
 
-    let opening_proof = fri_proof::<F, C, D>(
+    let opening_proof = timed!(
-        initial_merkle_trees,
-        lde_polynomial_coeffs,
-        lde_polynomial_values,
-        &mut challenger,
-        &fri_params,
         timing,
+        "compute openings proof",
+        PolynomialBatch::prove_openings(
+            &S::fri_instance(zeta, g, rate_bits),
+            initial_merkle_trees,
+            &mut challenger,
+            &fri_params,
+            timing,
+        )
     );
 
-    StarkProof {
+    Ok(StarkProof {
         trace_cap,
         openings,
         opening_proof,
-    }
+    })
+}
+
+/// Computes the quotient polynomials `(sum alpha^i C_i(x)) / Z_H(x)` for `alpha` in `alphas`,
+/// where the `C_i`s are the Stark constraints.
+// TODO: This won't work for the Fibonacci example because the constraints wrap around the subgroup.
+// The denominator should be the vanishing polynomial of `H` without its last element.
+fn compute_quotient_polys<F, C, S, const D: usize>(
+    stark: &S,
+    trace_commitment: &PolynomialBatch<F, C, D>,
+    public_inputs: [F; S::PUBLIC_INPUTS],
+    alphas: Vec<F>,
+    degree_bits: usize,
+    rate_bits: usize,
+) -> Vec<PolynomialCoeffs<F>>
+where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+    [(); S::COLUMNS]:,
+    [(); S::PUBLIC_INPUTS]:,
+{
+    let degree = 1 << degree_bits;
+    let points = F::two_adic_subgroup(degree_bits + rate_bits);
+
+    // Evaluation of the first Lagrange polynomial on the LDE domain.
+    let lagrange_first = {
+        let mut evals = PolynomialValues::new(vec![F::ZERO; degree]);
+        evals.values[0] = F::ONE;
+        evals.lde(rate_bits)
+    };
+    // Evaluation of the last Lagrange polynomial on the LDE domain.
+    let lagrange_last = {
+        let mut evals = PolynomialValues::new(vec![F::ZERO; degree]);
+        evals.values[degree - 1] = F::ONE;
+        evals.lde(rate_bits)
+    };
+
+    let z_h_on_coset = ZeroPolyOnCoset::new(degree_bits, rate_bits);
+
+    // Retrieve the LDE values at index `i`.
+    let get_at_index = |comm: &PolynomialBatch<F, C, D>, i: usize| -> [F; S::COLUMNS] {
+        comm.get_lde_values(i).try_into().unwrap()
+    };
+
+    let quotient_values = (0..degree << rate_bits)
+        .into_par_iter()
+        .map(|i| {
+            // TODO: Set `P` to a genuine `PackedField` here.
+            let mut consumer = ConstraintConsumer::<F>::new(
+                alphas.clone(),
+                lagrange_first.values[i],
+                lagrange_last.values[i],
+            );
+            let vars = StarkEvaluationVars::<F, F, { S::COLUMNS }, { S::PUBLIC_INPUTS }> {
+                local_values: &get_at_index(trace_commitment, i),
+                next_values: &get_at_index(trace_commitment, (i + 1) % (degree << rate_bits)),
+                public_inputs: &public_inputs,
+            };
+            stark.eval_packed_base(vars, &mut consumer);
+            // TODO: Fix this once we a genuine `PackedField`.
+            let mut constraints_evals = consumer.accumulators();
+            let denominator_inv = z_h_on_coset.eval_inverse(i);
+            for eval in &mut constraints_evals {
+                *eval *= denominator_inv;
+            }
+            constraints_evals
+        })
+        .collect::<Vec<_>>();
+
+    transpose(&quotient_values)
+        .into_par_iter()
+        .map(PolynomialValues::new)
+        .map(|values| values.coset_ifft(F::coset_shift()))
+        .collect()
 }

starky/src/stark.rs

@@ -1,5 +1,6 @@
 use plonky2::field::extension_field::{Extendable, FieldExtension};
 use plonky2::field::packed_field::PackedField;
+use plonky2::fri::structure::{FriBatchInfo, FriInstanceInfo, FriOracleInfo, FriPolynomialInfo};
 use plonky2::hash::hash_types::RichField;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 
@@ -8,7 +9,7 @@ use crate::vars::StarkEvaluationTargets;
 use crate::vars::StarkEvaluationVars;
 
 /// Represents a STARK system.
-pub trait Stark<F: RichField + Extendable<D>, const D: usize> {
+pub trait Stark<F: RichField + Extendable<D>, const D: usize>: Sync {
     /// The total number of columns in the trace.
     const COLUMNS: usize;
     /// The number of public inputs.
@@ -59,4 +60,28 @@ pub trait Stark<F: RichField + Extendable<D>, const D: usize> {
         vars: StarkEvaluationTargets<D, { Self::COLUMNS }, { Self::PUBLIC_INPUTS }>,
         yield_constr: &mut RecursiveConstraintConsumer<F, D>,
     );
+
+    /// Computes the FRI instance used to prove this Stark.
+    // TODO: Permutation polynomials.
+    fn fri_instance(
+        zeta: F::Extension,
+        g: F::Extension,
+        rate_bits: usize,
+    ) -> 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..1 << rate_bits);
+        let zeta_batch = FriBatchInfo {
+            point: zeta,
+            polynomials: [trace_info.clone(), quotient_info].concat(),
+        };
+        let zeta_right_batch = FriBatchInfo::<F, D> {
+            point: zeta * g,
+            polynomials: trace_info,
+        };
+        FriInstanceInfo {
+            oracles: vec![no_blinding_oracle; 3],
+            batches: vec![zeta_batch],
+        }
+    }
 }

system_zero/src/system_zero.rs

@@ -83,27 +83,33 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for SystemZero<F,
 
 #[cfg(test)]
 mod tests {
+    use anyhow::Result;
     use log::Level;
+    use plonky2::field::field_types::Field;
     use plonky2::field::goldilocks_field::GoldilocksField;
     use plonky2::plonk::config::PoseidonGoldilocksConfig;
     use plonky2::util::timing::TimingTree;
     use starky::config::StarkConfig;
     use starky::prover::prove;
+    use starky::stark::Stark;
 
     use crate::system_zero::SystemZero;
 
     #[test]
     #[ignore] // TODO
-    fn run() {
+    fn run() -> Result<()> {
         type F = GoldilocksField;
         type C = PoseidonGoldilocksConfig;
         const D: usize = 2;
 
         type S = SystemZero<F, D>;
         let system = S::default();
+        let public_inputs = [F::ZERO; S::PUBLIC_INPUTS];
         let config = StarkConfig::standard_fast_config();
         let mut timing = TimingTree::new("prove", Level::Debug);
         let trace = system.generate_trace();
-        prove::<F, C, S, D>(system, config, trace, &mut timing);
+        prove::<F, C, S, D>(system, config, trace, public_inputs, &mut timing)?;
+
+        Ok(())
     }
 }

util/src/lib.rs

@@ -11,6 +11,7 @@ use std::mem::size_of;
 use std::ptr::{swap, swap_nonoverlapping};
 
 mod transpose_util;
+
 use crate::transpose_util::transpose_in_place_square;
 
 pub fn bits_u64(n: u64) -> usize {