openvm_static_verifier/field/baby_bear/
extension.rs

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

openvm_static_verifier/field/baby_bear/
extension.rs
