openvm_static_verifier/field/baby_bear/
extension.rs

1use core::array;
2#[cfg(test)]
3use std::{cell::RefCell, vec::Vec};
4
5#[cfg(test)]
6use halo2_base::AssignedValue;
7use halo2_base::{
8    gates::range::RangeChip, halo2_proofs::halo2curves::bn256::Fr, safe_types::SafeBool, Context,
9};
10use itertools::Itertools;
11#[cfg(test)]
12use openvm_stark_sdk::openvm_stark_backend::p3_field::PrimeField64;
13use openvm_stark_sdk::{
14    openvm_stark_backend::p3_field::{
15        extension::{BinomialExtensionField, BinomiallyExtendable},
16        BasedVectorSpace, Field, PrimeCharacteristicRing,
17    },
18    p3_baby_bear::BabyBear,
19};
20
21use crate::{
22    field::baby_bear::{BabyBearChip, BabyBearWire, ReducedBabyBearWire},
23    utils::guarded_debug_assert_eq,
24};
25
26#[cfg(test)]
27pub(crate) struct RecordedExtBaseConst {
28    pub constant: u64,
29    pub cell: AssignedValue<Fr>,
30}
31
32#[cfg(test)]
33thread_local! {
34    static RECORDED_EXT_BASE_CONSTS: RefCell<Vec<RecordedExtBaseConst>> = const { RefCell::new(Vec::new()) };
35}
36
37#[cfg(test)]
38pub(crate) fn clear_recorded_ext_base_consts() {
39    RECORDED_EXT_BASE_CONSTS.with(|records| records.borrow_mut().clear());
40}
41
42#[cfg(test)]
43pub(crate) fn take_recorded_ext_base_consts() -> Vec<RecordedExtBaseConst> {
44    RECORDED_EXT_BASE_CONSTS.with(|records| records.borrow_mut().drain(..).collect())
45}
46
47// irred poly is x^4 - 11
48#[derive(Clone)]
49pub struct BabyBearExt4Chip {
50    pub base: BabyBearChip,
51}
52
53#[derive(Copy, Clone, Debug)]
54pub struct BabyBearExt4Wire(pub [BabyBearWire; 4]);
55
56/// An extension-field wire whose BabyBear basis coefficients are all reduced.
57///
58/// This is the extension-field analogue of `ReducedBabyBearWire`: it is safe for
59/// transcript/hash absorption coefficient-by-coefficient. Converting via
60/// `BabyBearExt4Wire::from` drops that evidence when the value is used by arithmetic
61/// helpers.
62#[derive(Copy, Clone, Debug)]
63pub struct ReducedBabyBearExt4Wire([ReducedBabyBearWire; 4]);
64pub type BabyBearExt4 = BinomialExtensionField<BabyBear, 4>;
65
66impl BabyBearExt4Wire {
67    pub fn to_extension_field(&self) -> BabyBearExt4 {
68        BabyBearExt4::from_basis_coefficients_fn(|i| self.0[i].to_baby_bear())
69    }
70}
71
72impl ReducedBabyBearExt4Wire {
73    pub fn coeffs(&self) -> &[ReducedBabyBearWire; 4] {
74        &self.0
75    }
76}
77
78impl From<ReducedBabyBearExt4Wire> for BabyBearExt4Wire {
79    /// Drops the canonicality evidence and returns the underlying arithmetic wire.
80    fn from(wire: ReducedBabyBearExt4Wire) -> Self {
81        BabyBearExt4Wire(wire.0.map(BabyBearWire::from))
82    }
83}
84
85impl From<&ReducedBabyBearExt4Wire> for BabyBearExt4Wire {
86    fn from(wire: &ReducedBabyBearExt4Wire) -> Self {
87        (*wire).into()
88    }
89}
90
91impl BabyBearExt4Chip {
92    pub fn new(base_chip: BabyBearChip) -> Self {
93        BabyBearExt4Chip { base: base_chip }
94    }
95
96    /// Loads each BabyBear coefficient and constrains only that its assigned
97    /// advice cell fits in 31 bits.
98    ///
99    /// The Rust input is canonicalized for the honest witness assignment, but the
100    /// circuit does not prove each advice cell is `< p`. Use
101    /// `load_reduced_witness` for transcript/hash inputs.
102    pub fn load_witness(&self, ctx: &mut Context<Fr>, value: BabyBearExt4) -> BabyBearExt4Wire {
103        let coeffs = value.as_basis_coefficients_slice();
104        BabyBearExt4Wire(array::from_fn(|i| self.base.load_witness(ctx, coeffs[i])))
105    }
106
107    /// Loads each coefficient and constrains it to the canonical BabyBear range.
108    pub fn load_reduced_witness(
109        &self,
110        ctx: &mut Context<Fr>,
111        value: BabyBearExt4,
112    ) -> ReducedBabyBearExt4Wire {
113        let coeffs = value.as_basis_coefficients_slice();
114        ReducedBabyBearExt4Wire(array::from_fn(|i| {
115            self.base.load_reduced_witness(ctx, coeffs[i])
116        }))
117    }
118
119    /// Loads canonical BabyBear constants for each coefficient and returns them
120    /// with reduced type evidence.
121    pub fn load_reduced_constant(
122        &self,
123        ctx: &mut Context<Fr>,
124        value: BabyBearExt4,
125    ) -> ReducedBabyBearExt4Wire {
126        let coeffs = value.as_basis_coefficients_slice();
127        // Constants are canonical by construction.
128        ReducedBabyBearExt4Wire(array::from_fn(|i| {
129            self.base.load_reduced_constant(ctx, coeffs[i])
130        }))
131    }
132    pub fn load_constant(&self, ctx: &mut Context<Fr>, value: BabyBearExt4) -> BabyBearExt4Wire {
133        let coeffs = value.as_basis_coefficients_slice();
134        BabyBearExt4Wire(array::from_fn(|i| self.base.load_constant(ctx, coeffs[i])))
135    }
136    pub fn add(
137        &self,
138        ctx: &mut Context<Fr>,
139        a: BabyBearExt4Wire,
140        b: BabyBearExt4Wire,
141    ) -> BabyBearExt4Wire {
142        BabyBearExt4Wire(
143            a.0.iter()
144                .zip(b.0.iter())
145                .map(|(a, b)| self.base.add(ctx, *a, *b))
146                .collect_vec()
147                .try_into()
148                .unwrap(),
149        )
150    }
151
152    pub fn neg(&self, ctx: &mut Context<Fr>, a: BabyBearExt4Wire) -> BabyBearExt4Wire {
153        BabyBearExt4Wire(
154            a.0.iter()
155                .map(|x| self.base.neg(ctx, *x))
156                .collect_vec()
157                .try_into()
158                .unwrap(),
159        )
160    }
161
162    pub fn sub(
163        &self,
164        ctx: &mut Context<Fr>,
165        a: BabyBearExt4Wire,
166        b: BabyBearExt4Wire,
167    ) -> BabyBearExt4Wire {
168        BabyBearExt4Wire(
169            a.0.iter()
170                .zip(b.0.iter())
171                .map(|(a, b)| self.base.sub(ctx, *a, *b))
172                .collect_vec()
173                .try_into()
174                .unwrap(),
175        )
176    }
177
178    pub fn scalar_mul(
179        &self,
180        ctx: &mut Context<Fr>,
181        a: BabyBearExt4Wire,
182        b: BabyBearWire,
183    ) -> BabyBearExt4Wire {
184        BabyBearExt4Wire(
185            a.0.iter()
186                .map(|x| self.base.mul(ctx, *x, b))
187                .collect_vec()
188                .try_into()
189                .unwrap(),
190        )
191    }
192
193    /// Fused `a * b + c` where `b` is a base-field scalar.
194    /// Uses `mul_add` gates to save cells vs separate `scalar_mul` + `add`.
195    pub fn scalar_mul_add(
196        &self,
197        ctx: &mut Context<Fr>,
198        a: BabyBearExt4Wire,
199        b: BabyBearWire,
200        c: BabyBearExt4Wire,
201    ) -> BabyBearExt4Wire {
202        BabyBearExt4Wire(
203            a.0.iter()
204                .zip(c.0.iter())
205                .map(|(ai, ci)| self.base.mul_add(ctx, *ai, b, *ci))
206                .collect_vec()
207                .try_into()
208                .unwrap(),
209        )
210    }
211
212    pub fn select(
213        &self,
214        ctx: &mut Context<Fr>,
215        cond: SafeBool<Fr>,
216        a: BabyBearExt4Wire,
217        b: BabyBearExt4Wire,
218    ) -> BabyBearExt4Wire {
219        BabyBearExt4Wire(
220            a.0.iter()
221                .zip(b.0.iter())
222                .map(|(a, b)| self.base.select(ctx, cond, *a, *b))
223                .collect_vec()
224                .try_into()
225                .unwrap(),
226        )
227    }
228
229    pub fn assert_zero(&self, ctx: &mut Context<Fr>, a: BabyBearExt4Wire) {
230        for x in a.0.iter() {
231            self.base.assert_zero(ctx, *x);
232        }
233    }
234
235    pub fn assert_equal(&self, ctx: &mut Context<Fr>, a: BabyBearExt4Wire, b: BabyBearExt4Wire) {
236        for (a, b) in a.0.iter().zip(b.0.iter()) {
237            self.base.assert_equal(ctx, *a, *b);
238        }
239    }
240
241    pub fn mul(
242        &self,
243        ctx: &mut Context<Fr>,
244        mut a: BabyBearExt4Wire,
245        mut b: BabyBearExt4Wire,
246    ) -> BabyBearExt4Wire {
247        let mut coeffs = Vec::with_capacity(7);
248        for s in 0..7 {
249            coeffs.push(self.base.special_inner_product(ctx, &mut a.0, &mut b.0, s));
250        }
251        let w = self
252            .base
253            .load_constant(ctx, <BabyBear as BinomiallyExtendable<4>>::W);
254        for i in 4..7 {
255            coeffs[i - 4] = self.base.mul_add(ctx, coeffs[i], w, coeffs[i - 4]);
256        }
257        coeffs.truncate(4);
258        let c = BabyBearExt4Wire(coeffs.try_into().unwrap());
259        guarded_debug_assert_eq!(
260            c.to_extension_field(),
261            a.to_extension_field() * b.to_extension_field()
262        );
263        c
264    }
265
266    pub fn div(
267        &self,
268        ctx: &mut Context<Fr>,
269        a: BabyBearExt4Wire,
270        b: BabyBearExt4Wire,
271    ) -> BabyBearExt4Wire {
272        let b_val = b.to_extension_field();
273        let b_inv_val = b_val.try_inverse().unwrap();
274        // Constrain b is non-zero by checking b * b_inv == 1
275        let b_inv = self.load_witness(ctx, b_inv_val);
276        let one = self.load_constant(ctx, BinomialExtensionField::<BabyBear, 4>::ONE);
277        let inv_prod = self.mul(ctx, b, b_inv);
278        self.assert_equal(ctx, inv_prod, one);
279
280        // Constrain a = b * c (mod p)
281        let c = self.load_witness(ctx, a.to_extension_field() * b_inv_val);
282        let prod = self.mul(ctx, b, c);
283        self.assert_equal(ctx, a, prod);
284
285        guarded_debug_assert_eq!(
286            c.to_extension_field(),
287            a.to_extension_field() / b.to_extension_field()
288        );
289        c
290    }
291
292    pub fn reduce_max_bits(&self, ctx: &mut Context<Fr>, a: BabyBearExt4Wire) -> BabyBearExt4Wire {
293        BabyBearExt4Wire(
294            a.0.into_iter()
295                .map(|x| self.base.reduce_max_bits(ctx, x))
296                .collect::<Vec<_>>()
297                .try_into()
298                .unwrap(),
299        )
300    }
301
302    pub fn base(&self) -> &BabyBearChip {
303        &self.base
304    }
305
306    pub fn range(&self) -> &RangeChip<Fr> {
307        self.base.range()
308    }
309
310    pub fn zero(&self, ctx: &mut Context<Fr>) -> BabyBearExt4Wire {
311        self.from_base_const(ctx, BabyBear::ZERO)
312    }
313
314    pub fn from_base_const(&self, ctx: &mut Context<Fr>, value: BabyBear) -> BabyBearExt4Wire {
315        let base_val = self.base.load_constant(ctx, value);
316        #[cfg(test)]
317        RECORDED_EXT_BASE_CONSTS.with(|records| {
318            records.borrow_mut().push(RecordedExtBaseConst {
319                constant: value.as_canonical_u64(),
320                cell: base_val.value,
321            });
322        });
323        let z = self.base.load_constant(ctx, BabyBear::ZERO);
324        BabyBearExt4Wire([base_val, z, z, z])
325    }
326
327    pub fn from_base_var(&self, ctx: &mut Context<Fr>, value: BabyBearWire) -> BabyBearExt4Wire {
328        let z = self.base.load_constant(ctx, BabyBear::ZERO);
329        BabyBearExt4Wire([value, z, z, z])
330    }
331
332    pub fn mul_base_const(
333        &self,
334        ctx: &mut Context<Fr>,
335        a: BabyBearExt4Wire,
336        c: BabyBear,
337    ) -> BabyBearExt4Wire {
338        let c_wire = self.base.load_constant(ctx, c);
339        self.scalar_mul(ctx, a, c_wire)
340    }
341
342    pub fn square(&self, ctx: &mut Context<Fr>, a: BabyBearExt4Wire) -> BabyBearExt4Wire {
343        self.mul(ctx, a, a)
344    }
345
346    pub fn pow_power_of_two(
347        &self,
348        ctx: &mut Context<Fr>,
349        a: BabyBearExt4Wire,
350        n: usize,
351    ) -> BabyBearExt4Wire {
352        let mut result = a;
353        for _ in 0..n {
354            result = self.square(ctx, result);
355        }
356        result
357    }
358}
openvm_static_verifier/field/baby_bear/extension.rs

openvm_static_verifier/field/baby_bear/
extension.rs
