openvm_deferral_circuit/canonicity/
air.rs

1use 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
11/// Sub-AIR to constrain that a field byte decomposition is canonical. Note:
12/// - It is assumed that each value has been range checked
13/// - eval returns a value to be range check
14pub 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}