1use std::{cell::RefCell, collections::HashMap, sync::Arc};
2
3use halo2_base::{
4 gates::{GateChip, GateInstructions, RangeChip, RangeInstructions},
5 halo2_proofs::{
6 arithmetic::Field as _,
7 halo2curves::{bn256::Fr, ff::PrimeField as _},
8 },
9 safe_types::SafeBool,
10 utils::{bigint_to_fe, biguint_to_fe, bit_length, fe_to_bigint, modulus, BigPrimeField},
11 AssignedValue, Context, QuantumCell,
12};
13use itertools::Itertools;
14use num_bigint::{BigInt, BigUint};
15use num_integer::Integer;
16use openvm_stark_sdk::{
17 openvm_stark_backend::p3_field::{Field, PrimeCharacteristicRing, PrimeField32, PrimeField64},
18 p3_baby_bear::BabyBear,
19};
20
21use super::BABY_BEAR_MODULUS_U64;
22use crate::utils::{guarded_debug_assert, guarded_debug_assert_eq};
23
24pub(crate) const BABYBEAR_MAX_BITS: usize = 31;
25const RESERVED_HIGH_BITS: usize = 2;
29
30#[derive(Copy, Clone, Debug)]
31pub struct BabyBearWire {
32 pub value: AssignedValue<Fr>,
40 pub max_bits: usize,
42}
43
44#[derive(Copy, Clone, Debug)]
51pub struct ReducedBabyBearWire(BabyBearWire);
52
53impl ReducedBabyBearWire {
54 pub fn value(&self) -> AssignedValue<Fr> {
55 self.0.value
56 }
57}
58
59impl From<ReducedBabyBearWire> for BabyBearWire {
60 fn from(wire: ReducedBabyBearWire) -> Self {
62 wire.0
63 }
64}
65
66impl From<&ReducedBabyBearWire> for BabyBearWire {
67 fn from(wire: &ReducedBabyBearWire) -> Self {
68 (*wire).into()
69 }
70}
71
72impl BabyBearWire {
73 pub fn to_baby_bear(&self) -> BabyBear {
74 let mut b_int = fe_to_bigint(self.value.value()) % BabyBear::ORDER_U32;
75 if b_int < BigInt::from(0) {
76 b_int += BabyBear::ORDER_U32;
77 }
78 BabyBear::from_u32(b_int.try_into().unwrap())
79 }
80
81 pub fn as_u64(&self) -> u64 {
82 PrimeField64::as_canonical_u64(&self.to_baby_bear())
83 }
84}
85
86#[derive(Clone, Debug)]
87pub struct BabyBearChip {
88 pub range: Arc<RangeChip<Fr>>,
89 const_cache: RefCell<HashMap<u64, BabyBearWire>>,
91}
92
93impl BabyBearChip {
94 pub fn new(range_chip: Arc<RangeChip<Fr>>) -> Self {
95 BabyBearChip {
96 range: range_chip,
97 const_cache: RefCell::new(HashMap::new()),
98 }
99 }
100
101 pub fn gate(&self) -> &GateChip<Fr> {
102 self.range.gate()
103 }
104
105 pub fn range(&self) -> &RangeChip<Fr> {
106 &self.range
107 }
108
109 pub fn load_witness(&self, ctx: &mut Context<Fr>, value: BabyBear) -> BabyBearWire {
116 let value = ctx.load_witness(Fr::from(PrimeField64::as_canonical_u64(&value)));
117 self.range.range_check(ctx, value, BABYBEAR_MAX_BITS);
118 BabyBearWire {
119 value,
120 max_bits: BABYBEAR_MAX_BITS,
121 }
122 }
123
124 pub fn load_reduced_witness(
126 &self,
127 ctx: &mut Context<Fr>,
128 value: BabyBear,
129 ) -> ReducedBabyBearWire {
130 let value = ctx.load_witness(Fr::from(PrimeField64::as_canonical_u64(&value)));
131 self.range
132 .check_less_than_safe(ctx, value, BABY_BEAR_MODULUS_U64);
133 ReducedBabyBearWire(BabyBearWire {
134 value,
135 max_bits: BABYBEAR_MAX_BITS,
136 })
137 }
138
139 pub fn load_constant(&self, ctx: &mut Context<Fr>, value: BabyBear) -> BabyBearWire {
140 let key = value.as_canonical_u64();
141 if let Some(&cached) = self.const_cache.borrow().get(&key) {
142 return cached;
143 }
144 let max_bits = bit_length(key);
145 let assigned = if value == BabyBear::ZERO {
146 ctx.load_zero()
147 } else {
148 ctx.load_constant(Fr::from(key))
149 };
150 let wire = BabyBearWire {
151 value: assigned,
152 max_bits,
153 };
154 self.const_cache.borrow_mut().insert(key, wire);
155 wire
156 }
157
158 pub fn load_reduced_constant(
160 &self,
161 ctx: &mut Context<Fr>,
162 value: BabyBear,
163 ) -> ReducedBabyBearWire {
164 ReducedBabyBearWire(self.load_constant(ctx, value))
166 }
167
168 pub fn reduce(&self, ctx: &mut Context<Fr>, a: BabyBearWire) -> BabyBearWire {
169 assert!(a.max_bits <= Fr::CAPACITY as usize - RESERVED_HIGH_BITS);
170 guarded_debug_assert!(fe_to_bigint(a.value.value()).bits() as usize <= a.max_bits);
171 let (_, r) = signed_div_mod(&self.range, ctx, a.value, a.max_bits);
172 let r = BabyBearWire {
173 value: r,
174 max_bits: BABYBEAR_MAX_BITS,
175 };
176 guarded_debug_assert_eq!(a.to_baby_bear(), r.to_baby_bear());
177 r
178 }
179
180 pub fn reduce_max_bits(&self, ctx: &mut Context<Fr>, a: BabyBearWire) -> BabyBearWire {
183 if a.max_bits > BABYBEAR_MAX_BITS {
184 self.reduce(ctx, a)
185 } else {
186 a
187 }
188 }
189
190 pub fn add(
191 &self,
192 ctx: &mut Context<Fr>,
193 mut a: BabyBearWire,
194 mut b: BabyBearWire,
195 ) -> BabyBearWire {
196 if a.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
197 a = self.reduce(ctx, a);
198 }
199 if b.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
200 b = self.reduce(ctx, b);
201 }
202 let value = self.gate().add(ctx, a.value, b.value);
203 let max_bits = a.max_bits.max(b.max_bits) + 1;
204 let c = BabyBearWire { value, max_bits };
205 guarded_debug_assert_eq!(c.to_baby_bear(), a.to_baby_bear() + b.to_baby_bear());
206 c
207 }
208
209 pub fn neg(&self, ctx: &mut Context<Fr>, a: BabyBearWire) -> BabyBearWire {
210 let value = self.gate().neg(ctx, a.value);
211 let b = BabyBearWire {
212 value,
213 max_bits: a.max_bits,
214 };
215 guarded_debug_assert_eq!(b.to_baby_bear(), -a.to_baby_bear());
216 b
217 }
218
219 pub fn sub(
220 &self,
221 ctx: &mut Context<Fr>,
222 mut a: BabyBearWire,
223 mut b: BabyBearWire,
224 ) -> BabyBearWire {
225 #[cfg(debug_assertions)]
226 let expected = a.to_baby_bear() - b.to_baby_bear();
227 if a.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
228 a = self.reduce(ctx, a);
229 }
230 if b.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
231 b = self.reduce(ctx, b);
232 }
233 let value = self.gate().sub(ctx, a.value, b.value);
234 let max_bits = a.max_bits.max(b.max_bits) + 1;
235 let c = BabyBearWire { value, max_bits };
236 guarded_debug_assert_eq!(c.to_baby_bear(), expected);
237 c
238 }
239
240 pub fn mul(
241 &self,
242 ctx: &mut Context<Fr>,
243 mut a: BabyBearWire,
244 mut b: BabyBearWire,
245 ) -> BabyBearWire {
246 if a.max_bits < b.max_bits {
247 std::mem::swap(&mut a, &mut b);
248 }
249 if a.max_bits + b.max_bits > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
250 a = self.reduce(ctx, a);
251 if a.max_bits + b.max_bits > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
252 b = self.reduce(ctx, b);
253 }
254 }
255 let value = self.gate().mul(ctx, a.value, b.value);
256 let max_bits = a.max_bits + b.max_bits;
257
258 let c = BabyBearWire { value, max_bits };
259 guarded_debug_assert_eq!(c.to_baby_bear(), a.to_baby_bear() * b.to_baby_bear());
260 c
261 }
262
263 pub fn mul_add(
264 &self,
265 ctx: &mut Context<Fr>,
266 mut a: BabyBearWire,
267 mut b: BabyBearWire,
268 mut c: BabyBearWire,
269 ) -> BabyBearWire {
270 if a.max_bits < b.max_bits {
271 std::mem::swap(&mut a, &mut b);
272 }
273 if a.max_bits + b.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
274 a = self.reduce(ctx, a);
275 if a.max_bits + b.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
276 b = self.reduce(ctx, b);
277 }
278 }
279 if c.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
280 c = self.reduce(ctx, c)
281 }
282 let value = self.gate().mul_add(ctx, a.value, b.value, c.value);
283 let max_bits = c.max_bits.max(a.max_bits + b.max_bits) + 1;
284
285 let d = BabyBearWire { value, max_bits };
286 guarded_debug_assert_eq!(
287 d.to_baby_bear(),
288 a.to_baby_bear() * b.to_baby_bear() + c.to_baby_bear()
289 );
290 d
291 }
292
293 pub fn div(
294 &self,
295 ctx: &mut Context<Fr>,
296 mut a: BabyBearWire,
297 mut b: BabyBearWire,
298 ) -> BabyBearWire {
299 let b_val = b.to_baby_bear();
300 let b_inv_val = b_val.try_inverse().unwrap();
301 let b_inv = self.load_witness(ctx, b_inv_val);
303 let one = self.load_constant(ctx, BabyBear::ONE);
304 let inv_prod = self.mul(ctx, b, b_inv);
305 self.assert_equal(ctx, inv_prod, one);
306
307 let mut c = self.load_witness(ctx, a.to_baby_bear() * b_inv_val);
309 if a.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
310 a = self.reduce(ctx, a);
311 }
312 if b.max_bits + c.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
313 b = self.reduce(ctx, b);
314 }
315 if b.max_bits + c.max_bits + 1 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS {
316 c = self.reduce(ctx, c);
317 }
318 let diff = self.gate().sub_mul(ctx, a.value, b.value, c.value);
319 let max_bits = a.max_bits.max(b.max_bits + c.max_bits) + 1;
320 self.assert_zero(
321 ctx,
322 BabyBearWire {
323 value: diff,
324 max_bits,
325 },
326 );
327 guarded_debug_assert_eq!(c.to_baby_bear(), a.to_baby_bear() / b.to_baby_bear());
328 c
329 }
330
331 pub(super) fn special_inner_product(
334 &self,
335 ctx: &mut Context<Fr>,
336 a: &mut [BabyBearWire],
337 b: &mut [BabyBearWire],
338 s: usize,
339 ) -> BabyBearWire {
340 assert!(a.len() == b.len());
341 assert!(a.len() == 4);
342 let mut max_bits = 0;
343 let lb = s.saturating_sub(3);
344 let ub = 4.min(s + 1);
345 let range = lb..ub;
346 let other_range = (s + 1 - ub)..(s + 1 - lb);
347 let len = if s < 3 { s + 1 } else { 7 - s };
348 for (i, (c, d)) in a[range.clone()]
349 .iter_mut()
350 .zip(b[other_range.clone()].iter_mut().rev())
351 .enumerate()
352 {
353 if c.max_bits + d.max_bits > Fr::CAPACITY as usize - RESERVED_HIGH_BITS - len + i {
354 if c.max_bits >= d.max_bits {
355 *c = self.reduce(ctx, *c);
356 if c.max_bits + d.max_bits
357 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS - len + i
358 {
359 *d = self.reduce(ctx, *d);
360 }
361 } else {
362 *d = self.reduce(ctx, *d);
363 if c.max_bits + d.max_bits
364 > Fr::CAPACITY as usize - RESERVED_HIGH_BITS - len + i
365 {
366 *c = self.reduce(ctx, *c);
367 }
368 }
369 }
370 if i == 0 {
371 max_bits = c.max_bits + d.max_bits;
372 } else {
373 max_bits = max_bits.max(c.max_bits + d.max_bits) + 1
374 }
375 }
376 let a_raw = a[range]
377 .iter()
378 .map(|a| QuantumCell::Existing(a.value))
379 .collect_vec();
380 let b_raw = b[other_range]
381 .iter()
382 .rev()
383 .map(|b| QuantumCell::Existing(b.value))
384 .collect_vec();
385 let prod = self.gate().inner_product(ctx, a_raw, b_raw);
386 BabyBearWire {
387 value: prod,
388 max_bits,
389 }
390 }
391
392 pub fn select(
393 &self,
394 ctx: &mut Context<Fr>,
395 cond: SafeBool<Fr>,
396 a: BabyBearWire,
397 b: BabyBearWire,
398 ) -> BabyBearWire {
399 let value = self.gate().select(ctx, a.value, b.value, *cond.as_ref());
400 let max_bits = a.max_bits.max(b.max_bits);
401 BabyBearWire { value, max_bits }
402 }
403
404 pub fn assert_zero(&self, ctx: &mut Context<Fr>, a: BabyBearWire) {
405 guarded_debug_assert_eq!(a.to_baby_bear(), BabyBear::ZERO);
406 assert!(a.max_bits <= Fr::CAPACITY as usize - RESERVED_HIGH_BITS);
407 let a_num_bits = a.max_bits;
408 let b: BigUint = BabyBear::ORDER_U32.into();
409 let a_val = fe_to_bigint(a.value.value());
410 assert!(a_val.bits() <= a_num_bits as u64);
411 let (div, _) = a_val.div_mod_floor(&b.clone().into());
414 let div = bigint_to_fe(&div);
415 ctx.assign_region(
416 [
417 QuantumCell::Constant(Fr::ZERO),
418 QuantumCell::Constant(biguint_to_fe(&b)),
419 QuantumCell::Witness(div),
420 a.value.into(),
421 ],
422 [0],
423 );
424 let div = ctx.get(-2);
425 let bound = (BigUint::from(1u32) << (a_num_bits as u32)) / &b;
427 let shifted_div =
428 self.range
429 .gate()
430 .add(ctx, div, QuantumCell::Constant(biguint_to_fe(&bound)));
431 guarded_debug_assert!(*shifted_div.value() < biguint_to_fe(&(&bound * 2u32 + 1u32)));
432 self.range
433 .range_check(ctx, shifted_div, (bound * 2u32 + 1u32).bits() as usize);
434 }
435
436 pub fn assert_equal(&self, ctx: &mut Context<Fr>, a: BabyBearWire, b: BabyBearWire) {
437 guarded_debug_assert_eq!(a.to_baby_bear(), b.to_baby_bear());
438 let diff = self.sub(ctx, a, b);
439 self.assert_zero(ctx, diff);
440 }
441
442 pub fn zero(&self, ctx: &mut Context<Fr>) -> BabyBearWire {
443 self.load_constant(ctx, BabyBear::ZERO)
444 }
445
446 pub fn one(&self, ctx: &mut Context<Fr>) -> BabyBearWire {
447 self.load_constant(ctx, BabyBear::ONE)
448 }
449
450 pub fn mul_const(&self, ctx: &mut Context<Fr>, a: BabyBearWire, c: BabyBear) -> BabyBearWire {
451 let c_wire = self.load_constant(ctx, c);
452 self.mul(ctx, a, c_wire)
453 }
454
455 pub fn square(&self, ctx: &mut Context<Fr>, a: BabyBearWire) -> BabyBearWire {
456 self.mul(ctx, a, a)
457 }
458
459 pub fn pow_power_of_two(
460 &self,
461 ctx: &mut Context<Fr>,
462 a: BabyBearWire,
463 n: usize,
464 ) -> BabyBearWire {
465 let mut result = a;
466 for _ in 0..n {
467 result = self.square(ctx, result);
468 }
469 result
470 }
471}
472
473fn signed_div_mod<F>(
492 range: &RangeChip<F>,
493 ctx: &mut Context<F>,
494 a: impl Into<QuantumCell<F>>,
495 a_num_bits: usize,
496) -> (AssignedValue<F>, AssignedValue<F>)
497where
498 F: BigPrimeField,
499{
500 assert!(a_num_bits <= F::CAPACITY as usize - RESERVED_HIGH_BITS);
501 let a = a.into();
567 let b = BigUint::from(BabyBear::ORDER_U32);
568 let a_val = fe_to_bigint(a.value());
569 assert!(a_val.bits() <= a_num_bits as u64);
570 let (div, rem) = a_val.div_mod_floor(&b.clone().into());
571 let [div, rem] = [div, rem].map(|v| bigint_to_fe(&v));
572 ctx.assign_region(
573 [
574 QuantumCell::Witness(rem),
575 QuantumCell::Constant(biguint_to_fe(&b)),
576 QuantumCell::Witness(div),
577 a,
578 ],
579 [0],
580 );
581 let rem = ctx.get(-4);
582 let div = ctx.get(-2);
583 let bound = ((BigUint::from(1u32) << a_num_bits) - 1u32).div_ceil(&b);
587 assert!((&bound * 4u32 + 2u32) * &b <= modulus::<F>());
588 let shifted_div = range
589 .gate()
590 .add(ctx, div, QuantumCell::Constant(biguint_to_fe(&bound)));
591 guarded_debug_assert!(*shifted_div.value() < biguint_to_fe(&(&bound * 2u32 + 1u32)));
592 range.range_check(ctx, shifted_div, (bound * 2u32 + 1u32).bits() as usize);
593 guarded_debug_assert!(*rem.value() < biguint_to_fe(&b));
594 range.check_big_less_than_safe(ctx, rem, b);
595 (div, rem)
596}