openvm-circuit-primitives 2.0.1

use std::borrow::{Borrow, BorrowMut};

use openvm_circuit_primitives_derive::AlignedBorrow;
use openvm_stark_backend::{
    any_air_arc_vec,
    p3_air::{Air, AirBuilder, BaseAir},
    p3_field::{Field, PrimeCharacteristicRing},
    p3_matrix::{dense::RowMajorMatrix, Matrix},
    p3_maybe_rayon::prelude::*,
    prover::AirProvingContext,
    utils::disable_debug_builder,
    BaseAirWithPublicValues, PartitionedBaseAir, StarkEngine, StarkTestError,
};
#[cfg(not(feature = "cuda"))]
use openvm_stark_sdk::config::baby_bear_poseidon2::F;
use test_case::test_case;
#[cfg(feature = "cuda")]
use {
    crate::cuda_abi::is_zero,
    openvm_cuda_backend::{
        base::DeviceMatrix, data_transporter::assert_eq_host_and_device_matrix, prelude::F,
    },
    openvm_cuda_common::copy::MemCopyH2D as _,
    openvm_stark_backend::p3_field::PrimeField32,
    openvm_stark_sdk::utils::create_seeded_rng,
    rand::Rng,
    std::sync::Arc,
};

use super::{IsZeroIo, IsZeroSubAir};
use crate::{
    utils::test_engine_small, ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
    TraceSubRowGenerator,
};

#[repr(C)]
#[derive(AlignedBorrow, StructReflection)]
pub struct IsZeroCols<T> {
    pub x: T,
    pub out: T,
    pub inv: T,
}

#[derive(Copy, Clone, ColumnsAir)]
#[columns_via(IsZeroCols<u8>)]
pub struct IsZeroTestAir(IsZeroSubAir);

impl<F: Field> BaseAirWithPublicValues<F> for IsZeroTestAir {}
impl<F: Field> PartitionedBaseAir<F> for IsZeroTestAir {}
impl<F: Field> BaseAir<F> for IsZeroTestAir {
    fn width(&self) -> usize {
        IsZeroCols::<F>::width()
    }
}
impl<AB> Air<AB> for IsZeroTestAir
where
    AB: AirBuilder<Var: Copy>,
    AB::F: Field,
{
    fn eval(&self, builder: &mut AB) {
        let main = builder.main();

        let local = main.row_slice(0).expect("window should have two elements");
        let local: &IsZeroCols<_> = (*local).borrow();
        let io = IsZeroIo::new(local.x.into(), local.out.into(), AB::Expr::ONE);

        self.0.eval(builder, (io, local.inv));
    }
}

pub struct IsZeroChip<F> {
    air: IsZeroTestAir,
    x: Vec<F>,
}

impl<F: Field> IsZeroChip<F> {
    pub fn new(x: Vec<F>) -> Self {
        Self {
            air: IsZeroTestAir(IsZeroSubAir),
            x,
        }
    }

    pub fn generate_trace(self) -> RowMajorMatrix<F> {
        let air = IsZeroSubAir;
        assert!(self.x.len().is_power_of_two());
        let width = IsZeroCols::<F>::width();
        let mut rows = F::zero_vec(width * self.x.len());
        rows.par_chunks_mut(width).zip(self.x).for_each(|(row, x)| {
            let row: &mut IsZeroCols<F> = row.borrow_mut();
            row.x = x;
            air.generate_subrow(x, (&mut row.inv, &mut row.out));
        });

        RowMajorMatrix::new(rows, width)
    }
}

#[test_case(97 ; "97 => 0")]
#[test_case(0 ; "0 => 1")]
fn test_single_is_zero(x: u32) {
    let chip = IsZeroChip::new(vec![F::from_u32(x)]);
    let air = chip.air;
    let trace = chip.generate_trace();

    assert_eq!(
        trace.get(0, 1).expect("matrix index out of bounds"),
        PrimeCharacteristicRing::from_bool(x == 0)
    );

    let traces = [trace]
        .into_iter()
        .map(AirProvingContext::simple_no_pis)
        .collect::<Vec<_>>();
    test_engine_small()
        .run_test(any_air_arc_vec![air], traces)
        .expect("Verification failed");
}

