1use derive_more::derive::From;
2use num_bigint::BigUint;
3use openvm_algebra_transpiler::Fp2Opcode;
4use openvm_circuit::{
5 arch::{SystemPort, VmExtension, VmInventory, VmInventoryBuilder, VmInventoryError},
6 system::phantom::PhantomChip,
7};
8use openvm_circuit_derive::{AnyEnum, InstructionExecutor};
9use openvm_circuit_primitives::bitwise_op_lookup::{
10 BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip,
11};
12use openvm_circuit_primitives_derive::{Chip, ChipUsageGetter};
13use openvm_instructions::{LocalOpcode, VmOpcode};
14use openvm_mod_circuit_builder::ExprBuilderConfig;
15use openvm_rv32_adapters::Rv32VecHeapAdapterChip;
16use openvm_stark_backend::p3_field::PrimeField32;
17use serde::{Deserialize, Serialize};
18use serde_with::{serde_as, DisplayFromStr};
19use strum::EnumCount;
20
21use crate::fp2_chip::{Fp2AddSubChip, Fp2MulDivChip};
22
23#[serde_as]
24#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)]
25pub struct Fp2Extension {
26 #[serde_as(as = "Vec<DisplayFromStr>")]
27 pub supported_modulus: Vec<BigUint>,
28}
29
30#[derive(ChipUsageGetter, Chip, InstructionExecutor, AnyEnum, From)]
31pub enum Fp2ExtensionExecutor<F: PrimeField32> {
32 Fp2AddSubRv32_32(Fp2AddSubChip<F, 2, 32>),
34 Fp2MulDivRv32_32(Fp2MulDivChip<F, 2, 32>),
35 Fp2AddSubRv32_48(Fp2AddSubChip<F, 6, 16>),
37 Fp2MulDivRv32_48(Fp2MulDivChip<F, 6, 16>),
38}
39
40#[derive(ChipUsageGetter, Chip, AnyEnum, From)]
41pub enum Fp2ExtensionPeriphery<F: PrimeField32> {
42 BitwiseOperationLookup(SharedBitwiseOperationLookupChip<8>),
43 Phantom(PhantomChip<F>),
45}
46
47impl<F: PrimeField32> VmExtension<F> for Fp2Extension {
48 type Executor = Fp2ExtensionExecutor<F>;
49 type Periphery = Fp2ExtensionPeriphery<F>;
50
51 fn build(
52 &self,
53 builder: &mut VmInventoryBuilder<F>,
54 ) -> Result<VmInventory<Self::Executor, Self::Periphery>, VmInventoryError> {
55 let mut inventory = VmInventory::new();
56 let SystemPort {
57 execution_bus,
58 program_bus,
59 memory_bridge,
60 } = builder.system_port();
61 let bitwise_lu_chip = if let Some(&chip) = builder
62 .find_chip::<SharedBitwiseOperationLookupChip<8>>()
63 .first()
64 {
65 chip.clone()
66 } else {
67 let bitwise_lu_bus = BitwiseOperationLookupBus::new(builder.new_bus_idx());
68 let chip = SharedBitwiseOperationLookupChip::new(bitwise_lu_bus);
69 inventory.add_periphery_chip(chip.clone());
70 chip
71 };
72 let offline_memory = builder.system_base().offline_memory();
73 let range_checker = builder.system_base().range_checker_chip.clone();
74 let address_bits = builder.system_config().memory_config.pointer_max_bits;
75
76 let addsub_opcodes = (Fp2Opcode::ADD as usize)..=(Fp2Opcode::SETUP_ADDSUB as usize);
77 let muldiv_opcodes = (Fp2Opcode::MUL as usize)..=(Fp2Opcode::SETUP_MULDIV as usize);
78
79 for (i, modulus) in self.supported_modulus.iter().enumerate() {
80 let bytes = modulus.bits().div_ceil(8);
82 let start_offset = Fp2Opcode::CLASS_OFFSET + i * Fp2Opcode::COUNT;
83
84 let config32 = ExprBuilderConfig {
85 modulus: modulus.clone(),
86 num_limbs: 32,
87 limb_bits: 8,
88 };
89 let config48 = ExprBuilderConfig {
90 modulus: modulus.clone(),
91 num_limbs: 48,
92 limb_bits: 8,
93 };
94 let adapter_chip_32 = Rv32VecHeapAdapterChip::new(
95 execution_bus,
96 program_bus,
97 memory_bridge,
98 address_bits,
99 bitwise_lu_chip.clone(),
100 );
101 let adapter_chip_48 = Rv32VecHeapAdapterChip::new(
102 execution_bus,
103 program_bus,
104 memory_bridge,
105 address_bits,
106 bitwise_lu_chip.clone(),
107 );
108
109 if bytes <= 32 {
110 let addsub_chip = Fp2AddSubChip::new(
111 adapter_chip_32.clone(),
112 config32.clone(),
113 start_offset,
114 range_checker.clone(),
115 offline_memory.clone(),
116 );
117 inventory.add_executor(
118 Fp2ExtensionExecutor::Fp2AddSubRv32_32(addsub_chip),
119 addsub_opcodes
120 .clone()
121 .map(|x| VmOpcode::from_usize(x + start_offset)),
122 )?;
123 let muldiv_chip = Fp2MulDivChip::new(
124 adapter_chip_32.clone(),
125 config32.clone(),
126 start_offset,
127 range_checker.clone(),
128 offline_memory.clone(),
129 );
130 inventory.add_executor(
131 Fp2ExtensionExecutor::Fp2MulDivRv32_32(muldiv_chip),
132 muldiv_opcodes
133 .clone()
134 .map(|x| VmOpcode::from_usize(x + start_offset)),
135 )?;
136 } else if bytes <= 48 {
137 let addsub_chip = Fp2AddSubChip::new(
138 adapter_chip_48.clone(),
139 config48.clone(),
140 start_offset,
141 range_checker.clone(),
142 offline_memory.clone(),
143 );
144 inventory.add_executor(
145 Fp2ExtensionExecutor::Fp2AddSubRv32_48(addsub_chip),
146 addsub_opcodes
147 .clone()
148 .map(|x| VmOpcode::from_usize(x + start_offset)),
149 )?;
150 let muldiv_chip = Fp2MulDivChip::new(
151 adapter_chip_48.clone(),
152 config48.clone(),
153 start_offset,
154 range_checker.clone(),
155 offline_memory.clone(),
156 );
157 inventory.add_executor(
158 Fp2ExtensionExecutor::Fp2MulDivRv32_48(muldiv_chip),
159 muldiv_opcodes
160 .clone()
161 .map(|x| VmOpcode::from_usize(x + start_offset)),
162 )?;
163 } else {
164 panic!("Modulus too large");
165 }
166 }
167
168 Ok(inventory)
169 }
170}