Recover constraint debugging

circuits/src/bitshift/stark.rs

@@ -1,108 +1,69 @@
-use std::marker::PhantomData;
+use core::fmt::Debug;
 
-use expr::{Expr, ExprBuilder, StarkFrameTyped};
+use expr::Expr;
 use mozak_circuits_derive::StarkNameDisplay;
-use plonky2::field::extension::{Extendable, FieldExtension};
-use plonky2::field::packed::PackedField;
-use plonky2::hash::hash_types::RichField;
-use plonky2::iop::ext_target::ExtensionTarget;
-use plonky2::plonk::circuit_builder::CircuitBuilder;
-use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
-use starky::evaluation_frame::StarkFrame;
-use starky::stark::Stark;
 
 use super::columns::BitshiftView;
-use crate::columns_view::{HasNamedColumns, NumberOfColumns};
-use crate::expr::{build_ext, build_packed, ConstraintBuilder};
+use crate::columns_view::NumberOfColumns;
+use crate::expr::{ConstraintBuilder, GenerateConstraints, StarkFrom, Vars};
 use crate::unstark::NoColumns;
 
 /// Bitshift Trace Constraints
 #[derive(Copy, Clone, Default, StarkNameDisplay)]
 #[allow(clippy::module_name_repetitions)]
-pub struct BitshiftStark<F, const D: usize> {
-    pub _f: PhantomData<F>,
-}
-
-impl<F, const D: usize> HasNamedColumns for BitshiftStark<F, D> {
-    type Columns = BitshiftView<F>;
-}
+pub struct BitshiftConstraints {}
 
 const COLUMNS: usize = BitshiftView::<()>::NUMBER_OF_COLUMNS;
 const PUBLIC_INPUTS: usize = 0;
 
-fn generate_constraints<'a, T: Copy>(
-    vars: &StarkFrameTyped<BitshiftView<Expr<'a, T>>, NoColumns<Expr<'a, T>>>,
-) -> ConstraintBuilder<Expr<'a, T>> {
-    let lv = vars.local_values.executed;
-    let nv = vars.next_values.executed;
-    let mut constraints = ConstraintBuilder::default();
-
-    // Constraints on shift amount
-    // They ensure:
-    //  1. Shift amount increases with each row by 0 or 1.
-    // (We allow increases of 0 in order to allow the table to add
-    //  multiple same value rows. This is needed when we have multiple
-    //  `SHL` or `SHR` operations with the same shift amount.)
-    //  2. We have shift amounts starting from 0 to max possible value of 31.
-    // (This is due to RISC-V max shift amount being 31.)
-
-    let diff = nv.amount - lv.amount;
-    // Check: initial amount value is set to 0
-    constraints.first_row(lv.amount);
-    // Check: amount value is increased by 1 or kept unchanged
-    constraints.transition(diff * (diff - 1));
-    // Check: last amount value is set to 31
-    constraints.last_row(lv.amount - 31);
-
-    // Constraints on multiplier
-    // They ensure:
-    //  1. Shift multiplier is multiplied by 2 only if amount increases.
-    //  2. We have shift multiplier from 1 to max possible value of 2^31.
-
-    // Check: initial multiplier value is set to 1 = 2^0
-    constraints.first_row(lv.multiplier - 1);
-    // Check: multiplier value is doubled if amount is increased
-    constraints.transition(nv.multiplier - (1 + diff) * lv.multiplier);
-    // Check: last multiplier value is set to 2^31
-    // (Note that based on the previous constraint, this is already
-    //  satisfied if the last amount value is 31. We leave it for readability.)
-    constraints.last_row(lv.multiplier - (1 << 31));
-
-    constraints
-}
-
-impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for BitshiftStark<F, D> {
-    type EvaluationFrame<FE, P, const D2: usize> = StarkFrame<P, P::Scalar, COLUMNS, PUBLIC_INPUTS>
-
-    where
-        FE: FieldExtension<D2, BaseField = F>,
-        P: PackedField<Scalar = FE>;
-    type EvaluationFrameTarget =
-        StarkFrame<ExtensionTarget<D>, ExtensionTarget<D>, COLUMNS, PUBLIC_INPUTS>;
-
-    fn eval_packed_generic<FE, P, const D2: usize>(
-        &self,
-        vars: &Self::EvaluationFrame<FE, P, D2>,
-        constraint_consumer: &mut ConstraintConsumer<P>,
-    ) where
-        FE: FieldExtension<D2, BaseField = F>,
-        P: PackedField<Scalar = FE>, {
-        let expr_builder = ExprBuilder::default();
-        let constraints = generate_constraints(&expr_builder.to_typed_starkframe(vars));
-        build_packed(constraints, constraint_consumer);
-    }
+#[allow(clippy::module_name_repetitions)]
+pub type BitshiftStark<F, const D: usize> =
+    StarkFrom<F, BitshiftConstraints, { D }, { COLUMNS }, { PUBLIC_INPUTS }>;
 
-    fn constraint_degree(&self) -> usize { 3 }
+impl GenerateConstraints<{ COLUMNS }, { PUBLIC_INPUTS }> for BitshiftConstraints {
+    type PublicInputs<E: Debug> = NoColumns<E>;
+    type View<E: Debug> = BitshiftView<E>;
 
-    fn eval_ext_circuit(
+    fn generate_constraints<'a, T: Debug + Copy>(
         &self,
-        circuit_builder: &mut CircuitBuilder<F, D>,
-        vars: &Self::EvaluationFrameTarget,
-        constraint_consumer: &mut RecursiveConstraintConsumer<F, D>,
-    ) {
-        let expr_builder = ExprBuilder::default();
-        let constraints = generate_constraints(&expr_builder.to_typed_starkframe(vars));
-        build_ext(constraints, circuit_builder, constraint_consumer);
+        vars: &Vars<'a, Self, T, COLUMNS, PUBLIC_INPUTS>,
+    ) -> ConstraintBuilder<Expr<'a, T>> {
+        let lv = vars.local_values.executed;
+        let nv = vars.next_values.executed;
+        let mut constraints = ConstraintBuilder::default();
+
+        // Constraints on shift amount
+        // They ensure:
+        //  1. Shift amount increases with each row by 0 or 1.
+        // (We allow increases of 0 in order to allow the table to add
+        //  multiple same value rows. This is needed when we have multiple
+        //  `SHL` or `SHR` operations with the same shift amount.)
+        //  2. We have shift amounts starting from 0 to max possible value of 31.
+        // (This is due to RISC-V max shift amount being 31.)
+
+        let diff = nv.amount - lv.amount;
+        // Check: initial amount value is set to 0
+        constraints.first_row(lv.amount);
+        // Check: amount value is increased by 1 or kept unchanged
+        constraints.transition(diff * (diff - 1));
+        // Check: last amount value is set to 31
+        constraints.last_row(lv.amount - 31);
+
+        // Constraints on multiplier
+        // They ensure:
+        //  1. Shift multiplier is multiplied by 2 only if amount increases.
+        //  2. We have shift multiplier from 1 to max possible value of 2^31.
+
+        // Check: initial multiplier value is set to 1 = 2^0
+        constraints.first_row(lv.multiplier - 1);
+        // Check: multiplier value is doubled if amount is increased
+        constraints.transition(nv.multiplier - (1 + diff) * lv.multiplier);
+        // Check: last multiplier value is set to 2^31
+        // (Note that based on the previous constraint, this is already
+        //  satisfied if the last amount value is 31. We leave it for readability.)
+        constraints.last_row(lv.multiplier - (1 << 31));
+
+        constraints
     }
 }
 

