1use halo2curves_axiom::ff::{Field, PrimeField};
2use num_bigint::BigUint;
3use num_traits::Num;
4use once_cell::sync::Lazy;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum FieldType {
8 K256Coordinate = 0,
9 K256Scalar = 1,
10 P256Coordinate = 2,
11 P256Scalar = 3,
12 BN254Coordinate = 4,
13 BN254Scalar = 5,
14 BLS12_381Coordinate = 6,
15 BLS12_381Scalar = 7,
16}
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum Operation {
20 Add = 0,
21 Sub = 1,
22 Mul = 2,
23 Div = 3,
24}
25
26static K256_COORD_MODULUS: Lazy<BigUint> =
28 Lazy::new(get_modulus_as_bigint_slow::<halo2curves_axiom::secq256k1::Fq>);
29static K256_SCALAR_MODULUS: Lazy<BigUint> =
30 Lazy::new(get_modulus_as_bigint_slow::<halo2curves_axiom::secq256k1::Fp>);
31static P256_COORD_MODULUS: Lazy<BigUint> =
32 Lazy::new(get_modulus_as_bigint_slow::<halo2curves_axiom::secp256r1::Fp>);
33static P256_SCALAR_MODULUS: Lazy<BigUint> =
34 Lazy::new(get_modulus_as_bigint_slow::<halo2curves_axiom::secp256r1::Fq>);
35static BN254_COORD_MODULUS: Lazy<BigUint> =
36 Lazy::new(get_modulus_as_bigint_slow::<halo2curves_axiom::bn256::Fq>);
37static BN254_SCALAR_MODULUS: Lazy<BigUint> =
38 Lazy::new(get_modulus_as_bigint_slow::<halo2curves_axiom::bn256::Fr>);
39static BLS12_381_COORD_MODULUS: Lazy<BigUint> =
40 Lazy::new(get_modulus_as_bigint_slow::<halo2curves_axiom::bls12_381::Fq>);
41static BLS12_381_SCALAR_MODULUS: Lazy<BigUint> =
42 Lazy::new(get_modulus_as_bigint_slow::<halo2curves_axiom::bls12_381::Fr>);
43
44fn get_modulus_as_bigint_slow<F: PrimeField>() -> BigUint {
45 BigUint::from_str_radix(F::MODULUS.trim_start_matches("0x"), 16).unwrap()
46}
47
48#[inline]
49pub fn get_field_type(modulus: &BigUint) -> Option<FieldType> {
50 if modulus == &*K256_COORD_MODULUS {
51 return Some(FieldType::K256Coordinate);
52 }
53
54 if modulus == &*K256_SCALAR_MODULUS {
55 return Some(FieldType::K256Scalar);
56 }
57
58 if modulus == &*P256_COORD_MODULUS {
59 return Some(FieldType::P256Coordinate);
60 }
61
62 if modulus == &*P256_SCALAR_MODULUS {
63 return Some(FieldType::P256Scalar);
64 }
65
66 if modulus == &*BN254_COORD_MODULUS {
67 return Some(FieldType::BN254Coordinate);
68 }
69
70 if modulus == &*BN254_SCALAR_MODULUS {
71 return Some(FieldType::BN254Scalar);
72 }
73
74 if modulus == &*BLS12_381_COORD_MODULUS {
75 return Some(FieldType::BLS12_381Coordinate);
76 }
77
78 if modulus == &*BLS12_381_SCALAR_MODULUS {
79 return Some(FieldType::BLS12_381Scalar);
80 }
81
82 None
83}
84
85#[inline]
86pub fn get_fp2_field_type(modulus: &BigUint) -> Option<FieldType> {
87 if modulus == &*BN254_COORD_MODULUS {
88 return Some(FieldType::BN254Coordinate);
89 }
90
91 if modulus == &*BLS12_381_COORD_MODULUS {
92 return Some(FieldType::BLS12_381Coordinate);
93 }
94
95 None
96}
97
98#[inline(always)]
99pub fn field_operation<
100 const FIELD: u8,
101 const BLOCKS: usize,
102 const BLOCK_SIZE: usize,
103 const OP: u8,
104>(
105 input_data: [[[u8; BLOCK_SIZE]; BLOCKS]; 2],
106) -> [[u8; BLOCK_SIZE]; BLOCKS] {
107 match FIELD {
108 x if x == FieldType::K256Coordinate as u8 => {
109 field_operation_256bit::<halo2curves_axiom::secq256k1::Fq, BLOCKS, BLOCK_SIZE, OP>(
110 input_data,
111 )
112 }
113 x if x == FieldType::K256Scalar as u8 => {
114 field_operation_256bit::<halo2curves_axiom::secq256k1::Fp, BLOCKS, BLOCK_SIZE, OP>(
115 input_data,
116 )
117 }
118 x if x == FieldType::P256Coordinate as u8 => {
119 field_operation_256bit::<halo2curves_axiom::secp256r1::Fp, BLOCKS, BLOCK_SIZE, OP>(
120 input_data,
121 )
122 }
123 x if x == FieldType::P256Scalar as u8 => {
124 field_operation_256bit::<halo2curves_axiom::secp256r1::Fq, BLOCKS, BLOCK_SIZE, OP>(
125 input_data,
126 )
127 }
128 x if x == FieldType::BN254Coordinate as u8 => {
129 field_operation_256bit::<halo2curves_axiom::bn256::Fq, BLOCKS, BLOCK_SIZE, OP>(
130 input_data,
131 )
132 }
133 x if x == FieldType::BN254Scalar as u8 => {
134 field_operation_256bit::<halo2curves_axiom::bn256::Fr, BLOCKS, BLOCK_SIZE, OP>(
135 input_data,
136 )
137 }
138 x if x == FieldType::BLS12_381Coordinate as u8 => {
139 field_operation_bls12_381_coordinate::<BLOCKS, BLOCK_SIZE, OP>(input_data)
140 }
141 x if x == FieldType::BLS12_381Scalar as u8 => {
142 field_operation_256bit::<halo2curves_axiom::bls12_381::Fr, BLOCKS, BLOCK_SIZE, OP>(
143 input_data,
144 )
145 }
146 _ => panic!("Unsupported field type: {FIELD}"),
147 }
148}
149
150#[inline(always)]
151pub fn fp2_operation<
152 const FIELD: u8,
153 const BLOCKS: usize,
154 const BLOCK_SIZE: usize,
155 const OP: u8,
156>(
157 input_data: [[[u8; BLOCK_SIZE]; BLOCKS]; 2],
158) -> [[u8; BLOCK_SIZE]; BLOCKS] {
159 match FIELD {
160 x if x == FieldType::BN254Coordinate as u8 => {
161 fp2_operation_bn254::<BLOCKS, BLOCK_SIZE, OP>(input_data)
162 }
163 x if x == FieldType::BLS12_381Coordinate as u8 => {
164 fp2_operation_bls12_381::<BLOCKS, BLOCK_SIZE, OP>(input_data)
165 }
166 _ => panic!("Unsupported field type for Fp2: {FIELD}"),
167 }
168}
169
170#[inline(always)]
171fn field_operation_256bit<
172 F: PrimeField<Repr = [u8; 32]>,
173 const BLOCKS: usize,
174 const BLOCK_SIZE: usize,
175 const OP: u8,
176>(
177 input_data: [[[u8; BLOCK_SIZE]; BLOCKS]; 2],
178) -> [[u8; BLOCK_SIZE]; BLOCKS] {
179 let a = blocks_to_field_element::<F>(input_data[0].as_flattened());
180 let b = blocks_to_field_element::<F>(input_data[1].as_flattened());
181 let c = match OP {
182 x if x == Operation::Add as u8 => a + b,
183 x if x == Operation::Sub as u8 => a - b,
184 x if x == Operation::Mul as u8 => a * b,
185 x if x == Operation::Div as u8 => a * b.invert().unwrap(),
186 _ => panic!("Unsupported operation: {OP}"),
187 };
188
189 let mut output = [[0u8; BLOCK_SIZE]; BLOCKS];
190 field_element_to_blocks(&c, &mut output);
191 output
192}
193
194#[inline(always)]
195fn field_operation_bls12_381_coordinate<
196 const BLOCKS: usize,
197 const BLOCK_SIZE: usize,
198 const OP: u8,
199>(
200 input_data: [[[u8; BLOCK_SIZE]; BLOCKS]; 2],
201) -> [[u8; BLOCK_SIZE]; BLOCKS] {
202 let a = blocks_to_field_element_bls12_381_coordinate(input_data[0].as_flattened());
203 let b = blocks_to_field_element_bls12_381_coordinate(input_data[1].as_flattened());
204 let c = match OP {
205 x if x == Operation::Add as u8 => a + b,
206 x if x == Operation::Sub as u8 => a - b,
207 x if x == Operation::Mul as u8 => a * b,
208 x if x == Operation::Div as u8 => a * b.invert().unwrap(),
209 _ => panic!("Unsupported operation: {OP}"),
210 };
211
212 let mut output = [[0u8; BLOCK_SIZE]; BLOCKS];
213 field_element_to_blocks_bls12_381_coordinate(&c, &mut output);
214 output
215}
216
217#[inline(always)]
218fn fp2_operation_bn254<const BLOCKS: usize, const BLOCK_SIZE: usize, const OP: u8>(
219 input_data: [[[u8; BLOCK_SIZE]; BLOCKS]; 2],
220) -> [[u8; BLOCK_SIZE]; BLOCKS] {
221 let a = blocks_to_fp2_bn254::<BLOCKS, BLOCK_SIZE>(input_data[0].as_ref());
222 let b = blocks_to_fp2_bn254::<BLOCKS, BLOCK_SIZE>(input_data[1].as_ref());
223 let c = match OP {
224 x if x == Operation::Add as u8 => a + b,
225 x if x == Operation::Sub as u8 => a - b,
226 x if x == Operation::Mul as u8 => a * b,
227 x if x == Operation::Div as u8 => a * b.invert().unwrap(),
228 _ => panic!("Unsupported operation: {OP}"),
229 };
230
231 let mut output = [[0u8; BLOCK_SIZE]; BLOCKS];
232 fp2_to_blocks_bn254(&c, &mut output);
233 output
234}
235
236#[inline(always)]
237fn fp2_operation_bls12_381<const BLOCKS: usize, const BLOCK_SIZE: usize, const OP: u8>(
238 input_data: [[[u8; BLOCK_SIZE]; BLOCKS]; 2],
239) -> [[u8; BLOCK_SIZE]; BLOCKS] {
240 let a = blocks_to_fp2_bls12_381::<BLOCKS, BLOCK_SIZE>(input_data[0].as_ref());
241 let b = blocks_to_fp2_bls12_381::<BLOCKS, BLOCK_SIZE>(input_data[1].as_ref());
242 let c = match OP {
243 x if x == Operation::Add as u8 => a + b,
244 x if x == Operation::Sub as u8 => a - b,
245 x if x == Operation::Mul as u8 => a * b,
246 x if x == Operation::Div as u8 => a * b.invert().unwrap(),
247 _ => panic!("Unsupported operation: {OP}"),
248 };
249
250 let mut output = [[0u8; BLOCK_SIZE]; BLOCKS];
251 fp2_to_blocks_bls12_381(&c, &mut output);
252 output
253}
254
255#[inline(always)]
256fn from_repr_with_reduction<F: PrimeField<Repr = [u8; 32]>>(bytes: [u8; 32]) -> F {
257 F::from_repr_vartime(bytes).unwrap_or_else(|| {
258 let modulus = get_modulus_as_bigint_slow::<F>();
260 let value = BigUint::from_bytes_le(&bytes);
261 let reduced = value % modulus;
262
263 let reduced_le_bytes = reduced.to_bytes_le();
264 let mut reduced_bytes = [0u8; 32];
265 reduced_bytes[..reduced_le_bytes.len()]
266 .copy_from_slice(&reduced_le_bytes[..reduced_le_bytes.len()]);
267
268 F::from_repr_vartime(reduced_bytes).unwrap()
269 })
270}
271
272#[inline(always)]
273fn from_repr_with_reduction_bls12_381_coordinate(bytes: [u8; 48]) -> blstrs::Fp {
274 blstrs::Fp::from_bytes_le(&bytes).unwrap_or_else(|| {
275 let modulus = BigUint::from_bytes_le(&blstrs::Fp::char());
277 let value = BigUint::from_bytes_le(&bytes);
278 let reduced = value % modulus;
279
280 let reduced_le_bytes = reduced.to_bytes_le();
281 let mut reduced_bytes = [0u8; 48];
282 reduced_bytes[..reduced_le_bytes.len()]
283 .copy_from_slice(&reduced_le_bytes[..reduced_le_bytes.len()]);
284
285 blstrs::Fp::from_bytes_le(&reduced_bytes).unwrap()
286 })
287}
288
289#[inline(always)]
290pub fn blocks_to_field_element<F: PrimeField<Repr = [u8; 32]>>(blocks: &[u8]) -> F {
291 debug_assert!(blocks.len() == 32);
292 let mut bytes = [0u8; 32];
293 bytes[..blocks.len()].copy_from_slice(&blocks[..blocks.len()]);
294
295 from_repr_with_reduction::<F>(bytes)
296}
297
298#[inline(always)]
299pub fn field_element_to_blocks<F: PrimeField<Repr = [u8; 32]>, const BLOCK_SIZE: usize>(
300 field_element: &F,
301 output: &mut [[u8; BLOCK_SIZE]],
302) {
303 debug_assert!(output.len() * BLOCK_SIZE == 32);
304 let bytes = field_element.to_repr();
305 let mut byte_idx = 0;
306
307 for block in output.iter_mut() {
308 for byte in block.iter_mut() {
309 *byte = if byte_idx < bytes.len() {
310 bytes[byte_idx]
311 } else {
312 0
313 };
314 byte_idx += 1;
315 }
316 }
317}
318
319#[inline(always)]
320pub fn blocks_to_field_element_bls12_381_coordinate(blocks: &[u8]) -> blstrs::Fp {
321 debug_assert!(blocks.len() == 48);
322 let mut bytes = [0u8; 48];
323 bytes[..blocks.len()].copy_from_slice(&blocks[..blocks.len()]);
324
325 from_repr_with_reduction_bls12_381_coordinate(bytes)
326}
327
328#[inline(always)]
329pub fn field_element_to_blocks_bls12_381_coordinate<const BLOCK_SIZE: usize>(
330 field_element: &blstrs::Fp,
331 output: &mut [[u8; BLOCK_SIZE]],
332) {
333 debug_assert!(output.len() * BLOCK_SIZE == 48);
334 let bytes = field_element.to_bytes_le();
335 let mut byte_idx = 0;
336
337 for block in output.iter_mut() {
338 for byte in block.iter_mut() {
339 *byte = if byte_idx < bytes.len() {
340 bytes[byte_idx]
341 } else {
342 0
343 };
344 byte_idx += 1;
345 }
346 }
347}
348
349#[inline(always)]
350fn blocks_to_fp2_bn254<const BLOCKS: usize, const BLOCK_SIZE: usize>(
351 blocks: &[[u8; BLOCK_SIZE]],
352) -> halo2curves_axiom::bn256::Fq2 {
353 let c0 = blocks_to_field_element::<halo2curves_axiom::bn256::Fq>(
354 blocks[..BLOCKS / 2].as_flattened(),
355 );
356 let c1 = blocks_to_field_element::<halo2curves_axiom::bn256::Fq>(
357 blocks[BLOCKS / 2..].as_flattened(),
358 );
359 halo2curves_axiom::bn256::Fq2::new(c0, c1)
360}
361
362#[inline(always)]
363fn fp2_to_blocks_bn254<const BLOCKS: usize, const BLOCK_SIZE: usize>(
364 fp2: &halo2curves_axiom::bn256::Fq2,
365 output: &mut [[u8; BLOCK_SIZE]; BLOCKS],
366) {
367 field_element_to_blocks::<halo2curves_axiom::bn256::Fq, BLOCK_SIZE>(
368 &fp2.c0,
369 &mut output[..BLOCKS / 2],
370 );
371 field_element_to_blocks::<halo2curves_axiom::bn256::Fq, BLOCK_SIZE>(
372 &fp2.c1,
373 &mut output[BLOCKS / 2..],
374 );
375}
376
377#[inline(always)]
378fn blocks_to_fp2_bls12_381<const BLOCKS: usize, const BLOCK_SIZE: usize>(
379 blocks: &[[u8; BLOCK_SIZE]],
380) -> blstrs::Fp2 {
381 let c0 = blocks_to_field_element_bls12_381_coordinate(blocks[..BLOCKS / 2].as_flattened());
382 let c1 = blocks_to_field_element_bls12_381_coordinate(blocks[BLOCKS / 2..].as_flattened());
383 blstrs::Fp2::new(c0, c1)
384}
385
386#[inline(always)]
387fn fp2_to_blocks_bls12_381<const BLOCKS: usize, const BLOCK_SIZE: usize>(
388 fp2: &blstrs::Fp2,
389 output: &mut [[u8; BLOCK_SIZE]; BLOCKS],
390) {
391 field_element_to_blocks_bls12_381_coordinate(&fp2.c0(), &mut output[..BLOCKS / 2]);
392 field_element_to_blocks_bls12_381_coordinate(&fp2.c1(), &mut output[BLOCKS / 2..]);
393}