zinc_uair/lib.rs
1//! UAIR description tools.
2
3pub mod collect_scalars;
4pub mod constraint_counter;
5pub mod degree_counter;
6pub mod do_nothing_builder;
7pub mod dummy_semiring;
8pub mod ideal;
9pub mod ideal_collector;
10pub mod lookup_types;
11
12use crate::ideal::Ideal;
13use crypto_primitives::{Semiring, SemiringConfig, SetElement};
14use std::borrow::Cow;
15use zinc_poly::{
16 mle::DenseMultilinearExtension,
17 univariate::{
18 binary::BinaryPoly,
19 dense::DensePolynomial,
20 dynamic::{DynamicPolynomial, HasDynamicPolynomialConfig},
21 },
22};
23use zinc_utils::{add, sub};
24
25pub use lookup_types::{LookupColumnSpec, LookupTableType};
26
27/// The abstract interface to constraint building logic.
28/// In essence it allows to create constraints modulo ideals.
29pub trait ConstraintBuilder {
30 /// The expressions the constraint builder operates on.
31 /// It is opaque from the PoV of an AIR: arithmetic operations on it are
32 /// provided by the expression config passed to
33 /// [`Uair::constrain_general`] alongside the builder.
34 type Expr: SetElement;
35 /// The type of ideals used by the constraint builder.
36 type Ideal: Ideal;
37 /// Ideals living over $F_{q_i}[X]$ for the prime tuple declared by
38 /// the surrounding [`UairSignature::primes`]. A single
39 /// `ConstraintBuilder` shares one runtime type for all primes; the prime
40 /// index is passed at the call site via
41 /// [`ConstraintBuilder::assert_in_fq_ideal`]. Builders that don't care
42 /// about $F_q[X]$-constraints (counters, collectors, etc.) set
43 /// this to `ImpossibleIdeal`.
44 type FqIdeal: Ideal;
45
46 /// Add a constraint saying that `expr` belongs to the ideal `ideal`.
47 fn assert_in_ideal(&mut self, expr: Self::Expr, ideal: &Self::Ideal);
48
49 /// Add a constraint saying that `expr` is equal to zero which is
50 /// the same as saying that `expr` belongs to the zero ideal.
51 fn assert_zero(&mut self, expr: Self::Expr);
52
53 /// Add a constraint saying that `expr`, after coefficient-wise reduction
54 /// mod $q_{\text{prime\_index}}$ (the paper's $\phi_{q_i}$), belongs to
55 /// the $F_{q_i}[X]$-ideal `ideal`.
56 ///
57 /// `prime_idx` indexes into [`UairSignature::primes`] and must be a
58 /// valid index for any UAIR that calls this method.
59 ///
60 /// # Ordering convention
61 ///
62 /// The order of constraints *within* each family must be stable across
63 /// `constrain_general` calls, so `count_constraints` /
64 /// `count_constraint_degrees` / `IdealCollector::{ideals, fq_ideals}`
65 /// line up per family.
66 ///
67 /// # Scope: projections of f_0 only
68 ///
69 /// `expr` is built only from the $Q[X]$-typed `up`/`down` rows passed to
70 /// [`Uair::constrain_general`] — i.e. the projection $\phi_{q_i}(\hat f_0)$
71 /// of the single integer trace. There is no separate $\hat f_i$ witness
72 /// typed natively in $F_{q_i}[X]$; $\phi_{q_i}$ is applied by the PIOP
73 /// layer at prove/verify time.
74 fn assert_in_fq_ideal(&mut self, prime_idx: usize, expr: Self::Expr, ideal: &Self::FqIdeal);
75}
76
77/// Specifies a shifted column
78/// `ShiftSpec { source_col: 0, shift_amount: 3 }` means
79/// "virtual column whose row i is the value of column 0 at row i+3
80/// (zero-padded beyond trace length)."
81///
82/// Multiple ShiftSpecs may reference the same source_col with
83/// different shift amounts.
84#[derive(Clone, Debug, PartialEq, Eq, Hash)]
85pub struct ShiftSpec {
86 /// Index of the committed column in the flattened trace
87 /// (binary_poly || arbitrary_poly || int, same indexing as
88 /// TraceRow::from_slice_with_layout).
89 source_col: usize,
90 /// Number of rows to shift by.
91 shift_amount: usize,
92}
93
94impl ShiftSpec {
95 pub fn new(source_col: usize, shift_amount: usize) -> Self {
96 assert!(shift_amount > 0, "shift must be non-zero");
97 Self {
98 source_col,
99 shift_amount,
100 }
101 }
102
103 pub fn source_col(&self) -> usize {
104 self.source_col
105 }
106
107 pub fn shift_amount(&self) -> usize {
108 self.shift_amount
109 }
110}
111
112// ---------------------------------------------------------------------------
113// BitOp virtual columns
114// ---------------------------------------------------------------------------
115
116/// An entry-wise `R`-linear endomorphism of the bounded-degree coefficient
117/// module `R^{<W}[X]` (cf. Section 2.1.1 of the Zinc+ paper) that defines a
118/// virtual column.
119///
120/// Per Lemma 2.3, any `R`-linear coordinate-wise map on `R^{<W}[X]` commutes
121/// with multilinear extension over the row hypercube. Consequently the column
122/// `T(v)` need not be committed: the prover materializes it during the
123/// constraint-aggregation sumcheck, and the verifier reconstructs its MLE
124/// evaluation at the final point `r_0` by applying `T` to the source
125/// column's lifted opening, its `W` `F_q`-coefficients, directly.
126///
127/// `Rot(c)` admits an alternative description as multiplication by `X^{W-c}`
128/// modulo `X^W - 1`, i.e. as an endomorphism of `R[X]/(X^W - 1)`. `ShR(c)` is
129/// pure zero-padding on coefficient indices and is *not* a quotient-ring
130/// operation; both, however, are `R`-linear maps on `R^{<W}[X]` and fall
131/// under the same Lemma 2.3 frame.
132///
133/// Bit-ops are defined only on binary_poly source columns.
134#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
135pub enum BitOp {
136 /// Right-rotation by `c` bit positions. The result's coefficient at
137 /// position `i` is the source's at `(i + c) mod W`, where `W` is the
138 /// cell width.
139 Rot(usize),
140 /// Right-shift by `c` bit positions. The result's coefficient at
141 /// position `i` is the source's at `i + c` if `i + c < W`, else zero.
142 ShR(usize),
143}
144
145impl BitOp {
146 /// The rotation / shift count.
147 pub fn count(&self) -> usize {
148 match self {
149 BitOp::Rot(c) | BitOp::ShR(c) => *c,
150 }
151 }
152
153 /// Apply the bit operation to a projected bit-polynomial cell.
154 pub fn transform<S: SemiringConfig, const D: usize>(
155 &self,
156 source: &DynamicPolynomial<S::Element>,
157 cfg: &S,
158 ) -> DynamicPolynomial<S::Element> {
159 let poly_cfg = cfg.dyn_poly_cfg();
160 match self {
161 BitOp::Rot(c) => poly_cfg.rotate_right::<D>(source, *c),
162 BitOp::ShR(c) => poly_cfg.shr::<D>(source, *c),
163 }
164 }
165}
166
167/// Specifies a bit-op virtual column.
168///
169/// `BitOpSpec { source_col: 0, op: BitOp::ShR(3) }` declares a virtual column
170/// whose row `i` is `ShR^3` applied entry-wise to the `i`-th cell of column 0.
171///
172/// `source_col` must reference a binary_poly column; bit-ops are only defined
173/// on bit-polynomial cells, i.e. elements of `R^{<W}[X]` with `{0,1}`
174/// coefficients.
175#[derive(Clone, Debug, PartialEq, Eq, Hash)]
176pub struct BitOpSpec {
177 /// Flat index of the binary_poly source column. Uses the same
178 /// `binary_poly || arbitrary_poly || int` indexing as `ShiftSpec`.
179 source_col: usize,
180 /// The bit-op applied entry-wise to the source column.
181 op: BitOp,
182}
183
184impl BitOpSpec {
185 pub fn new(source_col: usize, op: BitOp) -> Self {
186 assert!(op.count() > 0, "bit-op count must be non-zero");
187 Self { source_col, op }
188 }
189
190 pub fn source_col(&self) -> usize {
191 self.source_col
192 }
193
194 pub fn op(&self) -> BitOp {
195 self.op
196 }
197}
198
199// ---------------------------------------------------------------------------
200// Affine virtual columns for booleanity targets
201// ---------------------------------------------------------------------------
202
203/// One term in an affine virtual binary-polynomial expression.
204#[derive(Clone, Debug, PartialEq, Eq, Hash)]
205pub struct AffineVirtualTerm {
206 /// Flat total-column index of a binary-polynomial source. Binary-polynomial
207 /// columns form the prefix of the trace layout, with public columns before
208 /// witness columns.
209 source_col: usize,
210 /// Integer scalar applied to the source. The Ch/Maj constraints only need
211 /// small signed coefficients such as `1`, `-1`, and `-2`.
212 coefficient: i64,
213 /// Forward row shift. Zero reads the current row; a positive shift reads
214 /// `source_col[row + row_shift]`, with zero padding past the trace length.
215 row_shift: usize,
216}
217
218impl AffineVirtualTerm {
219 /// Construct an unshifted term `coefficient * source_col`.
220 pub fn new(source_col: usize, coefficient: i64) -> Self {
221 assert!(
222 coefficient != 0,
223 "affine virtual coefficient must be non-zero"
224 );
225 Self {
226 source_col,
227 coefficient,
228 row_shift: 0,
229 }
230 }
231
232 /// Construct a shifted term `coefficient * source_col[row + row_shift]`.
233 ///
234 /// [`UairSignature::with_affine_virtual_specs`] currently rejects shifted
235 /// terms until the protocol binds row-shifted source projections.
236 pub fn new_shifted(source_col: usize, coefficient: i64, row_shift: usize) -> Self {
237 assert!(row_shift > 0, "row shift must be non-zero");
238 Self {
239 row_shift,
240 ..Self::new(source_col, coefficient)
241 }
242 }
243
244 pub fn source_col(&self) -> usize {
245 self.source_col
246 }
247
248 pub fn coefficient(&self) -> i64 {
249 self.coefficient
250 }
251
252 pub fn row_shift(&self) -> usize {
253 self.row_shift
254 }
255}
256
257/// Declares an affine virtual binary-polynomial expression whose coefficients
258/// must be proved boolean.
259///
260/// The represented expression is
261///
262/// ```text
263/// ones_coefficient * 1_D + sum_i coefficient_i * source_i[row + shift_i]
264/// ```
265///
266/// where `1_D = 1 + X + ... + X^(D-1)` is the all-ones bit polynomial for the
267/// binary-poly cell width `D`. This covers the paper's Ch/Maj lookup targets,
268/// for example `a + b + c - 2m` and `(1_D - e) + g - 2u`.
269///
270/// Affine virtual specs are Q-side booleanity targets. They do not add entries
271/// to the `down` row, and they are not committed columns.
272#[derive(Clone, Debug, PartialEq, Eq, Hash)]
273pub struct AffineVirtualSpec {
274 terms: Vec<AffineVirtualTerm>,
275 ones_coefficient: i64,
276}
277
278impl AffineVirtualSpec {
279 /// Construct an affine virtual expression with no `1_D` offset.
280 pub fn new(terms: Vec<AffineVirtualTerm>) -> Self {
281 Self::with_ones_coefficient(terms, 0)
282 }
283
284 /// Construct an affine virtual expression with a scalar multiple of
285 /// `1_D` as its constant bit-polynomial offset.
286 pub fn with_ones_coefficient(terms: Vec<AffineVirtualTerm>, ones_coefficient: i64) -> Self {
287 assert!(
288 !terms.is_empty(),
289 "affine virtual spec must have at least one non-constant term"
290 );
291 Self {
292 terms,
293 ones_coefficient,
294 }
295 }
296
297 pub fn terms(&self) -> &[AffineVirtualTerm] {
298 &self.terms
299 }
300
301 pub fn ones_coefficient(&self) -> i64 {
302 self.ones_coefficient
303 }
304}
305
306// ---------------------------------------------------------------------------
307// Column layout types
308// ---------------------------------------------------------------------------
309
310/// Column counts per type (binary_poly, arbitrary_poly, int).
311/// Shared internals for the semantic newtype wrappers (Total, Public, Virtual,
312/// Witness)
313#[derive(Clone, Debug, Default)]
314pub struct ColumnLayout {
315 num_binary_poly_cols: usize,
316 num_arbitrary_poly_cols: usize,
317 num_int_cols: usize,
318}
319
320impl ColumnLayout {
321 pub fn new(
322 num_binary_poly_cols: usize,
323 num_arbitrary_poly_cols: usize,
324 num_int_cols: usize,
325 ) -> Self {
326 Self {
327 num_binary_poly_cols,
328 num_arbitrary_poly_cols,
329 num_int_cols,
330 }
331 }
332
333 pub fn num_binary_poly_cols(&self) -> usize {
334 self.num_binary_poly_cols
335 }
336
337 pub fn num_arbitrary_poly_cols(&self) -> usize {
338 self.num_arbitrary_poly_cols
339 }
340
341 pub fn num_int_cols(&self) -> usize {
342 self.num_int_cols
343 }
344
345 /// Maximum number of columns across the three types.
346 pub fn max_cols(&self) -> usize {
347 [
348 self.num_binary_poly_cols,
349 self.num_arbitrary_poly_cols,
350 self.num_int_cols,
351 ]
352 .into_iter()
353 .max()
354 .expect("the iterator is not empty")
355 }
356
357 /// The sum of the numbers of columns across all types.
358 #[allow(clippy::arithmetic_side_effects)]
359 pub fn cols(&self) -> usize {
360 self.num_binary_poly_cols + self.num_arbitrary_poly_cols + self.num_int_cols
361 }
362}
363
364macro_rules! column_layout_wrapper {
365 ($(#[$meta:meta])* $name:ident) => {
366 $(#[$meta])*
367 #[derive(Clone, Debug, Default)]
368 pub struct $name(ColumnLayout);
369
370 impl $name {
371 pub fn new(num_binary_poly_cols: usize, num_arbitrary_poly_cols: usize, num_int_cols: usize) -> Self {
372 Self(ColumnLayout::new(num_binary_poly_cols, num_arbitrary_poly_cols, num_int_cols))
373 }
374
375 pub fn num_binary_poly_cols(&self) -> usize { self.0.num_binary_poly_cols() }
376 pub fn num_arbitrary_poly_cols(&self) -> usize { self.0.num_arbitrary_poly_cols() }
377 pub fn num_int_cols(&self) -> usize { self.0.num_int_cols() }
378 pub fn max_cols(&self) -> usize { self.0.max_cols() }
379 pub fn cols(&self) -> usize { self.0.cols() }
380 pub fn as_column_layout(&self) -> &ColumnLayout { &self.0 }
381 }
382 };
383}
384
385column_layout_wrapper!(/// Layout of all trace columns (public + witness) per type.
386 TotalColumnLayout);
387column_layout_wrapper!(/// Layout of the public column subset.
388 PublicColumnLayout);
389column_layout_wrapper!(/// Layout of the virtual (shifted/down) columns.
390 VirtualColumnLayout);
391column_layout_wrapper!(/// Layout of the witness (total minus public) columns.
392 WitnessColumnLayout);
393
394// ---------------------------------------------------------------------------
395// UairSignature
396// ---------------------------------------------------------------------------
397
398/// The signature of a UAIR.
399///
400/// Public columns precede witness columns within each type group.
401/// The flattened trace ordering is:
402/// `[pub_bin, wit_bin, pub_arb, wit_arb, pub_int, wit_int]`.
403#[derive(Clone, Debug)]
404pub struct UairSignature<Prime: Semiring> {
405 /// Column-type layout of all (public + witness) columns.
406 total_cols: TotalColumnLayout,
407 /// Public column subset.
408 public_cols: PublicColumnLayout,
409 /// Witness column counts (total minus public) per type.
410 witness_cols: WitnessColumnLayout,
411 /// Shifted columns info sorted by `source_col`.
412 shifts: Vec<ShiftSpec>,
413 /// Bit-op virtual column specs, in insertion order. Each spec references a
414 /// binary_poly source column and contributes one extra entry to the
415 /// binary_poly slice of the down row, appended after the shifted entries.
416 bit_op_specs: Vec<BitOpSpec>,
417 /// Affine virtual expressions whose cells are Q-side booleanity targets.
418 /// These are not committed columns and do not contribute to `down_cols`.
419 affine_virtual_specs: Vec<AffineVirtualSpec>,
420 /// Column-type layout of the down row (shifted virtuals + bit-op virtuals).
421 down_cols: VirtualColumnLayout,
422 /// Lookup specifications: which trace columns are constrained against
423 /// which table types.
424 lookup_specs: Vec<LookupColumnSpec>,
425 /// Prime powers `(q_1, ..., q_n)` declared by this UAIR.
426 /// $F_{q_i}[X]$-constraints emitted via
427 /// [`ConstraintBuilder::assert_in_fq_ideal`] reference these by index.
428 /// Empty for $Q[X]$-only UAIRs.
429 primes: Vec<Prime>,
430}
431
432impl<Prime: Semiring> UairSignature<Prime> {
433 /// Create a new signature, sorting `shifts` by `source_col`.
434 pub fn new(
435 total_cols: TotalColumnLayout,
436 public_cols: PublicColumnLayout,
437 mut shifts: Vec<ShiftSpec>,
438 lookup_specs: Vec<LookupColumnSpec>,
439 ) -> Self {
440 for (name, pub_n, tot_n) in [
441 (
442 "binary_poly",
443 public_cols.num_binary_poly_cols(),
444 total_cols.num_binary_poly_cols(),
445 ),
446 (
447 "arbitrary_poly",
448 public_cols.num_arbitrary_poly_cols(),
449 total_cols.num_arbitrary_poly_cols(),
450 ),
451 ("int", public_cols.num_int_cols(), total_cols.num_int_cols()),
452 ] {
453 assert!(
454 pub_n <= tot_n,
455 "public {name}_cols ({pub_n}) > total ({tot_n})"
456 );
457 }
458
459 let num_cols = total_cols.cols();
460 for spec in &shifts {
461 assert!(
462 spec.source_col() < num_cols,
463 "ShiftSpec source_col {} out of range (total_cols = {}). \
464 source_col uses flat indexing: binary_poly || arbitrary_poly || int.",
465 spec.source_col(),
466 num_cols,
467 );
468 }
469
470 shifts.sort_by_key(|spec| spec.source_col());
471 let down_cols = Self::compute_down_layout(&total_cols, &shifts, &[]);
472 let witness_cols = WitnessColumnLayout::new(
473 sub!(
474 total_cols.num_binary_poly_cols(),
475 public_cols.num_binary_poly_cols()
476 ),
477 sub!(
478 total_cols.num_arbitrary_poly_cols(),
479 public_cols.num_arbitrary_poly_cols()
480 ),
481 sub!(total_cols.num_int_cols(), public_cols.num_int_cols()),
482 );
483
484 Self {
485 total_cols,
486 public_cols,
487 shifts,
488 bit_op_specs: Vec::new(),
489 affine_virtual_specs: Vec::new(),
490 down_cols,
491 witness_cols,
492 lookup_specs,
493 primes: Vec::new(),
494 }
495 }
496
497 /// Attach the prime-power tuple `(q_1, ..., q_n)` that
498 /// $F_{q_i}[X]$-constraints emitted by this UAIR live over.
499 pub fn with_primes(mut self, primes: Vec<Prime>) -> Self {
500 self.primes = primes;
501 self
502 }
503
504 /// Prime-power tuple `(q_1, ..., q_n)` declared by this UAIR. Empty for
505 /// UAIRs with $Q[X]$-only constraints.
506 pub fn primes(&self) -> &[Prime] {
507 &self.primes
508 }
509
510 /// Attach bit-op virtual column specs to the signature.
511 ///
512 /// Each spec must reference a binary_poly source column. The bit-op count
513 /// must be less than the binary-poly cell width `W`; materialization sites
514 /// check that bound with their const `DEGREE_PLUS_ONE` parameter.
515 ///
516 /// # Down-row ordering invariant
517 ///
518 /// Bit-op virtuals slot into the `binary_poly` slice of the down
519 /// `TraceRow`, *after* the shifted-binary entries and *before* any
520 /// non-binary entries. The full ordering of the down row is:
521 ///
522 /// ```text
523 /// [shifted_binary_poly..., bit_op_binary_poly..., shifted_arbitrary_poly..., shifted_int...]
524 /// ```
525 ///
526 /// This keeps `down` consistent with `ColumnLayout`'s
527 /// `binary_poly || arbitrary_poly || int` partitioning. Materialization
528 /// code in CPR / mp_eval must respect this order; appending bit-op evals
529 /// at the tail of `down_evals` would silently misalign constraint indices
530 /// on mixed-type shift UAIRs.
531 ///
532 /// Insertion order of `bit_op_specs` determines the position of each
533 /// bit-op virtual within its sub-slice.
534 pub fn with_bit_op_specs(mut self, bit_op_specs: Vec<BitOpSpec>) -> Self {
535 let binary_poly_end = self.total_cols.num_binary_poly_cols();
536 for spec in &bit_op_specs {
537 assert!(
538 spec.source_col() < binary_poly_end,
539 "BitOpSpec source_col {} is not a binary_poly column \
540 (binary_poly_end = {}). Bit-ops are only defined on the \
541 cell ring F_2[X]/(X^W).",
542 spec.source_col(),
543 binary_poly_end,
544 );
545 }
546 self.bit_op_specs = bit_op_specs;
547 self.down_cols =
548 Self::compute_down_layout(&self.total_cols, &self.shifts, &self.bit_op_specs);
549 self
550 }
551
552 /// Attach affine virtual booleanity specs to the signature.
553 ///
554 /// Each term in each spec must be unshifted and reference a binary_poly
555 /// column. Affine virtuals describe Q-side `{0,1}^{<D}[X]` membership
556 /// targets such as the Ch/Maj linear combinations in the Zinc+ paper. They
557 /// are not committed columns and do not affect the down-row layout.
558 pub fn with_affine_virtual_specs(
559 mut self,
560 affine_virtual_specs: Vec<AffineVirtualSpec>,
561 ) -> Self {
562 let binary_poly_end = self.total_cols.num_binary_poly_cols();
563 for (spec_index, spec) in affine_virtual_specs.iter().enumerate() {
564 for (term_index, term) in spec.terms().iter().enumerate() {
565 assert!(
566 term.source_col() < binary_poly_end,
567 "AffineVirtualTerm source_col {} is not a binary_poly column \
568 (binary_poly_end = {}). Affine virtual booleanity targets \
569 are only defined over binary_poly cells.",
570 term.source_col(),
571 binary_poly_end,
572 );
573 assert!(
574 term.row_shift() == 0,
575 "shifted affine virtual terms are not supported: spec {spec_index}, \
576 term {term_index}, row shift {}",
577 term.row_shift(),
578 );
579 }
580 }
581 self.affine_virtual_specs = affine_virtual_specs;
582 self
583 }
584
585 pub fn lookup_specs(&self) -> &[LookupColumnSpec] {
586 &self.lookup_specs
587 }
588
589 fn compute_down_layout(
590 total_cols: &TotalColumnLayout,
591 shifts: &[ShiftSpec],
592 bit_op_specs: &[BitOpSpec],
593 ) -> VirtualColumnLayout {
594 let binary_poly_end = total_cols.num_binary_poly_cols();
595 let arbitrary_poly_end = add!(binary_poly_end, total_cols.num_arbitrary_poly_cols());
596 let mut num_binary_poly = 0usize;
597 let mut num_arbitrary_poly = 0usize;
598 let mut num_int = 0usize;
599 for spec in shifts {
600 if spec.source_col() < binary_poly_end {
601 num_binary_poly = add!(num_binary_poly, 1);
602 } else if spec.source_col() < arbitrary_poly_end {
603 num_arbitrary_poly = add!(num_arbitrary_poly, 1);
604 } else {
605 num_int = add!(num_int, 1);
606 }
607 }
608 num_binary_poly = add!(num_binary_poly, bit_op_specs.len());
609 VirtualColumnLayout::new(num_binary_poly, num_arbitrary_poly, num_int)
610 }
611
612 pub fn total_cols(&self) -> &TotalColumnLayout {
613 &self.total_cols
614 }
615
616 pub fn public_cols(&self) -> &PublicColumnLayout {
617 &self.public_cols
618 }
619
620 /// Witness column counts (total minus public) per type.
621 pub fn witness_cols(&self) -> &WitnessColumnLayout {
622 &self.witness_cols
623 }
624
625 pub fn shifts(&self) -> &[ShiftSpec] {
626 &self.shifts
627 }
628
629 /// Bit-op virtual column specs, in insertion order. Each spec contributes
630 /// one binary_poly entry to the down row, appended after the shifted
631 /// entries.
632 pub fn bit_op_specs(&self) -> &[BitOpSpec] {
633 &self.bit_op_specs
634 }
635
636 /// Affine virtual booleanity specs, in insertion order.
637 pub fn affine_virtual_specs(&self) -> &[AffineVirtualSpec] {
638 &self.affine_virtual_specs
639 }
640
641 /// Number of row-shift virtual columns that precede bit-op virtuals in the
642 /// down-row ordering.
643 ///
644 /// The full down-row order is:
645 /// `[shifted_binary_poly..., bit_op_binary_poly...,
646 /// shifted_arbitrary_poly..., shifted_int...]`.
647 pub fn bit_op_down_offset(&self) -> usize {
648 let binary_poly_end = self.total_cols.num_binary_poly_cols();
649 self.shifts
650 .iter()
651 .take_while(|spec| spec.source_col() < binary_poly_end)
652 .count()
653 }
654
655 /// Column-type layout of the down row (shifted virtuals + bit-op virtuals).
656 pub fn down_cols(&self) -> &VirtualColumnLayout {
657 &self.down_cols
658 }
659
660 /// Build correctly-sized dummy up and down `TraceRow`s for static
661 /// analysis (constraint counting, degree counting, scalar/ideal
662 /// collection).
663 pub fn dummy_rows<T: Clone>(&self, val: T) -> (Vec<T>, Vec<T>) {
664 let up_size = self.total_cols.cols();
665 let down_size = self.down_cols.cols();
666 (vec![val.clone(); up_size], vec![val; down_size])
667 }
668}
669
670// ---------------------------------------------------------------------------
671// UairTrace
672// ---------------------------------------------------------------------------
673
674/// The trace of a UAIR execution (pre-projection).
675/// If owned, it contains the full trace, otherwise it contains a view on the
676/// full trace (e.g. only public columns).
677#[derive(Debug, Clone, Default)]
678pub struct UairTrace<
679 'a,
680 PolyCoeff: Clone,
681 Int: Clone,
682 const BINARY_POLY_DEGREE_PLUS_ONE: usize,
683 const ARBITRARY_POLY_DEGREE_PLUS_ONE: usize,
684> {
685 pub binary_poly: Cow<'a, [DenseMultilinearExtension<BinaryPoly<BINARY_POLY_DEGREE_PLUS_ONE>>]>,
686 pub arbitrary_poly: Cow<
687 'a,
688 [DenseMultilinearExtension<DensePolynomial<PolyCoeff, ARBITRARY_POLY_DEGREE_PLUS_ONE>>],
689 >,
690 pub int: Cow<'a, [DenseMultilinearExtension<Int>]>,
691}
692
693impl<PolyCoeff: Clone, Int: Clone, const DB: usize, const DA: usize>
694 UairTrace<'static, PolyCoeff, Int, DB, DA>
695{
696 /// Returns a sub-trace containing only public columns.
697 /// Returned trace is borrowed from the full trace.
698 pub fn public<Prime: Semiring>(
699 &self,
700 sig: &UairSignature<Prime>,
701 ) -> UairTrace<'_, PolyCoeff, Int, DB, DA> {
702 let p = sig.public_cols();
703 UairTrace {
704 binary_poly: Cow::Borrowed(&self.binary_poly[0..p.num_binary_poly_cols()]),
705 arbitrary_poly: Cow::Borrowed(&self.arbitrary_poly[0..p.num_arbitrary_poly_cols()]),
706 int: Cow::Borrowed(&self.int[0..p.num_int_cols()]),
707 }
708 }
709
710 /// Returns a sub-trace containing only witness columns.
711 /// Returned trace is borrowed from the full trace.
712 pub fn witness<Prime: Semiring>(
713 &self,
714 sig: &UairSignature<Prime>,
715 ) -> UairTrace<'_, PolyCoeff, Int, DB, DA> {
716 let p = sig.public_cols();
717 UairTrace {
718 binary_poly: Cow::Borrowed(&self.binary_poly[p.num_binary_poly_cols()..]),
719 arbitrary_poly: Cow::Borrowed(&self.arbitrary_poly[p.num_arbitrary_poly_cols()..]),
720 int: Cow::Borrowed(&self.int[p.num_int_cols()..]),
721 }
722 }
723}
724
725// ---------------------------------------------------------------------------
726// TraceRow
727// ---------------------------------------------------------------------------
728
729/// A view on a row of the trace.
730/// Contains references to cells of the trace
731/// of all types lying in the same trace row.
732#[derive(Clone, Copy)]
733pub struct TraceRow<'a, Expr> {
734 pub binary_poly: &'a [Expr],
735 pub arbitrary_poly: &'a [Expr],
736 pub int: &'a [Expr],
737}
738
739impl<'a, Expr> TraceRow<'a, Expr> {
740 /// Given a slice that represents a raw row of the trace,
741 /// creates a `TraceRow` from it.
742 /// Subdivides the slice according to the given column layout.
743 #[allow(clippy::arithmetic_side_effects)]
744 pub fn from_slice_with_layout(row: &'a [Expr], layout: &ColumnLayout) -> Self {
745 let num_binary_poly = layout.num_binary_poly_cols();
746 let num_arbitrary_poly = layout.num_arbitrary_poly_cols();
747 Self {
748 binary_poly: &row[0..num_binary_poly],
749 arbitrary_poly: &row[num_binary_poly..num_binary_poly + num_arbitrary_poly],
750 int: &row[num_binary_poly + num_arbitrary_poly..],
751 }
752 }
753}
754
755// ---------------------------------------------------------------------------
756// Uair trait
757// ---------------------------------------------------------------------------
758
759/// The trait that a universal AIR description has to implement.
760/// This must include all the constraint description logic of an UAIR.
761///
762/// One type might implement different UAIR logics for different underlying
763/// semirings hence the generic type parameter.
764pub trait Uair: Clone {
765 /// The ideal type the AIR operates with.
766 /// Since a `ConstraintBuilder` is "opaque" for a `Uair`
767 /// a `Uair` has to have a means to create ideals
768 /// so ideals are fixed by this associated types.
769 /// At the `constrain*` methods a `Uair` is given
770 /// a way to convert its own ideals into builder's ideals
771 /// via the `FromRef` trait.
772 type Ideal: Ideal;
773
774 /// The ideal type for $F_{q_i}[X]$-constraints emitted via
775 /// [`ConstraintBuilder::assert_in_fq_ideal`]. UAIRs that do not declare
776 /// any primes should set this to [`ideal::ImpossibleIdeal`].
777 type FqIdeal: Ideal;
778
779 /// The type of scalars of the UAIR.
780 /// For now, we assume they are of the type "arbitrary polynomials".
781 // Note: This is usually Z_32[X] (i.e. DensePolynomial<Ring, 32>), but according
782 // to @agareta, this in not always the case.
783 type Scalar: Semiring;
784
785 /// Type of primes defined in signature. Must be compatible with the field
786 /// type we're using.
787 type Prime: Semiring;
788
789 /// Signature of the UAIR.
790 ///
791 /// TODO: Consider caching the signature to avoid recomputing it at every
792 /// call site. Currently negligible since shifts are small (e.g. ~12 for
793 /// SHA/ECDSA), but may matter if signatures grow more expensive to
794 /// construct.
795 fn signature() -> UairSignature<Self::Prime>;
796
797 /// A general method for describing constraints.
798 ///
799 /// # Arguments
800 /// - `b`: a builder encapsulating the constraint storing logic. Its type
801 /// `B` has to have compatible `B::Ideal` with the `Self::Ideal`, i.e. it
802 /// must implement `FromRef<Self::Ideal>` trait.
803 /// - `expr_cfg`: the [`SemiringConfig`] providing arithmetic operations on
804 /// `B::Expr`. Per-family builder runs pass the family's config (e.g. a
805 /// [`zinc_poly::univariate::dynamic::DynamicPolynomialConfig`] over the
806 /// family's field, or the field config itself); static analyses pass a
807 /// `FixedConfig`.
808 /// - `up`: a `TraceRow` of expressions representing the current row of
809 /// UAIR.
810 /// - `down`: a `TraceRow` of expressions representing the shifted (down)
811 /// row of the UAIR. Its layout matches `UairSignature::down()`, which may
812 /// have fewer columns than `up` when only a subset of columns are
813 /// shifted.
814 /// - `from_ref`: a closure that turns the underlying ring `R` into
815 /// `B::Expr`. Sometimes (e.g. when dealing with random fields) it is
816 /// convenient to provide a closure instead of a `FromRef` implementation.
817 /// - `mbs`: a closure that allows to multiply expressions by `R`. Same
818 /// rationale as for `from_ref`.
819 /// - `ideal_from_ref`: a closure that turns a `Self::Ideal` into `B::Ideal`
820 /// for the $Q[X]$-ideal-membership family.
821 /// - `fq_ideal_from_ref`: a closure that turns a `Self::FqIdeal` into
822 /// `B::FqIdeal` for the new $F_{q_i}[X]$-ideal-membership family emitted
823 /// via [`ConstraintBuilder::assert_in_fq_ideal`]. UAIRs without
824 /// $F_q[X]$-constraints can ignore this closure.
825 #[allow(clippy::too_many_arguments)]
826 fn constrain_general<C, B, FromR, MulByScalar, IFromR, IFqFromR>(
827 b: &mut B,
828 expr_cfg: &C,
829 up: TraceRow<C::Element>,
830 down: TraceRow<C::Element>,
831 from_ref: FromR,
832 mbs: MulByScalar,
833 ideal_from_ref: IFromR,
834 fq_ideal_from_ref: IFqFromR,
835 ) where
836 C: SemiringConfig,
837 B: ConstraintBuilder<Expr = C::Element>,
838 FromR: Fn(&Self::Scalar) -> C::Element,
839 MulByScalar: Fn(&C::Element, &Self::Scalar) -> Option<C::Element>,
840 IFromR: Fn(&Self::Ideal) -> B::Ideal,
841 IFqFromR: Fn(&Self::FqIdeal) -> B::FqIdeal;
842}
843
844#[cfg(test)]
845mod tests {
846 use super::*;
847
848 fn signature_with_mixed_shifts() -> UairSignature<u64> {
849 UairSignature::new(
850 TotalColumnLayout::new(2, 1, 1),
851 PublicColumnLayout::new(0, 0, 0),
852 vec![
853 ShiftSpec::new(0, 1),
854 ShiftSpec::new(2, 1),
855 ShiftSpec::new(3, 1),
856 ],
857 vec![],
858 )
859 }
860
861 #[test]
862 fn bit_op_specs_extend_binary_down_layout() {
863 let specs = vec![
864 BitOpSpec::new(1, BitOp::ShR(3)),
865 BitOpSpec::new(0, BitOp::Rot(2)),
866 ];
867 let sig = signature_with_mixed_shifts().with_bit_op_specs(specs.clone());
868
869 assert_eq!(sig.bit_op_specs(), specs);
870 assert_eq!(sig.bit_op_specs()[0].source_col(), 1);
871 assert_eq!(sig.bit_op_specs()[0].op(), BitOp::ShR(3));
872 assert_eq!(sig.bit_op_specs()[0].op().count(), 3);
873 assert_eq!(sig.down_cols().num_binary_poly_cols(), 3);
874 assert_eq!(sig.down_cols().num_arbitrary_poly_cols(), 1);
875 assert_eq!(sig.down_cols().num_int_cols(), 1);
876 assert_eq!(sig.bit_op_down_offset(), 1);
877 }
878
879 #[test]
880 fn empty_bit_op_specs_keep_shift_only_down_layout() {
881 let sig = signature_with_mixed_shifts().with_bit_op_specs(vec![]);
882
883 assert!(sig.bit_op_specs().is_empty());
884 assert_eq!(sig.down_cols().num_binary_poly_cols(), 1);
885 assert_eq!(sig.down_cols().num_arbitrary_poly_cols(), 1);
886 assert_eq!(sig.down_cols().num_int_cols(), 1);
887 assert_eq!(sig.bit_op_down_offset(), 1);
888 }
889
890 #[test]
891 #[should_panic(expected = "bit-op count must be non-zero")]
892 fn bit_op_spec_rejects_zero_count() {
893 let _ = BitOpSpec::new(0, BitOp::Rot(0));
894 }
895
896 #[test]
897 #[should_panic(expected = "is not a binary_poly column")]
898 fn bit_op_specs_reject_non_binary_source() {
899 let _ =
900 signature_with_mixed_shifts().with_bit_op_specs(vec![BitOpSpec::new(2, BitOp::ShR(1))]);
901 }
902
903 #[test]
904 fn affine_virtual_specs_are_attached_in_order() {
905 let specs = vec![
906 AffineVirtualSpec::new(vec![
907 AffineVirtualTerm::new(0, 1),
908 AffineVirtualTerm::new(1, 1),
909 AffineVirtualTerm::new(0, -2),
910 ]),
911 AffineVirtualSpec::with_ones_coefficient(
912 vec![AffineVirtualTerm::new(0, -1), AffineVirtualTerm::new(1, 1)],
913 1,
914 ),
915 ];
916
917 let sig = signature_with_mixed_shifts().with_affine_virtual_specs(specs.clone());
918
919 assert_eq!(sig.affine_virtual_specs(), specs);
920 assert_eq!(sig.affine_virtual_specs()[0].terms()[1].source_col(), 1);
921 assert_eq!(sig.affine_virtual_specs()[0].terms()[1].coefficient(), 1);
922 assert_eq!(sig.affine_virtual_specs()[0].terms()[1].row_shift(), 0);
923 assert_eq!(sig.affine_virtual_specs()[1].ones_coefficient(), 1);
924 assert_eq!(sig.down_cols().num_binary_poly_cols(), 1);
925 assert_eq!(sig.down_cols().num_arbitrary_poly_cols(), 1);
926 assert_eq!(sig.down_cols().num_int_cols(), 1);
927 }
928
929 #[test]
930 fn empty_affine_virtual_specs_are_supported() {
931 let sig = signature_with_mixed_shifts().with_affine_virtual_specs(vec![]);
932
933 assert!(sig.affine_virtual_specs().is_empty());
934 assert_eq!(sig.down_cols().num_binary_poly_cols(), 1);
935 assert_eq!(sig.down_cols().num_arbitrary_poly_cols(), 1);
936 assert_eq!(sig.down_cols().num_int_cols(), 1);
937 }
938
939 #[test]
940 #[should_panic(expected = "affine virtual coefficient must be non-zero")]
941 fn affine_virtual_term_rejects_zero_coefficient() {
942 let _ = AffineVirtualTerm::new(0, 0);
943 }
944
945 #[test]
946 #[should_panic(expected = "affine virtual spec must have at least one non-constant term")]
947 fn affine_virtual_spec_rejects_constant_only_expression() {
948 let _ = AffineVirtualSpec::with_ones_coefficient(vec![], 1);
949 }
950
951 #[test]
952 #[should_panic(expected = "is not a binary_poly column")]
953 fn affine_virtual_specs_reject_non_binary_source() {
954 let _ =
955 signature_with_mixed_shifts().with_affine_virtual_specs(vec![AffineVirtualSpec::new(
956 vec![AffineVirtualTerm::new(2, 1)],
957 )]);
958 }
959
960 #[test]
961 #[should_panic(
962 expected = "shifted affine virtual terms are not supported: spec 0, term 0, row shift 1"
963 )]
964 fn affine_virtual_specs_reject_shifted_terms() {
965 let _ =
966 signature_with_mixed_shifts().with_affine_virtual_specs(vec![AffineVirtualSpec::new(
967 vec![AffineVirtualTerm::new_shifted(0, 1, 1)],
968 )]);
969 }
970}