circuits/src/columns_view.rs

@@ -27,10 +27,6 @@ pub(crate) const unsafe fn transmute_ref<T, U>(t: &T) -> &U {
     &*(std::ptr::from_ref::<T>(t).cast::<U>())
 }
 
-pub trait HasNamedColumns {
-    type Columns;
-}
-
 pub trait NumberOfColumns {
     const NUMBER_OF_COLUMNS: usize;
 }

circuits/src/cpu/stark.rs

@@ -1,35 +1,25 @@
-use std::marker::PhantomData;
+use core::fmt::Debug;
 
-use expr::{Expr, ExprBuilder, StarkFrameTyped};
+use expr::Expr;
 use mozak_circuits_derive::StarkNameDisplay;
-use plonky2::field::extension::{Extendable, FieldExtension};
-use plonky2::field::packed::PackedField;
-use plonky2::hash::hash_types::RichField;
-use plonky2::iop::ext_target::ExtensionTarget;
-use plonky2::plonk::circuit_builder::CircuitBuilder;
-use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
-use starky::evaluation_frame::StarkFrame;
-use starky::stark::Stark;
 
 use super::columns::{CpuState, OpSelectors};
 use super::{bitwise, branches, div, ecall, jalr, memory, mul, signed_comparison, sub};
-use crate::columns_view::{HasNamedColumns, NumberOfColumns};
+use crate::columns_view::NumberOfColumns;
 use crate::cpu::shift;
-use crate::expr::{build_ext, build_packed, ConstraintBuilder};
+use crate::expr::{ConstraintBuilder, GenerateConstraints, StarkFrom, Vars};
 use crate::unstark::NoColumns;
 
+// TODO: fix StarkNameDisplay?
+#[derive(Copy, Clone, Default, StarkNameDisplay)]
+#[allow(clippy::module_name_repetitions)]
+pub struct CpuConstraints {}
+
 /// A Gadget for CPU Instructions
 ///
 /// Instructions are either handled directly or through cross table lookup
-#[derive(Copy, Clone, Default, StarkNameDisplay)]
 #[allow(clippy::module_name_repetitions)]
