use p3_air::AirBuilder;
use p3_field::{AbstractField, Field};
use sp1_derive::AlignedBorrow;
use sp1_stark::air::SP1AirBuilder;
#[derive(AlignedBorrow, Default, Debug, Clone, Copy)]
#[repr(C)]
pub struct IsZeroOperation<T> {
pub inverse: T,
pub result: T,
}
impl<F: Field> IsZeroOperation<F> {
pub fn populate(&mut self, a: F) -> F {
let (inverse, result) =
if a.is_zero() { (F::zero(), F::one()) } else { (a.inverse(), F::zero()) };
self.inverse = inverse;
self.result = result;
let prod = inverse * a;
debug_assert!(prod == F::one() || prod.is_zero());
result
}
}
impl<F: Field> IsZeroOperation<F> {
pub fn eval<AB: SP1AirBuilder>(
builder: &mut AB,
a: AB::Expr,
cols: IsZeroOperation<AB::Var>,
is_real: AB::Expr,
) {
builder.assert_bool(is_real.clone());
builder.when(is_real.clone()).assert_bool(cols.result);
let one = AB::Expr::one();
let inverse = cols.inverse;
let is_zero = one.clone() - inverse * a.clone();
builder.when(is_real.clone()).assert_eq(is_zero, cols.result);
builder.when(is_real.clone()).assert_bool(cols.result);
builder.when(is_real.clone()).when(cols.result).assert_zero(a.clone());
}
}