spike(stark): builder-rewrite (Plan B) constraint capture → IR [compare vs #737]

crypto/stark/src/constraint_ir/builder.rs

@@ -0,0 +1,203 @@
+//! Explicit-builder capture front-end (Plan B).
+//!
+//! Where the symbolic-field front-end (Plan A) records IR by running a
+//! constraint's generic `evaluate` over recording field types, this front-end
+//! builds the same [`ConstraintProgram`] through an explicit [`IrBuilder`]:
+//! each constraint implements [`Capture`](super::Capture) and translates its
+//! `evaluate` body into builder calls (`main`, `add`, `sub`, `mul`, ...).
+//!
+//! No fake field, no thread-local arena. The builder hash-conses every node on
+//! `(Op, Dim)` and only emits leaves for columns the constraint actually reads,
+//! so captured programs are minimal.
+
+use std::collections::HashMap;
+
+use math::field::element::FieldElement;
+use math::field::goldilocks::GoldilocksField;
+
+use super::ir::{ConstraintProgram, Dim, Op};
+
+/// A handle to a node in an [`IrBuilder`]: its arena id and result dimension.
+///
+/// `Copy` so constraint bodies read like ordinary field arithmetic.
+#[derive(Clone, Copy, Debug)]
+pub struct Expr {
+    id: u32,
+    dim: Dim,
+}
+
+impl Expr {
+    /// The node's result dimension.
+    pub fn dim(self) -> Dim {
+        self.dim
+    }
+}
+
+/// Builds a [`ConstraintProgram`] from explicit node-construction calls.
+///
+/// Nodes are appended in topological order (id `i` references only `< i`) and
+/// hash-consed on `(Op, Dim)`, so structurally identical subexpressions share a
+/// single id. Base-field constants are additionally deduplicated by value via
+/// `const_cache`. Node id `0` is reserved for `Op::Const1(0)`, matching the
+/// interpreter's convention and Plan A's arena.
+pub struct IrBuilder {
+    nodes: Vec<Op>,
+    dims: Vec<Dim>,
+    cse: HashMap<(Op, Dim), u32>,
+    const_cache: HashMap<u64, u32>,
+    roots: Vec<u32>,
+}
+
+impl Default for IrBuilder {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl IrBuilder {
+    /// Create a builder with the reserved `Op::Const1(0)` node at id 0.
+    pub fn new() -> Self {
+        let mut b = IrBuilder {
+            nodes: Vec::new(),
+            dims: Vec::new(),
+            cse: HashMap::new(),
+            const_cache: HashMap::new(),
+            roots: Vec::new(),
+        };
+        // Reserve id 0 = Const1(0). `const_base(0)` will hash-cons to this.
+        let zero = b.push(Op::Const1(0), Dim::D1);
+        debug_assert_eq!(zero.id, 0);
+        b.const_cache.insert(0, 0);
+        b
+    }
+
+    /// Append (or reuse) a node with the given op and result dimension.
+    fn push(&mut self, op: Op, dim: Dim) -> Expr {
+        if let Some(&id) = self.cse.get(&(op, dim)) {
+            return Expr { id, dim };
+        }
+        let id = self.nodes.len() as u32;
+        self.nodes.push(op);
+        self.dims.push(dim);
+        self.cse.insert((op, dim), id);
+        Expr { id, dim }
+    }
+
+    // ---------------------------------------------------------------------
+    // Leaves
+    // ---------------------------------------------------------------------
+
+    /// A main-trace column read at the given frame `offset`, row 0.
+    pub fn main(&mut self, offset: u8, col: usize) -> Expr {
+        self.push(
+            Op::Var {
+                main: true,
+                offset,
+                row: 0,
+                col: col as u16,
+            },
+            Dim::D1,
+        )
+    }
+
+    /// An aux-trace column read at the given frame `offset`, row 0 (`D3`).
+    pub fn aux(&mut self, offset: u8, col: usize) -> Expr {
+        self.push(
+            Op::Var {
+                main: false,
+                offset,
+                row: 0,
+                col: col as u16,
+            },
+            Dim::D3,
+        )
+    }
+
+    // ---------------------------------------------------------------------
+    // Constants
+    // ---------------------------------------------------------------------
+
+    /// A base-field constant from a `u64`, reduced and deduplicated by value.
+    pub fn const_base(&mut self, v: u64) -> Expr {
+        let canon = *FieldElement::<GoldilocksField>::from(v).value();
+        if let Some(&id) = self.const_cache.get(&canon) {
+            return Expr { id, dim: Dim::D1 };
+        }
+        let e = self.push(Op::Const1(canon), Dim::D1);
+        self.const_cache.insert(canon, e.id);
+        e
+    }
+
+    /// A base-field constant from an `i64`; negatives map to `p - |v|`.
+    pub fn const_signed(&mut self, v: i64) -> Expr {
+        let canon = *FieldElement::<GoldilocksField>::from(v).value();
+        if let Some(&id) = self.const_cache.get(&canon) {
+            return Expr { id, dim: Dim::D1 };
+        }
+        let e = self.push(Op::Const1(canon), Dim::D1);
+        self.const_cache.insert(canon, e.id);
+        e
+    }
+
+    /// The base-field constant `1`.
+    pub fn one(&mut self) -> Expr {
+        self.const_base(1)
+    }
+
+    // ---------------------------------------------------------------------
+    // Arithmetic
+    // ---------------------------------------------------------------------
+
+    /// `a + b`. Result is `D1` only if both operands are `D1`.
+    pub fn add(&mut self, a: Expr, b: Expr) -> Expr {
+        let dim = Self::join(a.dim, b.dim);
+        self.push(Op::Add(a.id, b.id), dim)
+    }
+
+    /// `a - b`. Result is `D1` only if both operands are `D1`.
+    pub fn sub(&mut self, a: Expr, b: Expr) -> Expr {
+        let dim = Self::join(a.dim, b.dim);
+        self.push(Op::Sub(a.id, b.id), dim)
+    }
+
+    /// `a * b`. Result is `D1` only if both operands are `D1`.
+    pub fn mul(&mut self, a: Expr, b: Expr) -> Expr {
+        let dim = Self::join(a.dim, b.dim);
+        self.push(Op::Mul(a.id, b.id), dim)
+    }
+
+    /// `-a`. Preserves the operand's dimension.
+    pub fn neg(&mut self, a: Expr) -> Expr {
+        self.push(Op::Neg(a.id), a.dim)
+    }
+
+    /// Typing join: `(D1, D1) -> D1`; any `D3` operand -> `D3`.
+    fn join(a: Dim, b: Dim) -> Dim {
+        match (a, b) {
+            (Dim::D1, Dim::D1) => Dim::D1,
+            _ => Dim::D3,
+        }
+    }
+
+    // ---------------------------------------------------------------------
+    // Emit / finish
+    // ---------------------------------------------------------------------
+
+    /// Record `e` as the root for constraint `constraint_idx`.
+    ///
+    /// Roots are stored in emit order; the minimal spike emits exactly one root
+    /// per program, so `constraint_idx` is accepted for parity with the
+    /// production design but not used to index `roots` here.
+    pub fn emit(&mut self, _constraint_idx: usize, e: Expr) {
+        self.roots.push(e.id);
+    }
+
+    /// Consume the builder and produce the captured program.
+    pub fn finish(self) -> ConstraintProgram {
+        ConstraintProgram {
+            nodes: self.nodes,
+            dims: self.dims,
+            roots: self.roots,
+        }
+    }
+}

crypto/stark/src/constraint_ir/interp.rs

@@ -0,0 +1,97 @@
+//! CPU interpreter for a captured [`ConstraintProgram`].
+//!
+//! A single forward pass over the topologically ordered nodes evaluates each
+//! node into a [`Value`] (base `D1` or extension `D3`), reusing the real
+//! `FieldElement` arithmetic so per-op results are bit-identical to the boxed
+//! constraint path. Mixed-dimension ops auto-embed the `D1` operand into `D3`,
+//! mirroring the field tower's `F: IsSubFieldOf<E>` arithmetic.
+
+use math::field::element::FieldElement;
+use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as GoldilocksExtension;
+use math::field::goldilocks::GoldilocksField;
+
+use super::ir::{ConstraintProgram, Dim, Op};
+
+type Fp = FieldElement<GoldilocksField>;
+type Fp3 = FieldElement<GoldilocksExtension>;
+
+/// A node's computed value: base field (`D1`) or degree-3 extension (`D3`).
+#[derive(Clone, Copy, Debug)]
+enum Value {
+    D1(Fp),
+    D3(Fp3),
+}
+
+impl Value {
+    /// Promote to the extension field, embedding a base value if needed.
+    fn to_ext(self) -> Fp3 {
+        match self {
+            Value::D1(x) => x.to_extension::<GoldilocksExtension>(),
+            Value::D3(x) => x,
+        }
+    }
+
+    fn as_base(self) -> Fp {
+        match self {
+            Value::D1(x) => x,
+            Value::D3(_) => {
+                panic!("expected a base (D1) value but found an extension (D3) value")
+            }
+        }
+    }
+}
+
+/// Evaluate the program's single root over a base-field main row.
+///
+/// `main_row[col]` resolves `Var { main: true, col, .. }` leaves. The minimal
+/// algebraic constraint set only reads main columns at offset 0, row 0 and
+/// returns a base-field (`D1`) value, so this returns a `FieldElement<F>`.
+pub fn eval_program_base(prog: &ConstraintProgram, main_row: &[Fp]) -> Fp {
+    let mut values: Vec<Value> = Vec::with_capacity(prog.nodes.len());
+
+    for (i, op) in prog.nodes.iter().enumerate() {
+        let v = match *op {
+            Op::Const1(c) => Value::D1(Fp::from(c)),
+            Op::Const3([c0, c1, c2]) => {
+                Value::D3(Fp3::from_raw([Fp::from(c0), Fp::from(c1), Fp::from(c2)]))
+            }
+            Op::Var { main, row, col, .. } => {
+                assert!(main, "aux leaves are not part of the minimal algebraic set");
+                assert_eq!(row, 0, "minimal set reads row 0 only");
+                Value::D1(main_row[col as usize])
+            }
+            Op::Add(a, b) => binop(&values, a, b, prog.dims[i], |x, y| x + y, |x, y| x + y),
+            Op::Sub(a, b) => binop(&values, a, b, prog.dims[i], |x, y| x - y, |x, y| x - y),
+            Op::Mul(a, b) => binop(&values, a, b, prog.dims[i], |x, y| x * y, |x, y| x * y),
+            Op::Neg(a) => match (values[a as usize], prog.dims[i]) {
+                (Value::D1(x), Dim::D1) => Value::D1(-x),
+                (val, Dim::D3) => Value::D3(-val.to_ext()),
+                (Value::D3(x), Dim::D1) => Value::D3(-x), // dim mismatch, keep ext
+            },
+            Op::Embed(a) => Value::D3(values[a as usize].to_ext()),
+        };
+        values.push(v);
+    }
+
+    let root = prog.roots[0];
+    values[root as usize].as_base()
+}
+
+/// Apply a binary op, auto-embedding to the extension field when the result
+/// dimension is `D3` (or either operand is already `D3`).
+#[inline]
+fn binop(
+    values: &[Value],
+    a: u32,
+    b: u32,
+    result_dim: Dim,
+    base_op: impl Fn(Fp, Fp) -> Fp,
+    ext_op: impl Fn(Fp3, Fp3) -> Fp3,
+) -> Value {
+    let va = values[a as usize];
+    let vb = values[b as usize];
+    match (va, vb, result_dim) {
+        (Value::D1(x), Value::D1(y), Dim::D1) => Value::D1(base_op(x, y)),
+        _ => Value::D3(ext_op(va.to_ext(), vb.to_ext())),
+    }
+}

crypto/stark/src/constraint_ir/ir.rs

@@ -0,0 +1,80 @@
+//! Flat intermediate representation (IR) for captured transition constraints.
+//!
+//! A [`ConstraintProgram`] is a topologically ordered list of [`Op`] nodes plus
+//! a per-constraint root id. It is produced by the builder capture front-end
+//! (see [`crate::constraint_ir::builder`]) and consumed by the CPU interpreter
+//! (see [`crate::constraint_ir::interp`]).
+//!
+//! The IR is single-field over Goldilocks, with a [`Dim`] tag distinguishing
+//! base (`D1`, one `u64`) from the degree-3 extension (`D3`, three `u64`).
+
+/// Field-arithmetic dimension of a node's value: base Goldilocks (`D1`) or its
+/// degree-3 extension (`D3`).
+#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default)]
+pub enum Dim {
+    /// Base field (one Goldilocks `u64`).
+    #[default]
+    D1,
+    /// Degree-3 extension (`[u64; 3]`).
+    D3,
+}
+
+/// One IR instruction. Operand fields are `u32` ids into the program's `nodes`
+/// arena; a node with id `i` only references nodes with id `< i`.
+#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
+pub enum Op {
+    /// A base-field literal (already reduced mod the Goldilocks prime).
+    Const1(u64),
+    /// An extension-field literal `[c0, c1, c2]` (each component reduced).
+    Const3([u64; 3]),
+    /// A leaf read of a main-trace cell. `main` is always `true` for the
+    /// minimal algebraic set captured by the spike; aux reads would set it
+    /// `false`. `offset`/`row` select the frame step/row, `col` the column.
+    Var {
+        /// `true` for a main-trace column read, `false` for an aux read.
+        main: bool,
+        /// Frame step index (0-based).
+        offset: u8,
+        /// Row within the step.
+        row: u8,
+        /// Column index.
+        col: u16,
+    },
+    /// `nodes[a] + nodes[b]`.
+    Add(u32, u32),
+    /// `nodes[a] - nodes[b]`.
+    Sub(u32, u32),
+    /// `nodes[a] * nodes[b]`.
+    Mul(u32, u32),
+    /// `-nodes[a]`.
+    Neg(u32),
+    /// Embed a `D1` value into `D3` (`<F as IsSubFieldOf<E>>::embed`).
+    Embed(u32),
+}
+
+/// A captured program for one transition constraint (or a set of them).
+///
+/// `nodes` is topologically ordered (id `i` references only `< i`). `dims[i]`
+/// is the result dimension of `nodes[i]`. `roots[c]` is the node id of
+/// constraint `c`'s value.
+#[derive(Clone, Debug)]
+pub struct ConstraintProgram {
+    /// Topologically ordered instruction list.
+    pub nodes: Vec<Op>,
+    /// Per-node result dimension, parallel to `nodes`.
+    pub dims: Vec<Dim>,
+    /// Per-constraint root node ids.
+    pub roots: Vec<u32>,
+}
+
+impl ConstraintProgram {
+    /// Number of nodes in the program (an effectiveness measure for hash-consing).
+    pub fn len(&self) -> usize {
+        self.nodes.len()
+    }
+
+    /// Whether the program has no nodes.
+    pub fn is_empty(&self) -> bool {
+        self.nodes.is_empty()
+    }
+}