openvm_pairing_circuit/
fp12.rs

1use std::{array::from_fn, cell::RefCell, rc::Rc};
2
3use openvm_algebra_circuit::Fp2;
4use openvm_mod_circuit_builder::{ExprBuilder, FieldVariable};
5
6/// Field extension Fp12 defined with coefficients in Fp2.
7/// Represents the element `c0 + c1 w + ... + c5 w^5` in Fp12.
8/// Fp6-equivalent coefficients are c0: (c0, c2, c4), c1: (c1, c3, c5).
9pub struct Fp12 {
10    pub c: [Fp2; 6],
11}
12
13impl Fp12 {
14    pub fn new(builder: Rc<RefCell<ExprBuilder>>) -> Self {
15        let c = from_fn(|_| Fp2::new(builder.clone()));
16        Fp12 { c }
17    }
18
19    pub fn save(&mut self) -> [usize; 12] {
20        self.c
21            .each_mut()
22            .map(|c| c.save())
23            .concat()
24            .try_into()
25            .unwrap()
26    }
27
28    pub fn save_output(&mut self) {
29        for c in self.c.iter_mut() {
30            c.save_output();
31        }
32    }
33
34    pub fn add(&mut self, other: &mut Fp12) -> Fp12 {
35        Fp12 {
36            c: from_fn(|i| self.c[i].add(&mut other.c[i])),
37        }
38    }
39
40    pub fn sub(&mut self, other: &mut Fp12) -> Fp12 {
41        Fp12 {
42            c: from_fn(|i| self.c[i].sub(&mut other.c[i])),
43        }
44    }
45
46    pub fn mul(&mut self, other: &mut Fp12, xi: [isize; 2]) -> Fp12 {
47        let c = from_fn(|i| {
48            let mut sum = self.c[0].mul(&mut other.c[i]);
49            for j in 1..=5.min(i) {
50                let k = i - j;
51                sum = sum.add(&mut self.c[j].mul(&mut other.c[k]));
52            }
53            let mut sum_hi: Option<Fp2> = None;
54            for j in (i + 1)..=5 {
55                let k = 6 + i - j;
56                let mut term = self.c[j].mul(&mut other.c[k]);
57                if let Some(mut running_sum) = sum_hi {
58                    sum_hi = Some(running_sum.add(&mut term));
59                } else {
60                    sum_hi = Some(term);
61                }
62            }
63            if let Some(mut sum_hi) = sum_hi {
64                sum = sum.add(&mut sum_hi.int_mul(xi));
65            }
66            sum.save();
67            sum
68        });
69        Fp12 { c }
70    }
71
72    /// Multiply self by `x0 + x1 w + x2 w^2 + x3 w^3 + x4 w^4` in Fp12.
73    pub fn mul_by_01234(
74        &mut self,
75        x0: &mut Fp2,
76        x1: &mut Fp2,
77        x2: &mut Fp2,
78        x3: &mut Fp2,
79        x4: &mut Fp2,
80        xi: [isize; 2],
81    ) -> Fp12 {
82        let c0 = self.c[0].mul(x0).add(
83            &mut self.c[2]
84                .mul(x4)
85                .add(&mut self.c[3].mul(x3))
86                .add(&mut self.c[4].mul(x2))
87                .add(&mut self.c[5].mul(x1))
88                .int_mul(xi),
89        );
90
91        let c1 = self.c[0].mul(x1).add(&mut self.c[1].mul(x0)).add(
92            &mut self.c[3]
93                .mul(x4)
94                .add(&mut self.c[4].mul(x3))
95                .add(&mut self.c[5].mul(x2))
96                .int_mul(xi),
97        );
98
99        let c2 = self.c[0]
100            .mul(x2)
101            .add(&mut self.c[1].mul(x1))
102            .add(&mut self.c[2].mul(x0))
103            .add(&mut self.c[4].mul(x4).add(&mut self.c[5].mul(x3)).int_mul(xi));
104
105        let c3 = self.c[0]
106            .mul(x3)
107            .add(&mut self.c[1].mul(x2))
108            .add(&mut self.c[2].mul(x1))
109            .add(&mut self.c[3].mul(x0))
110            .add(&mut self.c[5].mul(x4).int_mul(xi));
111
112        let c4 = self.c[0]
113            .mul(x4)
114            .add(&mut self.c[1].mul(x3))
115            .add(&mut self.c[2].mul(x2))
116            .add(&mut self.c[3].mul(x1))
117            .add(&mut self.c[4].mul(x0));
118
119        let c5 = self.c[1]
120            .mul(x4)
121            .add(&mut self.c[2].mul(x3))
122            .add(&mut self.c[3].mul(x2))
123            .add(&mut self.c[4].mul(x1))
124            .add(&mut self.c[5].mul(x0));
125
126        Fp12 {
127            c: [c0, c1, c2, c3, c4, c5],
128        }
129    }
130
131    /// Multiply `self` by `x0 + x2 w^2 + x3 w^3 + x4 w^4 + x5 w^5` in Fp12.
132    pub fn mul_by_02345(
133        &mut self,
134        x0: &mut Fp2,
135        x2: &mut Fp2,
136        x3: &mut Fp2,
137        x4: &mut Fp2,
138        x5: &mut Fp2,
139        xi: [isize; 2],
140    ) -> Fp12 {
141        let c0 = self.c[0].mul(x0).add(
142            &mut self.c[1]
143                .mul(x5)
144                .add(&mut self.c[2].mul(x4))
145                .add(&mut self.c[3].mul(x3))
146                .add(&mut self.c[4].mul(x2))
147                .int_mul(xi),
148        );
149
150        let c1 = self.c[1].mul(x0).add(
151            &mut self.c[2]
152                .mul(x5)
153                .add(&mut self.c[3].mul(x4))
154                .add(&mut self.c[4].mul(x3))
155                .add(&mut self.c[5].mul(x2))
156                .int_mul(xi),
157        );
158
159        let c2 = self.c[0].mul(x2).add(&mut self.c[2].mul(x0)).add(
160            &mut self.c[3]
161                .mul(x5)
162                .add(&mut self.c[4].mul(x4))
163                .add(&mut self.c[5].mul(x3))
164                .int_mul(xi),
165        );
166
167        let c3 = self.c[0]
168            .mul(x3)
169            .add(&mut self.c[1].mul(x2))
170            .add(&mut self.c[3].mul(x0))
171            .add(&mut self.c[4].mul(x5).add(&mut self.c[5].mul(x4)).int_mul(xi));
172
173        let c4 = self.c[0]
174            .mul(x4)
175            .add(&mut self.c[1].mul(x3))
176            .add(&mut self.c[2].mul(x2))
177            .add(&mut self.c[4].mul(x0))
178            .add(&mut self.c[5].mul(x5).int_mul(xi));
179
180        let c5 = self.c[0]
181            .mul(x5)
182            .add(&mut self.c[1].mul(x4))
183            .add(&mut self.c[2].mul(x3))
184            .add(&mut self.c[3].mul(x2))
185            .add(&mut self.c[5].mul(x0));
186
187        Fp12 {
188            c: [c0, c1, c2, c3, c4, c5],
189        }
190    }
191
192    pub fn div(&mut self, _other: &mut Fp12, _xi: [isize; 2]) -> Fp12 {
193        unimplemented!()
194    }
195
196    pub fn scalar_mul(&mut self, fp: &mut FieldVariable) -> Fp12 {
197        Fp12 {
198            c: from_fn(|i| self.c[i].scalar_mul(fp)),
199        }
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use halo2curves_axiom::{bn256::Fq12, ff::Field};
206    use openvm_circuit_primitives::TraceSubRowGenerator;
207    use openvm_ecc_guest::algebra::field::FieldExtension;
208    use openvm_mod_circuit_builder::{test_utils::*, *};
209    use openvm_pairing_guest::bn254::BN254_MODULUS;
210    use openvm_stark_backend::{
211        p3_air::BaseAir, p3_field::FieldAlgebra, p3_matrix::dense::RowMajorMatrix,
212    };
213    use openvm_stark_sdk::{
214        any_rap_arc_vec, config::baby_bear_blake3::BabyBearBlake3Engine, engine::StarkFriEngine,
215        p3_baby_bear::BabyBear, utils::create_seeded_rng,
216    };
217
218    use super::*;
219
220    fn generate_random_fq12() -> Fq12 {
221        let mut rng = create_seeded_rng();
222        Fq12::random(&mut rng)
223    }
224
225    fn run_fp12_test(
226        x: Fq12,
227        y: Fq12,
228        xi: Option<[isize; 2]>,
229        fp12_fn_addsub: Option<impl Fn(&mut Fp12, &mut Fp12) -> Fp12>,
230        fp12_fn_mul: Option<impl Fn(&mut Fp12, &mut Fp12, [isize; 2]) -> Fp12>,
231        fq12_fn: impl Fn(&Fq12, &Fq12) -> Fq12,
232    ) {
233        if fp12_fn_addsub.is_none() && fp12_fn_mul.is_none() {
234            panic!("Either fp12_fn_addsub or fp12_fn_mul must be provided");
235        }
236        if fp12_fn_addsub.is_some() && fp12_fn_mul.is_some() {
237            panic!("Only one of fp12_fn_addsub or fp12_fn_mul must be provided");
238        }
239
240        let prime = BN254_MODULUS.clone();
241        let (range_checker, builder) = setup(&prime);
242
243        let mut x_fp12 = Fp12::new(builder.clone());
244        let mut y_fp12 = Fp12::new(builder.clone());
245        let mut r = if let Some(fp12_fn_addsub) = fp12_fn_addsub {
246            fp12_fn_addsub(&mut x_fp12, &mut y_fp12)
247        } else {
248            let fp12_fn_mul = fp12_fn_mul.unwrap();
249            fp12_fn_mul(&mut x_fp12, &mut y_fp12, xi.unwrap())
250        };
251        let indices = r.save();
252
253        let builder = builder.borrow().clone();
254        let air = FieldExpr::new(builder, range_checker.bus(), false);
255        let width = BaseAir::<BabyBear>::width(&air);
256
257        let x_fq12 = x;
258        let y_fq12 = y;
259        let r_fq12 = fq12_fn(&x_fq12, &y_fq12);
260        let mut inputs = bn254_fq12_to_biguint_vec(x_fq12);
261        inputs.extend(bn254_fq12_to_biguint_vec(y_fq12));
262
263        let mut row = BabyBear::zero_vec(width);
264        air.generate_subrow((&range_checker, inputs, vec![]), &mut row);
265        let FieldExprCols { vars, .. } = air.load_vars(&row);
266        let trace = RowMajorMatrix::new(row, width);
267        let range_trace = range_checker.generate_trace();
268
269        for (idx, v) in indices
270            .iter()
271            .zip(r_fq12.to_coeffs().into_iter().flat_map(|x| [x.c0, x.c1]))
272        {
273            assert_eq!(
274                evaluate_biguint(&vars[*idx], LIMB_BITS),
275                bn254_fq_to_biguint(v)
276            );
277        }
278
279        BabyBearBlake3Engine::run_simple_test_no_pis_fast(
280            any_rap_arc_vec![air, range_checker.air],
281            vec![trace, range_trace],
282        )
283        .expect("Verification failed");
284    }
285
286    #[test]
287    fn test_fp12_add() {
288        let x = generate_random_fq12();
289        let y = generate_random_fq12();
290        run_fp12_test(
291            x,
292            y,
293            None,
294            Some(Fp12::add),
295            None::<fn(&mut Fp12, &mut Fp12, [isize; 2]) -> Fp12>,
296            |x, y| x + y,
297        );
298    }
299
300    #[test]
301    fn test_fp12_sub() {
302        let x = generate_random_fq12();
303        let y = generate_random_fq12();
304        run_fp12_test(
305            x,
306            y,
307            None,
308            Some(Fp12::sub),
309            None::<fn(&mut Fp12, &mut Fp12, [isize; 2]) -> Fp12>,
310            |x, y| x - y,
311        );
312    }
313
314    #[test]
315    fn test_fp12_mul() {
316        let x = generate_random_fq12();
317        let y = generate_random_fq12();
318        let xi = [9, 1];
319        run_fp12_test(
320            x,
321            y,
322            Some(xi),
323            None::<fn(&mut Fp12, &mut Fp12) -> Fp12>,
324            Some(Fp12::mul),
325            |x, y| x * y,
326        );
327    }
328}