openvm_deferral_circuit/canonicity/
air.rs1use itertools::{izip, Itertools};
2use openvm_circuit_primitives::{utils::not, SubAir};
3use openvm_stark_backend::{
4 interaction::InteractionBuilder,
5 p3_air::AirBuilder,
6 p3_field::{PrimeCharacteristicRing, PrimeField32},
7};
8
9use super::{CanonicityAuxCols, CanonicityIo};
10
11pub struct CanonicitySubAir;
15
16impl<AB: InteractionBuilder> SubAir<AB> for CanonicitySubAir
17where
18 AB::F: PrimeField32,
19{
20 type AirContext<'a>
21 = (
22 CanonicityIo<AB::Expr>,
23 &'a CanonicityAuxCols<AB::Var>,
24 &'a mut AB::Expr,
25 )
26 where
27 AB::Expr: 'a,
28 AB::Var: 'a,
29 AB: 'a;
30
31 fn eval<'a>(
32 &'a self,
33 builder: &'a mut AB,
34 (io, aux, to_range_check): (
35 CanonicityIo<AB::Expr>,
36 &'a CanonicityAuxCols<AB::Var>,
37 &'a mut AB::Expr,
38 ),
39 ) where
40 AB::Var: 'a,
41 AB::Expr: 'a,
42 {
43 let order_be = AB::F::ORDER_U32
44 .to_le_bytes()
45 .into_iter()
46 .rev()
47 .map(AB::Expr::from_u8);
48
49 let mut prefix_sum = AB::Expr::ZERO;
50
51 for (x, y, &marker) in izip!(io.x, order_be, aux.diff_marker.iter()) {
52 let diff = y - x;
53 prefix_sum += marker.into();
54 builder.assert_bool(marker);
55 builder.when(marker).assert_one(io.count.clone());
56 builder
57 .when(io.count.clone())
58 .assert_zero(not::<AB::Expr>(prefix_sum.clone()) * diff.clone());
59 builder.when(marker).assert_eq(aux.diff_val, diff);
60 }
61
62 builder.assert_bool(prefix_sum.clone());
63 builder.when(io.count.clone()).assert_one(prefix_sum);
64
65 *to_range_check = AB::Expr::from(aux.diff_val) - AB::Expr::ONE;
66 }
67}
68
69impl CanonicitySubAir {
70 pub fn assert_canonicity<AB: InteractionBuilder>(
71 &self,
72 builder: &mut AB,
73 x: &[AB::Var],
74 aux: &CanonicityAuxCols<AB::Var>,
75 enabled: AB::Expr,
76 ) -> AB::Expr
77 where
78 AB::F: PrimeField32,
79 {
80 let io = CanonicityIo {
81 x: x.iter().rev().map(|b| (*b).into()).collect_array().unwrap(),
82 count: enabled,
83 };
84 let mut ret = AB::Expr::ZERO;
85 self.eval(builder, (io, aux, &mut ret));
86 ret
87 }
88}