zinc_uair/
ideal_collector.rs

1use crate::{
2    ConstraintBuilder, TraceRow, Uair,
3    dummy_semiring::{DUMMY_SEMIRING_CONFIG, DummySemiring, DummySemiringConfig},
4    ideal::{Ideal, IdealCheck, IdealCheckError},
5};
6use std::fmt::{Display, Formatter};
7use zinc_utils::{add, from_ref::FromRef};
8
9/// A `ConstraintBuilder` that collects ideals used in a `Uair`.
10///
11/// Both $Q[X]$-only-ideals (from [`ConstraintBuilder::assert_in_ideal`] /
12/// [`ConstraintBuilder::assert_zero`]) and new $F_{q_i}[X]$-ideals (from
13/// [`ConstraintBuilder::assert_in_fq_ideal`]) are kept in separate vectors, so
14/// downstream consumers (the PIOP layer) can dispatch them independently. The
15/// $F_q[X]$ ideals are wrapped in [`IdealOrZero`] to reuse the unconditional
16/// `IdealCheck` (over `DummySemiring`) proxy impl used during the collection;
17/// the `Zero` variant is never produced by the collection (there is no
18/// `assert_fq_zero`) but downstream consumers may construct it.
19pub struct IdealCollector<I: Ideal, IFq: Ideal> {
20    pub ideals: Vec<IdealOrZero<I>>,
21    /// $F_{q_i}[X]$-ideals indexed by their `prime_idx` into the
22    /// owning UAIR's [`crate::UairSignature::primes`] tuple. Empty for
23    /// UAIRs with $Q[X]$-only constraints.
24    pub fq_ideals: Vec<Vec<IdealOrZero<IFq>>>,
25}
26
27impl<I: Ideal, IFq: Ideal> IdealCollector<I, IFq> {
28    /// Create a new ideal collector
29    /// and hint the number of constraints
30    /// a target UAIR might have.
31    pub fn new(num_constraints: usize) -> Self {
32        Self {
33            ideals: Vec::with_capacity(num_constraints),
34            fq_ideals: Vec::new(),
35        }
36    }
37}
38
39/// Given a `Uair` and a hint of how many constraints
40/// it is going to have, creates an `IdealCollector`
41/// object and collects ideals from the `Uair`.
42pub fn collect_ideals<U: Uair>(num_constraints: usize) -> IdealCollector<U::Ideal, U::FqIdeal> {
43    let mut ideal_collector = IdealCollector::new(num_constraints);
44
45    let sig = U::signature();
46    let (up_dummy, down_dummy) = sig.dummy_rows(DummySemiring);
47    let up_row = TraceRow::from_slice_with_layout(&up_dummy, sig.total_cols().as_column_layout());
48    let down_row =
49        TraceRow::from_slice_with_layout(&down_dummy, sig.down_cols().as_column_layout());
50    U::constrain_general(
51        &mut ideal_collector,
52        &DUMMY_SEMIRING_CONFIG,
53        up_row,
54        down_row,
55        |_| DummySemiring,
56        |_, _| Some(DummySemiring),
57        IdealOrZero::from_ref,
58        IdealOrZero::from_ref,
59    );
60
61    ideal_collector
62}
63
64impl<I, IFq> ConstraintBuilder for IdealCollector<I, IFq>
65where
66    I: Ideal,
67    IFq: Ideal,
68{
69    type Expr = DummySemiring;
70    type Ideal = IdealOrZero<I>;
71    type FqIdeal = IdealOrZero<IFq>;
72
73    fn assert_in_ideal(&mut self, _expr: Self::Expr, ideal: &Self::Ideal) {
74        self.ideals.push(ideal.clone());
75    }
76
77    fn assert_zero(&mut self, _expr: Self::Expr) {
78        self.ideals.push(IdealOrZero::zero());
79    }
80
81    fn assert_in_fq_ideal(&mut self, prime_idx: usize, _expr: Self::Expr, ideal: &Self::FqIdeal) {
82        if self.fq_ideals.len() <= prime_idx {
83            self.fq_ideals.resize(add!(prime_idx, 1), Vec::new());
84        }
85        self.fq_ideals[prime_idx].push(ideal.clone());
86    }
87}
88
89/// Either a non-trivial ideal (from `assert_in_ideal`) or the zero ideal
90/// (from `assert_zero`).
91#[derive(Clone, Copy, Debug)]
92pub enum IdealOrZero<I: Ideal> {
93    NonZero(I),
94    Zero,
95}
96
97impl<I: Ideal> IdealOrZero<I> {
98    pub fn zero() -> Self {
99        IdealOrZero::Zero
100    }
101
102    /// Returns `true` if this is the zero ideal
103    /// (i.e., the ideal used by `assert_zero` constraints).
104    pub fn is_zero_ideal(&self) -> bool {
105        matches!(self, IdealOrZero::Zero)
106    }
107
108    pub fn map<I2: Ideal>(&self, f: impl FnOnce(&I) -> I2) -> IdealOrZero<I2> {
109        match self {
110            IdealOrZero::NonZero(ideal) => IdealOrZero::NonZero(f(ideal)),
111            IdealOrZero::Zero => IdealOrZero::Zero,
112        }
113    }
114}
115
116impl<I: Ideal> Display for IdealOrZero<I> {
117    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
118        write!(f, "IdealOrZero<")?;
119        match self {
120            Self::Zero => write!(f, "Zero")?,
121            Self::NonZero(ideal) => write!(f, "{ideal}")?,
122        }
123        write!(f, ">")?;
124        Ok(())
125    }
126}
127
128impl<I: Ideal> Ideal for IdealOrZero<I> {}
129
130impl<I: Ideal> FromRef<IdealOrZero<I>> for IdealOrZero<I> {
131    fn from_ref(value: &IdealOrZero<I>) -> Self {
132        value.clone()
133    }
134}
135
136impl<I: Ideal> FromRef<I> for IdealOrZero<I> {
137    fn from_ref(value: &I) -> Self {
138        IdealOrZero::NonZero(value.clone())
139    }
140}
141
142impl<I: Ideal> IdealCheck<DummySemiringConfig> for IdealOrZero<I> {
143    fn contains(
144        &self,
145        _cfg: &DummySemiringConfig,
146        _value: &DummySemiring,
147    ) -> Result<bool, IdealCheckError> {
148        // Do nothing.
149        Ok(true)
150    }
151}
zinc_uair/ideal_collector.rs

zinc_uair/
ideal_collector.rs
