openvm-ecc-circuit 2.0.1

use std::result::Result;

use openvm_algebra_circuit::{Rv32ModularConfig, Rv32ModularConfigExecutor, Rv32ModularCpuBuilder};
use openvm_circuit::{
    arch::{
        AirInventory, ChipInventoryError, InitFileGenerator, MatrixRecordArena, SystemConfig,
        VmBuilder, VmChipComplex, VmField, VmProverExtension,
    },
    system::SystemChipInventory,
};
use openvm_circuit_derive::VmConfig;
use openvm_cpu_backend::{CpuBackend, CpuDevice};
use openvm_stark_backend::{StarkEngine, StarkProtocolConfig, Val};
use serde::{Deserialize, Serialize};

mod weierstrass;
pub use weierstrass::*;

cfg_if::cfg_if! {
    if #[cfg(feature = "cuda")] {
        mod hybrid;
        pub use hybrid::*;
        pub use {
            EccHybridProverExt as EccProverExt,
            Rv32WeierstrassHybridBuilder as Rv32WeierstrassBuilder,
        };
    } else {
        pub use self::{
            EccCpuProverExt as EccProverExt,
            Rv32WeierstrassCpuBuilder as Rv32WeierstrassBuilder,
        };
    }
}
pub struct EccCpuProverExt;

#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)]
pub struct Rv32WeierstrassConfig {
    #[config(generics = true)]
    pub modular: Rv32ModularConfig,
    #[extension]
    pub weierstrass: WeierstrassExtension,
}

impl Rv32WeierstrassConfig {
    pub fn new(curves: Vec<CurveConfig>) -> Self {
        let primes: Vec<_> = curves
            .iter()
            .flat_map(|c| [c.modulus.clone(), c.scalar.clone()])
            .collect();
        Self {
            modular: Rv32ModularConfig::new(primes),
            weierstrass: WeierstrassExtension::new(curves),
        }
    }
}

impl InitFileGenerator for Rv32WeierstrassConfig {
    fn generate_init_file_contents(&self) -> Option<String> {
        Some(format!(
            "// This file is automatically generated by cargo openvm. Do not rename or edit.\n{}\n{}\n",
            self.modular.modular.generate_moduli_init(),
            self.weierstrass.generate_sw_init()
        ))
    }
}

#[derive(Clone)]
pub struct Rv32WeierstrassCpuBuilder;

impl<SC, E> VmBuilder<E> for Rv32WeierstrassCpuBuilder
where
    SC: StarkProtocolConfig,
    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
    Val<SC>: VmField,
    SC::EF: Ord,
{
    type VmConfig = Rv32WeierstrassConfig;
    type SystemChipInventory = SystemChipInventory<SC>;
    type RecordArena = MatrixRecordArena<Val<SC>>;

    fn create_chip_complex(
        &self,
        config: &Self::VmConfig,
        circuit: AirInventory<SC>,
        device_ctx: &openvm_stark_backend::EngineDeviceCtx<E>,
    ) -> Result<
        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
        ChipInventoryError,
    > {
        let mut chip_complex = VmBuilder::<E>::create_chip_complex(
            &Rv32ModularCpuBuilder,
            &config.modular,
            circuit,
            device_ctx,
        )?;
        let inventory = &mut chip_complex.inventory;
        VmProverExtension::<E, _, _>::extend_prover(
            &EccCpuProverExt,
            &config.weierstrass,
            inventory,
        )?;
        Ok(chip_complex)
    }
}