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::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}