Recover constraint debugging

circuits/src/bitshift/stark.rs

@@ -13,7 +13,7 @@ use starky::stark::Stark;
 
 use super::columns::BitshiftView;
 use crate::columns_view::{HasNamedColumns, NumberOfColumns};
-use crate::expr::{build_ext, build_packed, ConstraintBuilder};
+use crate::expr::{build_ext, build_packed, ConstraintBuilder, GenerateConstraints};
 use crate::unstark::NoColumns;
 
 /// Bitshift Trace Constraints
@@ -30,45 +30,50 @@ impl<F, const D: usize> HasNamedColumns for BitshiftStark<F, D> {
 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<'a, F, T: Copy, U, const D: usize>
+    GenerateConstraints<'a, T, T, BitshiftView<Expr<'a, T>>, NoColumns<U>>
+    for BitshiftStark<F, { D }>
+{
+    fn generate_constraints(
+        vars: &StarkFrameTyped<BitshiftView<Expr<'a, T>>, NoColumns<U>>,
+    ) -> 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> {
@@ -88,7 +93,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for BitshiftStark
         FE: FieldExtension<D2, BaseField = F>,
         P: PackedField<Scalar = FE>, {
         let expr_builder = ExprBuilder::default();
-        let constraints = generate_constraints(&expr_builder.to_typed_starkframe(vars));
+        let constraints = Self::generate_constraints(&expr_builder.to_typed_starkframe(vars));
         build_packed(constraints, constraint_consumer);
     }
 
@@ -101,7 +106,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for BitshiftStark
         constraint_consumer: &mut RecursiveConstraintConsumer<F, D>,
     ) {
         let expr_builder = ExprBuilder::default();
-        let constraints = generate_constraints(&expr_builder.to_typed_starkframe(vars));
+        let constraints = Self::generate_constraints(&expr_builder.to_typed_starkframe(vars));
         build_ext(constraints, circuit_builder, constraint_consumer);
     }
 }

circuits/src/cpu/stark.rs

@@ -15,7 +15,8 @@ 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::cpu::shift;
-use crate::expr::{build_ext, build_packed, ConstraintBuilder};
+use crate::expr::{build_ext, build_packed, ConstraintBuilder, GenerateConstraints};
+use crate::unstark::NoColumns;
 
 /// A Gadget for CPU Instructions
 ///
@@ -40,7 +41,7 @@ fn pc_ticks_up<'a, P: Copy>(lv: &CpuState<Expr<'a, P>>, cb: &mut ConstraintBuild
 /// Ie exactly one of them should be 1, and all others 0 in each row.
 /// See <https://en.wikipedia.org/wiki/One-hot>
 fn binary_selectors<'a, P: Copy>(
-    ops: &'a OpSelectors<Expr<'a, P>>,
+    ops: &OpSelectors<Expr<'a, P>>,
     cb: &mut ConstraintBuilder<Expr<'a, P>>,
 ) {
     // selectors have value 0 or 1.
@@ -77,35 +78,39 @@ 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, U>(
-    vars: &'a StarkFrameTyped<CpuState<Expr<'a, T>>, Vec<U>>,
-) -> 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<'a, F, T: Copy, U, const D: usize>
+    GenerateConstraints<'a, T, T, CpuState<Expr<'a, T>>, NoColumns<U>> for CpuStark<F, { D }>
+{
+    fn generate_constraints(
+        vars: &StarkFrameTyped<CpuState<Expr<'a, T>>, NoColumns<U>>,
+    ) -> 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> {
@@ -125,8 +130,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         FE: FieldExtension<D2, BaseField = F>,
         P: PackedField<Scalar = FE>, {
         let expr_builder = ExprBuilder::default();
-        let vars = expr_builder.to_typed_starkframe(vars);
-        let constraints = generate_constraints(&vars);
+        let constraints = Self::generate_constraints(&expr_builder.to_typed_starkframe(vars));
         build_packed(constraints, constraint_consumer);
     }
 
@@ -139,8 +143,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         constraint_consumer: &mut RecursiveConstraintConsumer<F, D>,
     ) {
         let expr_builder = ExprBuilder::default();
-        let vars = expr_builder.to_typed_starkframe(vars);
-        let constraints = generate_constraints(&vars);
+        let constraints = Self::generate_constraints(&expr_builder.to_typed_starkframe(vars));
         build_ext(constraints, circuit_builder, constraint_consumer);
     }
 }

circuits/src/cpu_skeleton/stark.rs

@@ -13,7 +13,7 @@ use starky::stark::Stark;
 
 use super::columns::CpuSkeleton;
 use crate::columns_view::{HasNamedColumns, NumberOfColumns};
-use crate::expr::{build_ext, build_packed, ConstraintBuilder};
+use crate::expr::{build_ext, build_packed, ConstraintBuilder, GenerateConstraints};
 use crate::stark::mozak_stark::PublicInputs;
 
 #[derive(Clone, Copy, Default, StarkNameDisplay)]
@@ -30,42 +30,47 @@ const COLUMNS: usize = CpuSkeleton::<()>::NUMBER_OF_COLUMNS;
 // Public inputs: [PC of the first row]
 const PUBLIC_INPUTS: usize = PublicInputs::<()>::NUMBER_OF_COLUMNS;
 
-fn generate_constraints<'a, T: Copy>(
-    vars: &StarkFrameTyped<CpuSkeleton<Expr<'a, T>>, PublicInputs<Expr<'a, T>>>,
-) -> ConstraintBuilder<Expr<'a, T>> {
-    let lv = vars.local_values;
-    let nv = vars.next_values;
-    let public_inputs = vars.public_inputs;
-    let mut constraints = ConstraintBuilder::default();
-
-    constraints.first_row(lv.pc - public_inputs.entry_point);
-    // Clock starts at 2. This is to differentiate
-    // execution clocks (2 and above) from
-    // clk values `0` and `1` which are reserved for
-    // elf initialisation and zero initialisation respectively.
-    constraints.first_row(lv.clk - 2);
-
-    let clock_diff = nv.clk - lv.clk;
-    constraints.transition(clock_diff.is_binary());
-
-    // clock only counts up when we are still running.
-    constraints.transition(clock_diff - lv.is_running);
-
-    // We start in running state.
-    constraints.first_row(lv.is_running - 1);
-
-    // We may transition to a non-running state.
-    constraints.transition(nv.is_running * (nv.is_running - lv.is_running));
-
-    // We end in a non-running state.
-    constraints.last_row(lv.is_running);
-
-    // NOTE: in our old CPU table we had constraints that made sure nothing
-    // changes anymore, once we are halted. We don't need those
-    // anymore: the only thing that can change are memory or registers.  And
-    // our CTLs make sure, that after we are halted, no more memory
-    // or register changes are allowed.
-    constraints
+impl<'a, F, T: Copy, const D: usize>
+    GenerateConstraints<'a, T, T, CpuSkeleton<Expr<'a, T>>, PublicInputs<Expr<'a, T>>>
+    for CpuSkeletonStark<F, { D }>
+{
+    fn generate_constraints(
+        vars: &StarkFrameTyped<CpuSkeleton<Expr<'a, T>>, PublicInputs<Expr<'a, T>>>,
+    ) -> ConstraintBuilder<Expr<'a, T>> {
+        let lv = vars.local_values;
+        let nv = vars.next_values;
+        let public_inputs = vars.public_inputs;
+        let mut constraints = ConstraintBuilder::default();
+
+        constraints.first_row(lv.pc - public_inputs.entry_point);
+        // Clock starts at 2. This is to differentiate
+        // execution clocks (2 and above) from
+        // clk values `0` and `1` which are reserved for
+        // elf initialisation and zero initialisation respectively.
+        constraints.first_row(lv.clk - 2);
+
+        let clock_diff = nv.clk - lv.clk;
+        constraints.transition(clock_diff.is_binary());
+
+        // clock only counts up when we are still running.
+        constraints.transition(clock_diff - lv.is_running);
+
+        // We start in running state.
+        constraints.first_row(lv.is_running - 1);
+
+        // We may transition to a non-running state.
+        constraints.transition(nv.is_running * (nv.is_running - lv.is_running));
+
+        // We end in a non-running state.
+        constraints.last_row(lv.is_running);
+
+        // NOTE: in our old CPU table we had constraints that made sure nothing
+        // changes anymore, once we are halted. We don't need those
+        // anymore: the only thing that can change are memory or registers.  And
+        // our CTLs make sure, that after we are halted, no more memory
+        // or register changes are allowed.
+        constraints
+    }
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuSkeletonStark<F, D> {
@@ -86,7 +91,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuSkeletonSt
         P: PackedField<Scalar = FE>, {
         let expr_builder = ExprBuilder::default();
         let vars = expr_builder.to_typed_starkframe(vars);
-        let constraints = generate_constraints(&vars);
+        let constraints = Self::generate_constraints(&vars);
         build_packed(constraints, consumer);
     }
 
@@ -97,7 +102,8 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuSkeletonSt
         consumer: &mut RecursiveConstraintConsumer<F, D>,
     ) {
         let eb = ExprBuilder::default();
-        let constraints = generate_constraints(&eb.to_typed_starkframe(vars));
+        let vars = eb.to_typed_starkframe(vars);
+        let constraints = Self::generate_constraints(&vars);
         build_ext(constraints, builder, consumer);
     }
 