#[test_case([0, 1, 2, 7], [1, 0, 0, 0] ; "0, 1, 2, 7 => 1, 0, 0, 0")]
#[test_case([97, 23, 179, 0], [0, 0, 0, 1] ; "97, 23, 179, 0 => 0, 0, 0, 1")]
fn test_vec_is_zero(x_vec: [u32; 4], expected: [u32; 4]) {
    let x_vec = x_vec
        .into_iter()
        .map(PrimeCharacteristicRing::from_u32)
        .collect();
    let chip = IsZeroChip::new(x_vec);
    let air = chip.air;
    let trace = chip.generate_trace();

    for (i, value) in expected.iter().enumerate() {
        assert_eq!(
            trace.values[3 * i + 1],
            PrimeCharacteristicRing::from_u32(*value)
        );
    }

    let traces = [trace]
        .into_iter()
        .map(AirProvingContext::simple_no_pis)
        .collect::<Vec<_>>();
    test_engine_small()
        .run_test(any_air_arc_vec![air], traces)
        .expect("Verification failed");
}

#[test_case(97 ; "97 => 0")]
#[test_case(0 ; "0 => 1")]
fn test_single_is_zero_fail(x: u32) {
    let x = PrimeCharacteristicRing::from_u32(x);
    let chip = IsZeroChip::new(vec![x]);
    let air = chip.air;
    let mut trace = chip.generate_trace();
    trace.values[1] = F::ONE - trace.values[1];

    disable_debug_builder();
    let traces = [trace]
        .into_iter()
        .map(AirProvingContext::simple_no_pis)
        .collect::<Vec<_>>();
    let result = test_engine_small().run_test(any_air_arc_vec![air], traces);
    assert!(matches!(result, Err(StarkTestError::Verifier(_))));
}

#[test_case([1, 2, 7, 0], [0, 0, 0, 1] ; "1, 2, 7, 0 => 0, 0, 0, 1")]
#[test_case([97, 0, 179, 0], [0, 1, 0, 1] ; "97, 0, 179, 0 => 0, 1, 0, 1")]
fn test_vec_is_zero_fail(x_vec: [u32; 4], _expected: [u32; 4]) {
    let x_vec: Vec<F> = x_vec.into_iter().map(F::from_u32).collect();
    let chip = IsZeroChip::new(x_vec);
    let air = chip.air;
    let mut trace = chip.generate_trace();

    disable_debug_builder();
    // Corrupt the first row's output to trigger a constraint failure
    trace.row_mut(0)[1] = F::ONE - trace.row_mut(0)[1];
    let traces = [trace]
        .into_iter()
        .map(AirProvingContext::simple_no_pis)
        .collect::<Vec<_>>();
    let result = test_engine_small().run_test(any_air_arc_vec![air], traces);
    assert!(matches!(result, Err(StarkTestError::Verifier(_))));
}

#[cfg(feature = "cuda")]
#[test]
fn test_cuda_is_zero_against_cpu_full() {
    let device_ctx = crate::utils::test_device_ctx();
    let mut rng = create_seeded_rng();
    for log_height in 1..=16 {
        let n = 1 << log_height;
        let vec_x: Vec<F> = (0..n)
            .map(|_| {
                if rng.random_bool(0.5) {
                    0 // 50% chance to be zero
                } else {
                    rng.random_range(0..F::ORDER_U32) // 50% chance to be random
                }
            })
            .map(F::from_u32)
            .collect();

        let input_buffer = vec_x.as_slice().to_device_on(&device_ctx).unwrap();
        let output = DeviceMatrix::<F>::with_capacity_on(n, 2, &device_ctx);
        unsafe {
            is_zero::dummy_tracegen(output.buffer(), &input_buffer, device_ctx.stream.as_raw())
                .unwrap();
        };

        let cpu_matrix = Arc::new(RowMajorMatrix::<F>::new(
            vec_x
                .iter()
                .flat_map(|x| {
                    let cur_x = *x;
                    let mut cur_inv = F::ZERO;
                    let mut cur_out = F::ONE;
                    IsZeroSubAir.generate_subrow(cur_x, (&mut cur_inv, &mut cur_out));
                    [cur_inv, cur_out]
                })
                .collect::<Vec<_>>(),
            2,
        ));

        assert_eq_host_and_device_matrix(cpu_matrix, &output, &device_ctx);
    }
}