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::PrimeCharacteristicRing, 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,
216    };
217    use rand08::{rngs::StdRng, SeedableRng};
218
219    use super::*;
220
221    fn generate_random_fq12() -> Fq12 {
222        let mut rng = StdRng::from_entropy();
223        Fq12::random(&mut rng)
224    }
225
226    fn run_fp12_test(
227        x: Fq12,
228        y: Fq12,
229        xi: Option<[isize; 2]>,
230        fp12_fn_addsub: Option<impl Fn(&mut Fp12, &mut Fp12) -> Fp12>,
231        fp12_fn_mul: Option<impl Fn(&mut Fp12, &mut Fp12, [isize; 2]) -> Fp12>,
232        fq12_fn: impl Fn(&Fq12, &Fq12) -> Fq12,
233    ) {
234        if fp12_fn_addsub.is_none() && fp12_fn_mul.is_none() {
235            panic!("Either fp12_fn_addsub or fp12_fn_mul must be provided");
236        }
237        if fp12_fn_addsub.is_some() && fp12_fn_mul.is_some() {
238            panic!("Only one of fp12_fn_addsub or fp12_fn_mul must be provided");
239        }
240
241        let prime = BN254_MODULUS.clone();
242        let (range_checker, builder) = setup(&prime);
243
244        let mut x_fp12 = Fp12::new(builder.clone());
245        let mut y_fp12 = Fp12::new(builder.clone());
246        let mut r = if let Some(fp12_fn_addsub) = fp12_fn_addsub {
247            fp12_fn_addsub(&mut x_fp12, &mut y_fp12)
248        } else {
249            let fp12_fn_mul = fp12_fn_mul.unwrap();
250            fp12_fn_mul(&mut x_fp12, &mut y_fp12, xi.unwrap())
251        };
252        let indices = r.save();
253
254        let builder = builder.borrow().clone();
255        let air = FieldExpr::new(builder, range_checker.bus(), false);
256        let width = BaseAir::<BabyBear>::width(&air);
257
258        let x_fq12 = x;
259        let y_fq12 = y;
260        let r_fq12 = fq12_fn(&x_fq12, &y_fq12);
261        let mut inputs = bn254_fq12_to_biguint_vec(x_fq12);
262        inputs.extend(bn254_fq12_to_biguint_vec(y_fq12));
263
264        let mut row = BabyBear::zero_vec(width);
265        air.generate_subrow((&range_checker, inputs, vec![]), &mut row);
266        let FieldExprCols { vars, .. } = air.load_vars(&row);
267        let trace = RowMajorMatrix::new(row, width);
268        let range_trace = range_checker.generate_trace();
269
270        for (idx, v) in indices
271            .iter()
272            .zip(r_fq12.to_coeffs().into_iter().flat_map(|x| [x.c0, x.c1]))
273        {
274            assert_eq!(
275                evaluate_biguint(&vars[*idx], LIMB_BITS),
276                bn254_fq_to_biguint(v)
277            );
278        }
279
280        BabyBearBlake3Engine::run_simple_test_no_pis_fast(
281            any_rap_arc_vec![air, range_checker.air],
282            vec![trace, range_trace],
283        )
284        .expect("Verification failed");
285    }
286
287    #[test]
288    fn test_fp12_add() {
289        let x = generate_random_fq12();
290        let y = generate_random_fq12();
291        run_fp12_test(
292            x,
293            y,
294            None,
295            Some(Fp12::add),
296            None::<fn(&mut Fp12, &mut Fp12, [isize; 2]) -> Fp12>,
297            |x, y| x + y,
298        );
299    }
300
301    #[test]
302    fn test_fp12_sub() {
303        let x = generate_random_fq12();
304        let y = generate_random_fq12();
305        run_fp12_test(
306            x,
307            y,
308            None,
309            Some(Fp12::sub),
310            None::<fn(&mut Fp12, &mut Fp12, [isize; 2]) -> Fp12>,
311            |x, y| x - y,
312        );
313    }
314
315    #[test]
316    fn test_fp12_mul() {
317        let x = generate_random_fq12();
318        let y = generate_random_fq12();
319        let xi = [9, 1];
320        run_fp12_test(
321            x,
322            y,
323            Some(xi),
324            None::<fn(&mut Fp12, &mut Fp12) -> Fp12>,
325            Some(Fp12::mul),
326            |x, y| x * y,
327        );
328    }
329}