diff --git a/plonky3/src/circuit_builder.rs b/plonky3/src/circuit_builder.rs index 84ec64a38..9c47a905d 100644 --- a/plonky3/src/circuit_builder.rs +++ b/plonky3/src/circuit_builder.rs @@ -6,7 +6,10 @@ //! everywhere save for at row j is constructed to constrain s * (pub - x) on //! every row. -use std::{any::TypeId, collections::BTreeMap}; +use std::{ + any::TypeId, + collections::{BTreeMap, HashSet}, +}; use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder}; use p3_field::AbstractField; @@ -14,16 +17,46 @@ use p3_goldilocks::Goldilocks; use p3_matrix::{dense::RowMajorMatrix, Matrix}; use powdr_ast::analyzed::{ AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression, - AlgebraicUnaryOperation, AlgebraicUnaryOperator, Analyzed, IdentityKind, PolynomialType, + AlgebraicUnaryOperation, AlgebraicUnaryOperator, Analyzed, Identity, IdentityKind, + PolynomialType, SelectedExpressions, }; use powdr_executor::witgen::WitgenCallback; -use powdr_number::{FieldElement, GoldilocksField, LargeInt}; +use powdr_number::{DegreeType, FieldElement, GoldilocksField, LargeInt}; pub type Val = p3_goldilocks::Goldilocks; +/// A description of the constraint system. +/// All of the data is derived from the analyzed PIL, but is materialized +/// here for performance reasons. +struct ConstraintSystem { + identities: Vec>>>, + publics: Vec<(String, usize, usize)>, + commitment_count: usize, + constant_count: usize, + degrees: HashSet, +} + +impl From<&Analyzed> for ConstraintSystem { + fn from(analyzed: &Analyzed) -> Self { + let identities = analyzed.identities_with_inlined_intermediate_polynomials(); + let publics = analyzed.get_publics(); + let commitment_count = analyzed.commitment_count(); + let constant_count = analyzed.constant_count(); + let degrees = analyzed.degrees(); + + Self { + identities, + publics, + commitment_count, + constant_count, + degrees, + } + } +} + pub(crate) struct PowdrCircuit<'a, T> { - /// The analyzed PIL - analyzed: &'a Analyzed, + /// The constraint system description + constraint_system: ConstraintSystem, /// The value of the witness columns, if set witness: Option<&'a [(String, Vec)]>, /// Callback to augment the witness in the later stages @@ -37,7 +70,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> { pub fn generate_trace_rows(&self) -> RowMajorMatrix { // an iterator over all columns, committed then fixed let witness = self.witness().iter(); - let degrees = self.analyzed.degrees(); + let degrees = &self.constraint_system.degrees; let values = match degrees.len() { 1 => { @@ -78,7 +111,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> { } Self { - analyzed, + constraint_system: analyzed.into(), witness: None, _witgen_callback: None, #[cfg(debug_assertions)] @@ -100,8 +133,8 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> { .map(|(name, values)| (name, values)) .collect::>(); - self.analyzed - .get_publics() + self.constraint_system + .publics .iter() .map(|(col_name, _, idx)| { let vals = *witness.get(&col_name).unwrap(); @@ -111,7 +144,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> { } pub(crate) fn with_witness(self, witness: &'a [(String, Vec)]) -> Self { - assert_eq!(witness.len(), self.analyzed.commitment_count()); + assert_eq!(witness.len(), self.constraint_system.commitment_count); Self { witness: Some(witness), ..self @@ -149,7 +182,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> { match poly_id.ptype { PolynomialType::Committed => { assert!( - r.poly_id.id < self.analyzed.commitment_count() as u64, + r.poly_id.id < self.constraint_system.commitment_count as u64, "Plonky3 expects `poly_id` to be contiguous" ); let row = main.row_slice(r.next as usize); @@ -157,7 +190,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> { } PolynomialType::Constant => { assert!( - r.poly_id.id < self.analyzed.constant_count() as u64, + r.poly_id.id < self.constraint_system.constant_count as u64, "Plonky3 expects `poly_id` to be contiguous" ); let row = fixed.row_slice(r.next as usize); @@ -206,11 +239,11 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> { impl<'a, T: FieldElement> BaseAir for PowdrCircuit<'a, T> { fn width(&self) -> usize { - self.analyzed.commitment_count() + self.constraint_system.commitment_count } fn preprocessed_width(&self) -> usize { - self.analyzed.constant_count() + self.analyzed.publics_count() + self.constraint_system.constant_count + self.constraint_system.publics.len() } fn preprocessed_trace(&self) -> Option> { @@ -230,38 +263,35 @@ impl<'a, T: FieldElement, AB: AirBuilderWithPublicValues + PairBuilder> let main = builder.main(); let fixed = builder.preprocessed(); let pi = builder.public_values(); - let publics = self.analyzed.get_publics(); - assert_eq!(publics.len(), pi.len()); + assert_eq!(self.constraint_system.publics.len(), pi.len()); let local = main.row_slice(0); // public constraints - let public_vals_by_id = publics + let public_vals_by_id = self + .constraint_system + .publics .iter() .zip(pi.to_vec()) .map(|((id, _, _), val)| (id, val)) .collect::::PublicVar>>(); let fixed_local = fixed.row_slice(0); - let public_offset = self.analyzed.constant_count(); + let public_offset = self.constraint_system.constant_count; - publics - .iter() - .enumerate() - .for_each(|(index, (pub_id, col_id, _))| { + self.constraint_system.publics.iter().enumerate().for_each( + |(index, (pub_id, col_id, _))| { let selector = fixed_local[public_offset + index]; let witness_col = local[*col_id]; let public_value = public_vals_by_id[pub_id]; // constraining s(i) * (pub[i] - x(i)) = 0 builder.assert_zero(selector * (public_value.into() - witness_col)); - }); + }, + ); // circuit constraints - for identity in &self - .analyzed - .identities_with_inlined_intermediate_polynomials() - { + for identity in &self.constraint_system.identities { match identity.kind { IdentityKind::Polynomial => { assert_eq!(identity.left.expressions.len(), 0);