circuits/src/expr.rs

@@ -1,7 +1,7 @@
 use std::panic::Location;
 
 pub use expr::PureEvaluator;
-use expr::{BinOp, Cached, Evaluator, Expr, UnaOp};
+use expr::{BinOp, Cached, Evaluator, Expr, StarkFrameTyped, UnaOp};
 use plonky2::field::extension::{Extendable, FieldExtension};
 use plonky2::field::packed::PackedField;
 use plonky2::hash::hash_types::RichField;
@@ -175,7 +175,15 @@ pub fn build_packed<F, FE, P, const D: usize, const D2: usize>(
         .map(|c| c.map(|constraint| evaluator.eval(constraint)))
         .collect::<Vec<_>>();
 
+    let mozak_stark_debug = std::env::var("MOZAK_STARK_DEBUG").is_ok();
     for c in evaluated {
+        if mozak_stark_debug && !c.term.is_zeros() {
+            log::error!(
+                "ConstraintConsumer - DEBUG trace (non-zero-constraint): {}",
+                c.location
+            );
+        }
+
         (match c.constraint_type {
             ConstraintType::FirstRow => ConstraintConsumer::constraint_first_row,
             ConstraintType::Always => ConstraintConsumer::constraint,
@@ -184,3 +192,9 @@ pub fn build_packed<F, FE, P, const D: usize, const D2: usize>(
         })(yield_constr, c.term);
     }
 }
+
+pub trait GenerateConstraints<'a, T, U, View, PublicInputs> {
+    fn generate_constraints(
+        vars: &StarkFrameTyped<View, PublicInputs>,
+    ) -> ConstraintBuilder<Expr<'a, T>>;
+}