openvm_ecc_circuit/extension/
weierstrass.rs

1use std::sync::Arc;
2
3use hex_literal::hex;
4use lazy_static::lazy_static;
5use num_bigint::BigUint;
6use num_traits::{FromPrimitive, Zero};
7use once_cell::sync::Lazy;
8use openvm_circuit::{
9    arch::{
10        AirInventory, AirInventoryError, ChipInventory, ChipInventoryError, ExecutionBridge,
11        ExecutorInventoryBuilder, ExecutorInventoryError, RowMajorMatrixArena, VmCircuitExtension,
12        VmExecutionExtension, VmProverExtension, DEFAULT_BLOCK_SIZE,
13    },
14    system::{memory::SharedMemoryHelper, SystemPort},
15};
16use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor};
17use openvm_circuit_primitives::{
18    bitwise_op_lookup::{
19        BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip,
20        SharedBitwiseOperationLookupChip,
21    },
22    var_range::VariableRangeCheckerBus,
23};
24use openvm_cpu_backend::{CpuBackend, CpuDevice};
25use openvm_ecc_transpiler::Rv32WeierstrassOpcode;
26use openvm_instructions::{LocalOpcode, VmOpcode};
27use openvm_mod_circuit_builder::ExprBuilderConfig;
28use openvm_stark_backend::{p3_field::PrimeField32, StarkEngine, StarkProtocolConfig, Val};
29use serde::{Deserialize, Serialize};
30use serde_with::{serde_as, DisplayFromStr};
31use strum::EnumCount;
32
33use crate::{
34    get_ec_addne_air, get_ec_addne_chip, get_ec_addne_executor, get_ec_double_air,
35    get_ec_double_chip, get_ec_double_executor, EcAddNeExecutor, EcDoubleExecutor, EccCpuProverExt,
36    WeierstrassAir, ECC_BLOCKS_32, ECC_BLOCKS_48, NUM_LIMBS_32, NUM_LIMBS_48,
37};
38
39#[serde_as]
40#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)]
41pub struct CurveConfig {
42    /// The name of the curve struct as defined by moduli_declare.
43    pub struct_name: String,
44    /// The coordinate modulus of the curve.
45    #[serde_as(as = "DisplayFromStr")]
46    pub modulus: BigUint,
47    /// The scalar field modulus of the curve.
48    #[serde_as(as = "DisplayFromStr")]
49    pub scalar: BigUint,
50    /// The coefficient a of y^2 = x^3 + ax + b.
51    #[serde_as(as = "DisplayFromStr")]
52    pub a: BigUint,
53    /// The coefficient b of y^2 = x^3 + ax + b.
54    #[serde_as(as = "DisplayFromStr")]
55    pub b: BigUint,
56}
57
58pub static SECP256K1_CONFIG: Lazy<CurveConfig> = Lazy::new(|| CurveConfig {
59    struct_name: SECP256K1_ECC_STRUCT_NAME.to_string(),
60    modulus: SECP256K1_MODULUS.clone(),
61    scalar: SECP256K1_ORDER.clone(),
62    a: BigUint::zero(),
63    b: BigUint::from_u8(7u8).unwrap(),
64});
65
66pub static P256_CONFIG: Lazy<CurveConfig> = Lazy::new(|| CurveConfig {
67    struct_name: P256_ECC_STRUCT_NAME.to_string(),
68    modulus: P256_MODULUS.clone(),
69    scalar: P256_ORDER.clone(),
70    a: BigUint::from_bytes_le(&P256_A),
71    b: BigUint::from_bytes_le(&P256_B),
72});
73
74#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)]
75pub struct WeierstrassExtension {
76    pub supported_curves: Vec<CurveConfig>,
77}
78
79impl WeierstrassExtension {
80    pub fn generate_sw_init(&self) -> String {
81        let supported_curves = self
82            .supported_curves
83            .iter()
84            .map(|curve_config| format!("\"{}\"", curve_config.struct_name))
85            .collect::<Vec<String>>()
86            .join(", ");
87
88        format!("openvm_ecc_guest::sw_macros::sw_init! {{ {supported_curves} }}")
89    }
90}
91
92#[derive(Clone, AnyEnum, Executor, MeteredExecutor, PreflightExecutor)]
93#[cfg_attr(
94    feature = "aot",
95    derive(
96        openvm_circuit_derive::AotExecutor,
97        openvm_circuit_derive::AotMeteredExecutor
98    )
99)]
100pub enum WeierstrassExtensionExecutor {
101    // 32 limbs prime
102    EcAddNeRv32_32(EcAddNeExecutor<ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>),
103    EcDoubleRv32_32(EcDoubleExecutor<ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>),
104    // 48 limbs prime
105    EcAddNeRv32_48(EcAddNeExecutor<ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>),
106    EcDoubleRv32_48(EcDoubleExecutor<ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>),
107}
108
109impl<F: PrimeField32> VmExecutionExtension<F> for WeierstrassExtension {
110    type Executor = WeierstrassExtensionExecutor;
111
112    fn extend_execution(
113        &self,
114        inventory: &mut ExecutorInventoryBuilder<F, WeierstrassExtensionExecutor>,
115    ) -> Result<(), ExecutorInventoryError> {
116        let pointer_max_bits = inventory.pointer_max_bits();
117        // TODO: somehow get the range checker bus from `ExecutorInventory`
118        let dummy_range_checker_bus = VariableRangeCheckerBus::new(u16::MAX, 16);
119        for (i, curve) in self.supported_curves.iter().enumerate() {
120            let start_offset =
121                Rv32WeierstrassOpcode::CLASS_OFFSET + i * Rv32WeierstrassOpcode::COUNT;
122            let bytes = curve.modulus.bits().div_ceil(8) as usize;
123
124            if bytes <= NUM_LIMBS_32 {
125                let config = ExprBuilderConfig {
126                    modulus: curve.modulus.clone(),
127                    num_limbs: NUM_LIMBS_32,
128                    limb_bits: 8,
129                };
130                let addne = get_ec_addne_executor(
131                    config.clone(),
132                    dummy_range_checker_bus,
133                    pointer_max_bits,
134                    start_offset,
135                );
136
137                inventory.add_executor(
138                    WeierstrassExtensionExecutor::EcAddNeRv32_32(addne),
139                    ((Rv32WeierstrassOpcode::EC_ADD_NE as usize)
140                        ..=(Rv32WeierstrassOpcode::SETUP_EC_ADD_NE as usize))
141                        .map(|x| VmOpcode::from_usize(x + start_offset)),
142                )?;
143
144                let double = get_ec_double_executor(
145                    config,
146                    dummy_range_checker_bus,
147                    pointer_max_bits,
148                    start_offset,
149                    curve.a.clone(),
150                );
151
152                inventory.add_executor(
153                    WeierstrassExtensionExecutor::EcDoubleRv32_32(double),
154                    ((Rv32WeierstrassOpcode::EC_DOUBLE as usize)
155                        ..=(Rv32WeierstrassOpcode::SETUP_EC_DOUBLE as usize))
156                        .map(|x| VmOpcode::from_usize(x + start_offset)),
157                )?;
158            } else if bytes <= NUM_LIMBS_48 {
159                let config = ExprBuilderConfig {
160                    modulus: curve.modulus.clone(),
161                    num_limbs: NUM_LIMBS_48,
162                    limb_bits: 8,
163                };
164                let addne = get_ec_addne_executor(
165                    config.clone(),
166                    dummy_range_checker_bus,
167                    pointer_max_bits,
168                    start_offset,
169                );
170
171                inventory.add_executor(
172                    WeierstrassExtensionExecutor::EcAddNeRv32_48(addne),
173                    ((Rv32WeierstrassOpcode::EC_ADD_NE as usize)
174                        ..=(Rv32WeierstrassOpcode::SETUP_EC_ADD_NE as usize))
175                        .map(|x| VmOpcode::from_usize(x + start_offset)),
176                )?;
177
178                let double = get_ec_double_executor(
179                    config,
180                    dummy_range_checker_bus,
181                    pointer_max_bits,
182                    start_offset,
183                    curve.a.clone(),
184                );
185
186                inventory.add_executor(
187                    WeierstrassExtensionExecutor::EcDoubleRv32_48(double),
188                    ((Rv32WeierstrassOpcode::EC_DOUBLE as usize)
189                        ..=(Rv32WeierstrassOpcode::SETUP_EC_DOUBLE as usize))
190                        .map(|x| VmOpcode::from_usize(x + start_offset)),
191                )?;
192            } else {
193                panic!("Modulus too large");
194            }
195        }
196
197        Ok(())
198    }
199}
200
201impl<SC: StarkProtocolConfig> VmCircuitExtension<SC> for WeierstrassExtension {
202    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
203        let SystemPort {
204            execution_bus,
205            program_bus,
206            memory_bridge,
207        } = inventory.system().port();
208
209        let exec_bridge = ExecutionBridge::new(execution_bus, program_bus);
210        let range_checker_bus = inventory.range_checker().bus;
211        let pointer_max_bits = inventory.pointer_max_bits();
212
213        let bitwise_lu = {
214            // A trick to get around Rust's borrow rules
215            let existing_air = inventory.find_air::<BitwiseOperationLookupAir<8>>().next();
216            if let Some(air) = existing_air {
217                air.bus
218            } else {
219                let bus = BitwiseOperationLookupBus::new(inventory.new_bus_idx());
220                let air = BitwiseOperationLookupAir::<8>::new(bus);
221                inventory.add_air(air);
222                air.bus
223            }
224        };
225        for (i, curve) in self.supported_curves.iter().enumerate() {
226            let start_offset =
227                Rv32WeierstrassOpcode::CLASS_OFFSET + i * Rv32WeierstrassOpcode::COUNT;
228            let bytes = curve.modulus.bits().div_ceil(8) as usize;
229
230            if bytes <= NUM_LIMBS_32 {
231                let config = ExprBuilderConfig {
232                    modulus: curve.modulus.clone(),
233                    num_limbs: NUM_LIMBS_32,
234                    limb_bits: 8,
235                };
236
237                let addne = get_ec_addne_air::<ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
238                    exec_bridge,
239                    memory_bridge,
240                    config.clone(),
241                    range_checker_bus,
242                    bitwise_lu,
243                    pointer_max_bits,
244                    start_offset,
245                );
246                inventory.add_air(addne);
247
248                let double = get_ec_double_air::<ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
249                    exec_bridge,
250                    memory_bridge,
251                    config,
252                    range_checker_bus,
253                    bitwise_lu,
254                    pointer_max_bits,
255                    start_offset,
256                    curve.a.clone(),
257                );
258                inventory.add_air(double);
259            } else if bytes <= NUM_LIMBS_48 {
260                let config = ExprBuilderConfig {
261                    modulus: curve.modulus.clone(),
262                    num_limbs: NUM_LIMBS_48,
263                    limb_bits: 8,
264                };
265
266                let addne = get_ec_addne_air::<ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
267                    exec_bridge,
268                    memory_bridge,
269                    config.clone(),
270                    range_checker_bus,
271                    bitwise_lu,
272                    pointer_max_bits,
273                    start_offset,
274                );
275                inventory.add_air(addne);
276
277                let double = get_ec_double_air::<ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
278                    exec_bridge,
279                    memory_bridge,
280                    config,
281                    range_checker_bus,
282                    bitwise_lu,
283                    pointer_max_bits,
284                    start_offset,
285                    curve.a.clone(),
286                );
287                inventory.add_air(double);
288            } else {
289                panic!("Modulus too large");
290            }
291        }
292
293        Ok(())
294    }
295}
296
297// This implementation is specific to CpuBackend because the lookup chips (VariableRangeChecker,
298// BitwiseOperationLookupChip) are specific to CpuBackend.
299impl<SC, E, RA> VmProverExtension<E, RA, WeierstrassExtension> for EccCpuProverExt
300where
301    SC: StarkProtocolConfig,
302    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
303    RA: RowMajorMatrixArena<Val<SC>>,
304    Val<SC>: PrimeField32,
305    SC::EF: Ord,
306{
307    fn extend_prover(
308        &self,
309        extension: &WeierstrassExtension,
310        inventory: &mut ChipInventory<SC, RA, CpuBackend<SC>>,
311    ) -> Result<(), ChipInventoryError> {
312        let range_checker = inventory.range_checker()?.clone();
313        let timestamp_max_bits = inventory.timestamp_max_bits();
314        let pointer_max_bits = inventory.airs().pointer_max_bits();
315        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
316        let bitwise_lu = {
317            let existing_chip = inventory
318                .find_chip::<SharedBitwiseOperationLookupChip<8>>()
319                .next();
320            if let Some(chip) = existing_chip {
321                chip.clone()
322            } else {
323                let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
324                let chip = Arc::new(BitwiseOperationLookupChip::new(air.bus));
325                inventory.add_periphery_chip(chip.clone());
326                chip
327            }
328        };
329        for curve in extension.supported_curves.iter() {
330            let bytes = curve.modulus.bits().div_ceil(8) as usize;
331
332            if bytes <= NUM_LIMBS_32 {
333                let config = ExprBuilderConfig {
334                    modulus: curve.modulus.clone(),
335                    num_limbs: NUM_LIMBS_32,
336                    limb_bits: 8,
337                };
338
339                inventory.next_air::<WeierstrassAir<2, ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>>()?;
340                let addne = get_ec_addne_chip::<Val<SC>, ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
341                    config.clone(),
342                    mem_helper.clone(),
343                    range_checker.clone(),
344                    bitwise_lu.clone(),
345                    pointer_max_bits,
346                );
347                inventory.add_executor_chip(addne);
348
349                inventory.next_air::<WeierstrassAir<1, ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>>()?;
350                let double = get_ec_double_chip::<Val<SC>, ECC_BLOCKS_32, DEFAULT_BLOCK_SIZE>(
351                    config,
352                    mem_helper.clone(),
353                    range_checker.clone(),
354                    bitwise_lu.clone(),
355                    pointer_max_bits,
356                    curve.a.clone(),
357                );
358                inventory.add_executor_chip(double);
359            } else if bytes <= NUM_LIMBS_48 {
360                let config = ExprBuilderConfig {
361                    modulus: curve.modulus.clone(),
362                    num_limbs: NUM_LIMBS_48,
363                    limb_bits: 8,
364                };
365
366                inventory.next_air::<WeierstrassAir<2, ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>>()?;
367                let addne = get_ec_addne_chip::<Val<SC>, ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
368                    config.clone(),
369                    mem_helper.clone(),
370                    range_checker.clone(),
371                    bitwise_lu.clone(),
372                    pointer_max_bits,
373                );
374                inventory.add_executor_chip(addne);
375
376                inventory.next_air::<WeierstrassAir<1, ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>>()?;
377                let double = get_ec_double_chip::<Val<SC>, ECC_BLOCKS_48, DEFAULT_BLOCK_SIZE>(
378                    config,
379                    mem_helper.clone(),
380                    range_checker.clone(),
381                    bitwise_lu.clone(),
382                    pointer_max_bits,
383                    curve.a.clone(),
384                );
385                inventory.add_executor_chip(double);
386            } else {
387                panic!("Modulus too large");
388            }
389        }
390
391        Ok(())
392    }
393}
394
395// Convenience constants for constructors
396lazy_static! {
397    // The constants are taken from: https://en.bitcoin.it/wiki/Secp256k1
398    pub static ref SECP256K1_MODULUS: BigUint = BigUint::from_bytes_be(&hex!(
399        "FFFFFFFF FFFFFFFF FFFFFFFF FFFFFFFF FFFFFFFF FFFFFFFF FFFFFFFE FFFFFC2F"
400    ));
401    pub static ref SECP256K1_ORDER: BigUint = BigUint::from_bytes_be(&hex!(
402        "FFFFFFFF FFFFFFFF FFFFFFFF FFFFFFFE BAAEDCE6 AF48A03B BFD25E8C D0364141"
403    ));
404}
405
406lazy_static! {
407    // The constants are taken from: https://neuromancer.sk/std/secg/secp256r1
408    pub static ref P256_MODULUS: BigUint = BigUint::from_bytes_be(&hex!(
409        "ffffffff00000001000000000000000000000000ffffffffffffffffffffffff"
410    ));
411    pub static ref P256_ORDER: BigUint = BigUint::from_bytes_be(&hex!(
412        "ffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551"
413    ));
414}
415// little-endian
416const P256_A: [u8; 32] = hex!("fcffffffffffffffffffffff00000000000000000000000001000000ffffffff");
417// little-endian
418const P256_B: [u8; 32] = hex!("4b60d2273e3cce3bf6b053ccb0061d65bc86987655bdebb3e7933aaad835c65a");
419
420pub const SECP256K1_ECC_STRUCT_NAME: &str = "Secp256k1Point";
421pub const P256_ECC_STRUCT_NAME: &str = "P256Point";
openvm_ecc_circuit/extension/weierstrass.rs

openvm_ecc_circuit/extension/
weierstrass.rs
