openvm_algebra_circuit/
fields.rs

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
26// Cached modulus constants to avoid repeated string parsing
27static 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        // Reduce modulo the field's modulus for non-canonical representations
259        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        // Reduce modulo the field's modulus for non-canonical representations
276        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}