openvm_circuit_primitives/is_zero/
mod.rs

1use derive_new::new;
2use openvm_circuit_primitives_derive::AlignedBorrow;
3use openvm_stark_backend::{p3_air::AirBuilder, p3_field::Field};
4
5use crate::{SubAir, TraceSubRowGenerator};
6
7#[cfg(test)]
8pub mod tests;
9
10#[repr(C)]
11#[derive(Copy, Clone, Debug, new)]
12pub struct IsZeroIo<T> {
13    pub x: T,
14    /// The boolean output, constrained to equal (x == 0) when `condition != 0`..
15    pub out: T,
16    /// Constraints only hold when `condition != 0`. When `condition == 0`, setting all trace values
17    /// to zero still passes the constraints.
18    pub condition: T,
19}
20
21#[repr(C)]
22#[derive(AlignedBorrow, Copy, Clone, Debug, new)]
23pub struct IsZeroAuxCols<T> {
24    pub inv: T,
25}
26
27/// An Air that constraints that checks if a number equals 0
28#[derive(Copy, Clone)]
29pub struct IsZeroSubAir;
30
31impl<AB: AirBuilder> SubAir<AB> for IsZeroSubAir {
32    /// (io, inv)
33    type AirContext<'a>
34        = (IsZeroIo<AB::Expr>, AB::Var)
35    where
36        AB::Expr: 'a,
37        AB::Var: 'a,
38        AB: 'a;
39
40    fn eval<'a>(&'a self, builder: &'a mut AB, (io, inv): (IsZeroIo<AB::Expr>, AB::Var))
41    where
42        AB::Var: 'a,
43        AB::Expr: 'a,
44    {
45        // We always assert this, even when `condition == 0`, because x = 0, out = 0 will pass.
46        builder.assert_zero(io.x.clone() * io.out.clone());
47        builder.when(io.condition).assert_one(io.out + io.x * inv);
48    }
49}
50
51impl<F: Field> TraceSubRowGenerator<F> for IsZeroSubAir {
52    /// `x`
53    type TraceContext<'a> = F;
54    /// `(inv, out)`
55    type ColsMut<'a> = (&'a mut F, &'a mut F);
56
57    fn generate_subrow<'a>(&'a self, x: F, (inv, out): (&'a mut F, &'a mut F)) {
58        *out = F::from_bool(x.is_zero());
59        *inv = x.try_inverse().unwrap_or(F::ZERO);
60    }
61}