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 crypto_primitives::Semiring;
13use std::borrow::Cow;
14use zinc_poly::{
15    mle::DenseMultilinearExtension,
16    univariate::{binary::BinaryPoly, dense::DensePolynomial},
17};
18use zinc_utils::{UNCHECKED, add, from_ref::FromRef, mul_by_scalar::MulByScalar, sub};
19
20use crate::ideal::{Ideal, IdealCheck};
21
22pub use lookup_types::{LookupColumnSpec, LookupTableType};
23
24/// The abstract interface to constraint building logic.
25/// In essence it allows to create constraints modulo ideals.
26pub trait ConstraintBuilder {
27    /// The expressions the constraint builder operates on.
28    /// It is opaque from the PoV of an AIR apart from
29    /// the fact that arithmetic operations are available on it
30    /// and one can check if an expression is in an ideal.
31    type Expr: Semiring;
32    /// The type of ideals used by the constraint builder.
33    type Ideal: Ideal + IdealCheck<Self::Expr>;
34
35    /// Add a constraint saying that `expr` belongs to the ideal `ideal`.
36    fn assert_in_ideal(&mut self, expr: Self::Expr, ideal: &Self::Ideal);
37
38    /// Add a constraint saying that `expr` is equal to zero which is
39    /// the same as saying that `expr` belongs to the zero ideal.
40    fn assert_zero(&mut self, expr: Self::Expr);
41}
42
43/// Specifies a shifted column
44/// `ShiftSpec { source_col: 0, shift_amount: 3 }` means
45/// "virtual column whose row i is the value of column 0 at row i+3
46/// (zero-padded beyond trace length)."
47///
48/// Multiple ShiftSpecs may reference the same source_col with
49/// different shift amounts.
50#[derive(Clone, Debug, PartialEq, Eq, Hash)]
51pub struct ShiftSpec {
52    /// Index of the committed column in the flattened trace
53    /// (binary_poly || arbitrary_poly || int, same indexing as
54    /// TraceRow::from_slice_with_layout).
55    source_col: usize,
56    /// Number of rows to shift by.
57    shift_amount: usize,
58}
59
60impl ShiftSpec {
61    pub fn new(source_col: usize, shift_amount: usize) -> Self {
62        assert!(shift_amount > 0, "shift must be non-zero");
63        Self {
64            source_col,
65            shift_amount,
66        }
67    }
68
69    pub fn source_col(&self) -> usize {
70        self.source_col
71    }
72
73    pub fn shift_amount(&self) -> usize {
74        self.shift_amount
75    }
76}
77
78// ---------------------------------------------------------------------------
79// Column layout types
80// ---------------------------------------------------------------------------
81
82/// Column counts per type (binary_poly, arbitrary_poly, int).
83/// Shared internals for the semantic newtype wrappers (Total, Public, Virtual,
84/// Witness)
85#[derive(Clone, Debug, Default)]
86pub struct ColumnLayout {
87    num_binary_poly_cols: usize,
88    num_arbitrary_poly_cols: usize,
89    num_int_cols: usize,
90}
91
92impl ColumnLayout {
93    pub fn new(
94        num_binary_poly_cols: usize,
95        num_arbitrary_poly_cols: usize,
96        num_int_cols: usize,
97    ) -> Self {
98        Self {
99            num_binary_poly_cols,
100            num_arbitrary_poly_cols,
101            num_int_cols,
102        }
103    }
104
105    pub fn num_binary_poly_cols(&self) -> usize {
106        self.num_binary_poly_cols
107    }
108
109    pub fn num_arbitrary_poly_cols(&self) -> usize {
110        self.num_arbitrary_poly_cols
111    }
112
113    pub fn num_int_cols(&self) -> usize {
114        self.num_int_cols
115    }
116
117    /// Maximum number of columns across the three types.
118    pub fn max_cols(&self) -> usize {
119        [
120            self.num_binary_poly_cols,
121            self.num_arbitrary_poly_cols,
122            self.num_int_cols,
123        ]
124        .into_iter()
125        .max()
126        .expect("the iterator is not empty")
127    }
128
129    /// The sum of the numbers of columns across all types.
130    #[allow(clippy::arithmetic_side_effects)]
131    pub fn cols(&self) -> usize {
132        self.num_binary_poly_cols + self.num_arbitrary_poly_cols + self.num_int_cols
133    }
134}
135
136macro_rules! column_layout_wrapper {
137    ($(#[$meta:meta])* $name:ident) => {
138        $(#[$meta])*
139        #[derive(Clone, Debug, Default)]
140        pub struct $name(ColumnLayout);
141
142        impl $name {
143            pub fn new(num_binary_poly_cols: usize, num_arbitrary_poly_cols: usize, num_int_cols: usize) -> Self {
144                Self(ColumnLayout::new(num_binary_poly_cols, num_arbitrary_poly_cols, num_int_cols))
145            }
146
147            pub fn num_binary_poly_cols(&self) -> usize { self.0.num_binary_poly_cols() }
148            pub fn num_arbitrary_poly_cols(&self) -> usize { self.0.num_arbitrary_poly_cols() }
149            pub fn num_int_cols(&self) -> usize { self.0.num_int_cols() }
150            pub fn max_cols(&self) -> usize { self.0.max_cols() }
151            pub fn cols(&self) -> usize { self.0.cols() }
152            pub fn as_column_layout(&self) -> &ColumnLayout { &self.0 }
153        }
154    };
155}
156
157column_layout_wrapper!(/// Layout of all trace columns (public + witness) per type.
158    TotalColumnLayout);
159column_layout_wrapper!(/// Layout of the public column subset.
160    PublicColumnLayout);
161column_layout_wrapper!(/// Layout of the virtual (shifted/down) columns.
162    VirtualColumnLayout);
163column_layout_wrapper!(/// Layout of the witness (total minus public) columns.
164    WitnessColumnLayout);
165
166// ---------------------------------------------------------------------------
167// UairSignature
168// ---------------------------------------------------------------------------
169
170/// The signature of a UAIR.
171///
172/// Public columns precede witness columns within each type group.
173/// The flattened trace ordering is:
174/// `[pub_bin, wit_bin, pub_arb, wit_arb, pub_int, wit_int]`.
175#[derive(Clone, Debug)]
176pub struct UairSignature {
177    /// Column-type layout of all (public + witness) columns.
178    total_cols: TotalColumnLayout,
179    /// Public column subset.
180    public_cols: PublicColumnLayout,
181    /// Witness column counts (total minus public) per type.
182    witness_cols: WitnessColumnLayout,
183    /// Shifted columns info sorted by `source_col`.
184    shifts: Vec<ShiftSpec>,
185    /// Column-type layout of the shifted (down) row.
186    down_cols: VirtualColumnLayout,
187    /// Lookup specifications: which trace columns are constrained against
188    /// which table types.
189    lookup_specs: Vec<LookupColumnSpec>,
190}
191
192impl UairSignature {
193    /// Create a new signature, sorting `shifts` by `source_col`.
194    pub fn new(
195        total_cols: TotalColumnLayout,
196        public_cols: PublicColumnLayout,
197        mut shifts: Vec<ShiftSpec>,
198        lookup_specs: Vec<LookupColumnSpec>,
199    ) -> Self {
200        for (name, pub_n, tot_n) in [
201            (
202                "binary_poly",
203                public_cols.num_binary_poly_cols(),
204                total_cols.num_binary_poly_cols(),
205            ),
206            (
207                "arbitrary_poly",
208                public_cols.num_arbitrary_poly_cols(),
209                total_cols.num_arbitrary_poly_cols(),
210            ),
211            ("int", public_cols.num_int_cols(), total_cols.num_int_cols()),
212        ] {
213            assert!(
214                pub_n <= tot_n,
215                "public {name}_cols ({pub_n}) > total ({tot_n})"
216            );
217        }
218
219        let num_cols = total_cols.cols();
220        for spec in &shifts {
221            assert!(
222                spec.source_col() < num_cols,
223                "ShiftSpec source_col {} out of range (total_cols = {}). \
224                 source_col uses flat indexing: binary_poly || arbitrary_poly || int.",
225                spec.source_col(),
226                num_cols,
227            );
228        }
229
230        shifts.sort_by_key(|spec| spec.source_col());
231        let down_cols = Self::compute_down_layout(&total_cols, &shifts);
232        let witness_cols = WitnessColumnLayout::new(
233            sub!(
234                total_cols.num_binary_poly_cols(),
235                public_cols.num_binary_poly_cols()
236            ),
237            sub!(
238                total_cols.num_arbitrary_poly_cols(),
239                public_cols.num_arbitrary_poly_cols()
240            ),
241            sub!(total_cols.num_int_cols(), public_cols.num_int_cols()),
242        );
243
244        Self {
245            total_cols,
246            public_cols,
247            shifts,
248            down_cols,
249            witness_cols,
250            lookup_specs,
251        }
252    }
253
254    pub fn lookup_specs(&self) -> &[LookupColumnSpec] {
255        &self.lookup_specs
256    }
257
258    fn compute_down_layout(
259        total_cols: &TotalColumnLayout,
260        shifts: &[ShiftSpec],
261    ) -> VirtualColumnLayout {
262        let binary_poly_end = total_cols.num_binary_poly_cols();
263        let arbitrary_poly_end = add!(binary_poly_end, total_cols.num_arbitrary_poly_cols());
264        let mut num_binary_poly = 0usize;
265        let mut num_arbitrary_poly = 0usize;
266        let mut num_int = 0usize;
267        for spec in shifts {
268            if spec.source_col() < binary_poly_end {
269                num_binary_poly = add!(num_binary_poly, 1);
270            } else if spec.source_col() < arbitrary_poly_end {
271                num_arbitrary_poly = add!(num_arbitrary_poly, 1);
272            } else {
273                num_int = add!(num_int, 1);
274            }
275        }
276        VirtualColumnLayout::new(num_binary_poly, num_arbitrary_poly, num_int)
277    }
278
279    pub fn total_cols(&self) -> &TotalColumnLayout {
280        &self.total_cols
281    }
282
283    pub fn public_cols(&self) -> &PublicColumnLayout {
284        &self.public_cols
285    }
286
287    /// Witness column counts (total minus public) per type.
288    pub fn witness_cols(&self) -> &WitnessColumnLayout {
289        &self.witness_cols
290    }
291
292    pub fn shifts(&self) -> &[ShiftSpec] {
293        &self.shifts
294    }
295
296    /// Column-type layout of the shifted (down) row.
297    pub fn down_cols(&self) -> &VirtualColumnLayout {
298        &self.down_cols
299    }
300
301    /// Build correctly-sized dummy up and down `TraceRow`s for static
302    /// analysis (constraint counting, degree counting, scalar/ideal
303    /// collection).
304    pub fn dummy_rows<T: Clone>(&self, val: T) -> (Vec<T>, Vec<T>) {
305        let up_size = self.total_cols.cols();
306        let down_size = self.down_cols.cols();
307        (vec![val.clone(); up_size], vec![val; down_size])
308    }
309}
310
311// ---------------------------------------------------------------------------
312// UairTrace
313// ---------------------------------------------------------------------------
314
315/// The trace of a UAIR execution (pre-projection).
316/// If owned, it contains the full trace, otherwise it contains a view on the
317/// full trace (e.g. only public columns).
318#[derive(Debug, Clone, Default)]
319pub struct UairTrace<'a, PolyCoeff: Clone, Int: Clone, const D: usize> {
320    pub binary_poly: Cow<'a, [DenseMultilinearExtension<BinaryPoly<D>>]>,
321    pub arbitrary_poly: Cow<'a, [DenseMultilinearExtension<DensePolynomial<PolyCoeff, D>>]>,
322    pub int: Cow<'a, [DenseMultilinearExtension<Int>]>,
323}
324
325impl<PolyCoeff: Clone, Int: Clone, const D: usize> UairTrace<'static, PolyCoeff, Int, D> {
326    /// Returns a sub-trace containing only public columns.
327    /// Returned trace is borrowed from the full trace.
328    pub fn public(&self, sig: &UairSignature) -> UairTrace<'_, PolyCoeff, Int, D> {
329        let p = sig.public_cols();
330        UairTrace {
331            binary_poly: Cow::Borrowed(&self.binary_poly[0..p.num_binary_poly_cols()]),
332            arbitrary_poly: Cow::Borrowed(&self.arbitrary_poly[0..p.num_arbitrary_poly_cols()]),
333            int: Cow::Borrowed(&self.int[0..p.num_int_cols()]),
334        }
335    }
336
337    /// Returns a sub-trace containing only witness columns.
338    /// Returned trace is borrowed from the full trace.
339    pub fn witness(&self, sig: &UairSignature) -> UairTrace<'_, PolyCoeff, Int, D> {
340        let p = sig.public_cols();
341        UairTrace {
342            binary_poly: Cow::Borrowed(&self.binary_poly[p.num_binary_poly_cols()..]),
343            arbitrary_poly: Cow::Borrowed(&self.arbitrary_poly[p.num_arbitrary_poly_cols()..]),
344            int: Cow::Borrowed(&self.int[p.num_int_cols()..]),
345        }
346    }
347}
348
349// ---------------------------------------------------------------------------
350// TraceRow
351// ---------------------------------------------------------------------------
352
353/// A view on a row of the trace.
354/// Contains references to cells of the trace
355/// of all types lying in the same trace row.
356#[derive(Clone, Copy)]
357pub struct TraceRow<'a, Expr> {
358    pub binary_poly: &'a [Expr],
359    pub arbitrary_poly: &'a [Expr],
360    pub int: &'a [Expr],
361}
362
363impl<'a, Expr> TraceRow<'a, Expr> {
364    /// Given a slice that represents a raw row of the trace,
365    /// creates a `TraceRow` from it.
366    /// Subdivides the slice according to the given column layout.
367    #[allow(clippy::arithmetic_side_effects)]
368    pub fn from_slice_with_layout(row: &'a [Expr], layout: &ColumnLayout) -> Self {
369        let num_binary_poly = layout.num_binary_poly_cols();
370        let num_arbitrary_poly = layout.num_arbitrary_poly_cols();
371        Self {
372            binary_poly: &row[0..num_binary_poly],
373            arbitrary_poly: &row[num_binary_poly..num_binary_poly + num_arbitrary_poly],
374            int: &row[num_binary_poly + num_arbitrary_poly..],
375        }
376    }
377}
378
379// ---------------------------------------------------------------------------
380// Uair trait
381// ---------------------------------------------------------------------------
382
383/// The trait that a universal AIR description has to implement.
384/// This must include all the constraint description logic of an UAIR.
385///
386/// One type might implement different UAIR logics for different underlying
387/// semirings hence the generic type parameter.
388pub trait Uair: Clone {
389    /// The ideal type the AIR operates with.
390    /// Since a `ConstraintBuilder` is "opaque" for a `Uair`
391    /// a `Uair` has to have a means to create ideals
392    /// so ideals are fixed by this associated types.
393    /// At the `constrain*` methods a `Uair` is given
394    /// a way to convert its own ideals into builder's ideals
395    /// via the `FromRef` trait.
396    type Ideal: Ideal;
397
398    /// The type of scalars of the UAIR.
399    /// For now, we assume they are of
400    /// the type "arbitrary polynomials".
401    // Note: This is usually Z_32[X] (i.e. DensePolynomial<Ring, 32>), but according
402    // to @agareta, this in not always the case.
403    type Scalar: Semiring;
404
405    /// Signature of the UAIR.
406    ///
407    /// TODO: Consider caching the signature to avoid recomputing it at every
408    /// call site. Currently negligible since shifts are small (e.g. ~12 for
409    /// SHA/ECDSA), but may matter if signatures grow more expensive to
410    /// construct.
411    fn signature() -> UairSignature;
412
413    /// A general method for describing constraints.
414    ///
415    /// # Arguments
416    /// - `b`: a builder encapsulating the constraint storing logic. Its type
417    ///   `B` has to have compatible `B::Ideal` with the `Self::Ideal`, i.e. it
418    ///   must implement `FromRef<Self::Ideal>` trait.
419    /// - `up`: a `TraceRow` of expressions representing the current row of
420    ///   UAIR.
421    /// - `down`: a `TraceRow` of expressions representing the shifted (down)
422    ///   row of the UAIR. Its layout matches `UairSignature::down()`, which may
423    ///   have fewer columns than `up` when only a subset of columns are
424    ///   shifted.
425    /// - `from_ref`: a closure that turns the underlying ring `R` into
426    ///   `B::Expr`. Sometimes (e.g. when dealing with random fields) it is
427    ///   convenient to provide a closure instead of a `FromRef` implementation.
428    /// - `mbs`: a closure that allows to multiply expressions by `R`. Same
429    ///   rationale as for `from_ref`.
430    fn constrain_general<B, FromR, MulByScalar, IFromR>(
431        b: &mut B,
432        up: TraceRow<B::Expr>,
433        down: TraceRow<B::Expr>,
434        from_ref: FromR,
435        mbs: MulByScalar,
436        ideal_from_ref: IFromR,
437    ) where
438        B: ConstraintBuilder,
439        FromR: Fn(&Self::Scalar) -> B::Expr,
440        MulByScalar: Fn(&B::Expr, &Self::Scalar) -> Option<B::Expr>,
441        IFromR: Fn(&Self::Ideal) -> B::Ideal;
442
443    // Same as `constrain_general` but `from_ref` and `mbs`
444    // come from the trait implementations.
445    fn constrain<B>(b: &mut B, up: TraceRow<B::Expr>, down: TraceRow<B::Expr>)
446    where
447        B: ConstraintBuilder,
448        B::Expr: FromRef<Self::Scalar> + for<'b> MulByScalar<&'b Self::Scalar>,
449        B::Ideal: FromRef<Self::Ideal>,
450    {
451        Self::constrain_general(
452            b,
453            up,
454            down,
455            B::Expr::from_ref,
456            |x, y| B::Expr::mul_by_scalar::<UNCHECKED>(x, y),
457            B::Ideal::from_ref,
458        )
459    }
460}
zinc_uair/lib.rs

zinc_uair/
lib.rs