-pub struct CpuStark<F, const D: usize> {
-    pub _f: PhantomData<F>,
-}
-
-impl<F, const D: usize> HasNamedColumns for CpuStark<F, D> {
-    type Columns = CpuState<F>;
-}
+pub type CpuStark<F, const D: usize> = StarkFrom<F, CpuConstraints, { D }, COLUMNS, PUBLIC_INPUTS>;
 
 /// Ensure that if opcode is straight line, then program counter is incremented
 /// by 4.
@@ -78,69 +68,40 @@ fn populate_op2_value<'a, P: Copy>(
 const COLUMNS: usize = CpuState::<()>::NUMBER_OF_COLUMNS;
 const PUBLIC_INPUTS: usize = 0;
 
-fn generate_constraints<'a, T: Copy>(
-    vars: &StarkFrameTyped<CpuState<Expr<'a, T>>, NoColumns<Expr<'a, T>>>,
-) -> ConstraintBuilder<Expr<'a, T>> {
-    let lv = &vars.local_values;
-    let mut constraints = ConstraintBuilder::default();
-
-    pc_ticks_up(lv, &mut constraints);
-
-    binary_selectors(&lv.inst.ops, &mut constraints);
-
-    // Registers
-    populate_op2_value(lv, &mut constraints);
-
-    // ADD is now handled by its own table.
-    constraints.always(lv.inst.ops.add);
-    sub::constraints(lv, &mut constraints);
-    bitwise::constraints(lv, &mut constraints);
-    branches::comparison_constraints(lv, &mut constraints);
-    branches::constraints(lv, &mut constraints);
-    memory::constraints(lv, &mut constraints);
-    signed_comparison::signed_constraints(lv, &mut constraints);
-    signed_comparison::slt_constraints(lv, &mut constraints);
-    shift::constraints(lv, &mut constraints);
-    div::constraints(lv, &mut constraints);
-    mul::constraints(lv, &mut constraints);
-    jalr::constraints(lv, &mut constraints);
-    ecall::constraints(lv, &mut constraints);
-
-    constraints
-}
-
-impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D> {
-    type EvaluationFrame<FE, P, const D2: usize> = StarkFrame<P, P::Scalar, COLUMNS, PUBLIC_INPUTS>
-
-    where
-        FE: FieldExtension<D2, BaseField = F>,
-        P: PackedField<Scalar = FE>;
-    type EvaluationFrameTarget =
-        StarkFrame<ExtensionTarget<D>, ExtensionTarget<D>, COLUMNS, PUBLIC_INPUTS>;
-
-    fn eval_packed_generic<FE, P, const D2: usize>(
-        &self,
-        vars: &Self::EvaluationFrame<FE, P, D2>,
-        constraint_consumer: &mut ConstraintConsumer<P>,
-    ) where
-        FE: FieldExtension<D2, BaseField = F>,
-        P: PackedField<Scalar = FE>, {
-        let expr_builder = ExprBuilder::default();
-        let constraints = generate_constraints(&expr_builder.to_typed_starkframe(vars));
-        build_packed(constraints, constraint_consumer);
-    }
-
-    fn constraint_degree(&self) -> usize { 3 }
+impl GenerateConstraints<{ COLUMNS }, { PUBLIC_INPUTS }> for CpuConstraints {
+    type PublicInputs<E: Debug> = NoColumns<E>;
+    type View<E: Debug> = CpuState<E>;
 
-    fn eval_ext_circuit(
+    fn generate_constraints<'a, T: Debug + Copy>(
         &self,
-        circuit_builder: &mut CircuitBuilder<F, D>,
-        vars: &Self::EvaluationFrameTarget,
-        constraint_consumer: &mut RecursiveConstraintConsumer<F, D>,
-    ) {
-        let expr_builder = ExprBuilder::default();
-        let constraints = generate_constraints(&expr_builder.to_typed_starkframe(vars));
-        build_ext(constraints, circuit_builder, constraint_consumer);
+        vars: &Vars<'a, Self, T, COLUMNS, PUBLIC_INPUTS>,
+    ) -> ConstraintBuilder<Expr<'a, T>> {
+        let lv = &vars.local_values;
+        let mut constraints = ConstraintBuilder::default();
+
+        pc_ticks_up(lv, &mut constraints);
+
+        binary_selectors(&lv.inst.ops, &mut constraints);
+
+        // Registers
+        populate_op2_value(lv, &mut constraints);
+
+        // ADD is now handled by its own table.
+        constraints.always(lv.inst.ops.add);
+        sub::constraints(lv, &mut constraints);
+        bitwise::constraints(lv, &mut constraints);
+        branches::comparison_constraints(lv, &mut constraints);
+        branches::constraints(lv, &mut constraints);
+        memory::constraints(lv, &mut constraints);
+        signed_comparison::signed_constraints(lv, &mut constraints);
+        signed_comparison::slt_constraints(lv, &mut constraints);
+        shift::constraints(lv, &mut constraints);
+        div::constraints(lv, &mut constraints);
+        mul::constraints(lv, &mut constraints);
+        jalr::constraints(lv, &mut constraints);
+        ecall::constraints(lv, &mut constraints);
+
+        constraints
     }
 }