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