openvm_algebra_circuit/fp2_chip/
addsub.rs

1use std::{
2    cell::RefCell,
3    rc::Rc,
4    sync::{Arc, Mutex},
5};
6
7use openvm_algebra_transpiler::Fp2Opcode;
8use openvm_circuit::{arch::VmChipWrapper, system::memory::OfflineMemory};
9use openvm_circuit_derive::InstructionExecutor;
10use openvm_circuit_primitives::var_range::{
11    SharedVariableRangeCheckerChip, VariableRangeCheckerBus,
12};
13use openvm_circuit_primitives_derive::{Chip, ChipUsageGetter};
14use openvm_mod_circuit_builder::{
15    ExprBuilder, ExprBuilderConfig, FieldExpr, FieldExpressionCoreChip,
16};
17use openvm_rv32_adapters::Rv32VecHeapAdapterChip;
18use openvm_stark_backend::p3_field::PrimeField32;
19
20use crate::Fp2;
21
22// Input: Fp2 * 2
23// Output: Fp2
24#[derive(Chip, ChipUsageGetter, InstructionExecutor)]
25pub struct Fp2AddSubChip<F: PrimeField32, const BLOCKS: usize, const BLOCK_SIZE: usize>(
26    pub  VmChipWrapper<
27        F,
28        Rv32VecHeapAdapterChip<F, 2, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>,
29        FieldExpressionCoreChip,
30    >,
31);
32
33impl<F: PrimeField32, const BLOCKS: usize, const BLOCK_SIZE: usize>
34    Fp2AddSubChip<F, BLOCKS, BLOCK_SIZE>
35{
36    pub fn new(
37        adapter: Rv32VecHeapAdapterChip<F, 2, BLOCKS, BLOCKS, BLOCK_SIZE, BLOCK_SIZE>,
38        config: ExprBuilderConfig,
39        offset: usize,
40        range_checker: SharedVariableRangeCheckerChip,
41        offline_memory: Arc<Mutex<OfflineMemory<F>>>,
42    ) -> Self {
43        let (expr, is_add_flag, is_sub_flag) = fp2_addsub_expr(config, range_checker.bus());
44        let core = FieldExpressionCoreChip::new(
45            expr,
46            offset,
47            vec![
48                Fp2Opcode::ADD as usize,
49                Fp2Opcode::SUB as usize,
50                Fp2Opcode::SETUP_ADDSUB as usize,
51            ],
52            vec![is_add_flag, is_sub_flag],
53            range_checker,
54            "Fp2AddSub",
55            false,
56        );
57        Self(VmChipWrapper::new(adapter, core, offline_memory))
58    }
59}
60
61pub fn fp2_addsub_expr(
62    config: ExprBuilderConfig,
63    range_bus: VariableRangeCheckerBus,
64) -> (FieldExpr, usize, usize) {
65    config.check_valid();
66    let builder = ExprBuilder::new(config, range_bus.range_max_bits);
67    let builder = Rc::new(RefCell::new(builder));
68
69    let mut x = Fp2::new(builder.clone());
70    let mut y = Fp2::new(builder.clone());
71    let add = x.add(&mut y);
72    let sub = x.sub(&mut y);
73
74    let is_add_flag = builder.borrow_mut().new_flag();
75    let is_sub_flag = builder.borrow_mut().new_flag();
76    let diff = Fp2::select(is_sub_flag, &sub, &x);
77    let mut z = Fp2::select(is_add_flag, &add, &diff);
78    z.save_output();
79
80    let builder = builder.borrow().clone();
81    (
82        FieldExpr::new(builder, range_bus, true),
83        is_add_flag,
84        is_sub_flag,
85    )
86}
87
88#[cfg(test)]
89mod tests {
90
91    use halo2curves_axiom::{bn256::Fq2, ff::Field};
92    use itertools::Itertools;
93    use openvm_algebra_transpiler::Fp2Opcode;
94    use openvm_circuit::arch::testing::{VmChipTestBuilder, BITWISE_OP_LOOKUP_BUS};
95    use openvm_circuit_primitives::bitwise_op_lookup::{
96        BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip,
97    };
98    use openvm_instructions::{riscv::RV32_CELL_BITS, LocalOpcode};
99    use openvm_mod_circuit_builder::{
100        test_utils::{biguint_to_limbs, bn254_fq2_to_biguint_vec, bn254_fq_to_biguint},
101        ExprBuilderConfig,
102    };
103    use openvm_pairing_guest::bn254::BN254_MODULUS;
104    use openvm_rv32_adapters::{rv32_write_heap_default, Rv32VecHeapAdapterChip};
105    use openvm_stark_backend::p3_field::FieldAlgebra;
106    use openvm_stark_sdk::p3_baby_bear::BabyBear;
107    use rand::{rngs::StdRng, SeedableRng};
108
109    use super::Fp2AddSubChip;
110
111    const NUM_LIMBS: usize = 32;
112    const LIMB_BITS: usize = 8;
113    type F = BabyBear;
114
115    #[test]
116    fn test_fp2_addsub() {
117        let mut tester: VmChipTestBuilder<F> = VmChipTestBuilder::default();
118        let modulus = BN254_MODULUS.clone();
119        let config = ExprBuilderConfig {
120            modulus: modulus.clone(),
121            num_limbs: NUM_LIMBS,
122            limb_bits: LIMB_BITS,
123        };
124        let bitwise_bus = BitwiseOperationLookupBus::new(BITWISE_OP_LOOKUP_BUS);
125        let bitwise_chip = SharedBitwiseOperationLookupChip::<RV32_CELL_BITS>::new(bitwise_bus);
126        let adapter = Rv32VecHeapAdapterChip::<F, 2, 2, 2, NUM_LIMBS, NUM_LIMBS>::new(
127            tester.execution_bus(),
128            tester.program_bus(),
129            tester.memory_bridge(),
130            tester.address_bits(),
131            bitwise_chip.clone(),
132        );
133        let mut chip = Fp2AddSubChip::new(
134            adapter,
135            config,
136            Fp2Opcode::CLASS_OFFSET,
137            tester.range_checker(),
138            tester.offline_memory_mutex_arc(),
139        );
140
141        let mut rng = StdRng::seed_from_u64(42);
142        let x = Fq2::random(&mut rng);
143        let y = Fq2::random(&mut rng);
144        let inputs = [x.c0, x.c1, y.c0, y.c1].map(bn254_fq_to_biguint);
145
146        let expected_sum = bn254_fq2_to_biguint_vec(x + y);
147        let r_sum = chip
148            .0
149            .core
150            .expr()
151            .execute_with_output(inputs.to_vec(), vec![true, false]);
152        assert_eq!(r_sum.len(), 2);
153        assert_eq!(r_sum[0], expected_sum[0]);
154        assert_eq!(r_sum[1], expected_sum[1]);
155
156        let expected_sub = bn254_fq2_to_biguint_vec(x - y);
157        let r_sub = chip
158            .0
159            .core
160            .expr()
161            .execute_with_output(inputs.to_vec(), vec![false, true]);
162        assert_eq!(r_sub.len(), 2);
163        assert_eq!(r_sub[0], expected_sub[0]);
164        assert_eq!(r_sub[1], expected_sub[1]);
165
166        let x_limbs = inputs[0..2]
167            .iter()
168            .map(|x| {
169                biguint_to_limbs::<NUM_LIMBS>(x.clone(), LIMB_BITS)
170                    .map(BabyBear::from_canonical_u32)
171            })
172            .collect_vec();
173        let y_limbs = inputs[2..4]
174            .iter()
175            .map(|x| {
176                biguint_to_limbs::<NUM_LIMBS>(x.clone(), LIMB_BITS)
177                    .map(BabyBear::from_canonical_u32)
178            })
179            .collect_vec();
180        let modulus =
181            biguint_to_limbs::<NUM_LIMBS>(modulus, LIMB_BITS).map(BabyBear::from_canonical_u32);
182        let zero = [BabyBear::ZERO; NUM_LIMBS];
183        let setup_instruction = rv32_write_heap_default(
184            &mut tester,
185            vec![modulus, zero],
186            vec![zero; 2],
187            chip.0.core.air.offset + Fp2Opcode::SETUP_ADDSUB as usize,
188        );
189        let instruction1 = rv32_write_heap_default(
190            &mut tester,
191            x_limbs.clone(),
192            y_limbs.clone(),
193            chip.0.core.air.offset + Fp2Opcode::ADD as usize,
194        );
195        let instruction2 = rv32_write_heap_default(
196            &mut tester,
197            x_limbs,
198            y_limbs,
199            chip.0.core.air.offset + Fp2Opcode::SUB as usize,
200        );
201        tester.execute(&mut chip, &setup_instruction);
202        tester.execute(&mut chip, &instruction1);
203        tester.execute(&mut chip, &instruction2);
204        let tester = tester.build().load(chip).load(bitwise_chip).finalize();
205        tester.simple_test().expect("Verification failed");
206    }
207}