1use std::{array::from_fn, cell::RefCell, rc::Rc};
2
3use openvm_algebra_circuit::Fp2;
4use openvm_mod_circuit_builder::{ExprBuilder, FieldVariable};
5
6pub 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 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 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}