openvm_algebra_circuit/fp2_chip/
muldiv.rs

1use std::{cell::RefCell, rc::Rc};
2
3use openvm_algebra_transpiler::Fp2Opcode;
4use openvm_circuit::{
5    arch::ExecutionBridge,
6    system::memory::{offline_checker::MemoryBridge, SharedMemoryHelper},
7};
8use openvm_circuit_primitives::{
9    bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip},
10    var_range::{SharedVariableRangeCheckerChip, VariableRangeCheckerBus},
11};
12use openvm_instructions::riscv::RV32_CELL_BITS;
13use openvm_mod_circuit_builder::{
14    ExprBuilder, ExprBuilderConfig, FieldExpr, FieldExpressionCoreAir, FieldExpressionExecutor,
15    FieldExpressionFiller, SymbolicExpr,
16};
17use openvm_rv32_adapters::{
18    Rv32VecHeapAdapterAir, Rv32VecHeapAdapterExecutor, Rv32VecHeapAdapterFiller,
19};
20
21use super::{Fp2Air, Fp2Chip, Fp2Executor};
22use crate::{FieldExprVecHeapExecutor, Fp2};
23
24pub fn fp2_muldiv_expr(
25    config: ExprBuilderConfig,
26    range_bus: VariableRangeCheckerBus,
27) -> (FieldExpr, usize, usize) {
28    config.check_valid();
29    let builder = ExprBuilder::new(config, range_bus.range_max_bits);
30    let builder = Rc::new(RefCell::new(builder));
31
32    let x = Fp2::new(builder.clone());
33    let mut y = Fp2::new(builder.clone());
34    let is_mul_flag = builder.borrow_mut().new_flag();
35    let is_div_flag = builder.borrow_mut().new_flag();
36    let (z_idx, mut z) = Fp2::new_var(builder.clone());
37
38    let mut lvar = Fp2::select(is_mul_flag, &x, &z);
39
40    let mut rvar = Fp2::select(is_mul_flag, &z, &x);
41    let fp2_constraint = lvar.mul(&mut y).sub(&mut rvar);
42    // When it's SETUP op, the constraints is z * y - x = 0, it still works as:
43    // x.c0 = x.c1 = p == 0, y.c0 = y.c1 = 0, so whatever z is, z * 0 - 0 = 0
44
45    z.save_output();
46    builder
47        .borrow_mut()
48        .set_constraint(z_idx.0, fp2_constraint.c0.expr);
49    builder
50        .borrow_mut()
51        .set_constraint(z_idx.1, fp2_constraint.c1.expr);
52
53    // Compute expression has to be done manually at the SymbolicExpr level.
54    // Otherwise it saves the quotient and introduces new variables.
55    let compute_z0_div = (&x.c0.expr * &y.c0.expr + &x.c1.expr * &y.c1.expr)
56        / (&y.c0.expr * &y.c0.expr + &y.c1.expr * &y.c1.expr);
57    let compute_z0_mul = &x.c0.expr * &y.c0.expr - &x.c1.expr * &y.c1.expr;
58    let compute_z0 = SymbolicExpr::Select(
59        is_mul_flag,
60        Box::new(compute_z0_mul),
61        Box::new(SymbolicExpr::Select(
62            is_div_flag,
63            Box::new(compute_z0_div),
64            Box::new(x.c0.expr.clone()),
65        )),
66    );
67    let compute_z1_div = (&x.c1.expr * &y.c0.expr - &x.c0.expr * &y.c1.expr)
68        / (&y.c0.expr * &y.c0.expr + &y.c1.expr * &y.c1.expr);
69    let compute_z1_mul = &x.c1.expr * &y.c0.expr + &x.c0.expr * &y.c1.expr;
70    let compute_z1 = SymbolicExpr::Select(
71        is_mul_flag,
72        Box::new(compute_z1_mul),
73        Box::new(SymbolicExpr::Select(
74            is_div_flag,
75            Box::new(compute_z1_div),
76            Box::new(x.c1.expr),
77        )),
78    );
79    builder.borrow_mut().set_compute(z_idx.0, compute_z0);
80    builder.borrow_mut().set_compute(z_idx.1, compute_z1);
81
82    let builder = builder.borrow().clone();
83    (
84        FieldExpr::new(builder, range_bus, true),
85        is_mul_flag,
86        is_div_flag,
87    )
88}
89
90// Input: Fp2 * 2
91// Output: Fp2
92
93fn gen_base_expr(
94    config: ExprBuilderConfig,
95    range_checker_bus: VariableRangeCheckerBus,
96) -> (FieldExpr, Vec<usize>, Vec<usize>) {
97    let (expr, is_mul_flag, is_div_flag) = fp2_muldiv_expr(config, range_checker_bus);
98
99    let local_opcode_idx = vec![
100        Fp2Opcode::MUL as usize,
101        Fp2Opcode::DIV as usize,
102        Fp2Opcode::SETUP_MULDIV as usize,
103    ];
104    let opcode_flag_idx = vec![is_mul_flag, is_div_flag];
105
106    (expr, local_opcode_idx, opcode_flag_idx)
107}
108
109pub fn get_fp2_muldiv_air<const BLOCKS: usize, const BLOCK_SIZE: usize>(
110    exec_bridge: ExecutionBridge,
111    mem_bridge: MemoryBridge,
112    config: ExprBuilderConfig,
113    range_checker_bus: VariableRangeCheckerBus,
114    bitwise_lookup_bus: BitwiseOperationLookupBus,
115    pointer_max_bits: usize,
116    offset: usize,
117) -> Fp2Air<BLOCKS, BLOCK_SIZE> {
118    let (expr, local_opcode_idx, opcode_flag_idx) = gen_base_expr(config, range_checker_bus);
119    Fp2Air::new(
120        Rv32VecHeapAdapterAir::new(
121            exec_bridge,
122            mem_bridge,
123            bitwise_lookup_bus,
124            pointer_max_bits,
125        ),
126        FieldExpressionCoreAir::new(expr, offset, local_opcode_idx, opcode_flag_idx),
127    )
128}
129
130pub fn get_fp2_muldiv_step<const BLOCKS: usize, const BLOCK_SIZE: usize>(
131    config: ExprBuilderConfig,
132    range_checker_bus: VariableRangeCheckerBus,
133    pointer_max_bits: usize,
134    offset: usize,
135) -> Fp2Executor<BLOCKS, BLOCK_SIZE> {
136    let (expr, local_opcode_idx, opcode_flag_idx) = gen_base_expr(config, range_checker_bus);
137
138    FieldExprVecHeapExecutor(FieldExpressionExecutor::new(
139        Rv32VecHeapAdapterExecutor::new(pointer_max_bits),
140        expr,
141        offset,
142        local_opcode_idx,
143        opcode_flag_idx,
144        "Fp2MulDiv",
145    ))
146}
147
148pub fn get_fp2_muldiv_chip<F, const BLOCKS: usize, const BLOCK_SIZE: usize>(
149    config: ExprBuilderConfig,
150    mem_helper: SharedMemoryHelper<F>,
151    range_checker: SharedVariableRangeCheckerChip,
152    bitwise_lookup_chip: SharedBitwiseOperationLookupChip<RV32_CELL_BITS>,
153    pointer_max_bits: usize,
154) -> Fp2Chip<F, BLOCKS, BLOCK_SIZE> {
155    let (expr, local_opcode_idx, opcode_flag_idx) = gen_base_expr(config, range_checker.bus());
156    Fp2Chip::new(
157        FieldExpressionFiller::new(
158            Rv32VecHeapAdapterFiller::new(pointer_max_bits, bitwise_lookup_chip),
159            expr,
160            local_opcode_idx,
161            opcode_flag_idx,
162            range_checker,
163            false,
164        ),
165        mem_helper,
166    )
167}