1use alloc::format;
6use alloc::string::ToString;
7use alloc::vec::Vec;
8use core::array;
9use core::fmt::{self, Display, Formatter};
10use core::iter::{Product, Sum};
11use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
12
13use itertools::Itertools;
14use num_bigint::BigUint;
15use p3_util::{as_base_slice, as_base_slice_mut, reconstitute_from_base};
16
17use super::packed_cubic_extension::PackedCubicTrinomialExtensionField;
18use super::{ExtField, HasFrobenius, HasTwoAdicCubicExtension};
19use crate::extension::{CubicTrinomial, CubicTrinomialExtendable, ExtensionAlgebra};
20use crate::field::Field;
21use crate::{
22 Algebra, ExtensionField, PackedFieldExtension, PrimeCharacteristicRing, RawDataSerializable,
23 TwoAdicField, field_to_array,
24};
25
26pub type CubicTrinomialExtensionField<F, A = F> = ExtField<F, 3, CubicTrinomial, A>;
32
33impl<F: Copy> CubicTrinomialExtensionField<F, F> {
34 #[inline]
41 pub const fn new_array<const N: usize>(input: [[F; 3]; N]) -> [Self; N] {
42 const { assert!(N > 0) }
43 let mut output = [Self::new(input[0]); N];
44 let mut i = 1;
45 while i < N {
46 output[i] = Self::new(input[i]);
47 i += 1;
48 }
49 output
50 }
51}
52
53impl<F: CubicTrinomialExtendable> ExtensionField<F> for CubicTrinomialExtensionField<F>
54where
55 PackedCubicTrinomialExtensionField<F, F::Packing>: PackedFieldExtension<F, Self>,
56{
57 type ExtensionPacking = PackedCubicTrinomialExtensionField<F, F::Packing>;
58
59 #[inline]
60 fn is_in_basefield(&self) -> bool {
61 self.value[1..].iter().all(F::is_zero)
62 }
63
64 #[inline]
65 fn as_base(&self) -> Option<F> {
66 <Self as ExtensionField<F>>::is_in_basefield(self).then(|| self.value[0])
67 }
68}
69
70impl<F: CubicTrinomialExtendable> HasFrobenius<F> for CubicTrinomialExtensionField<F> {
71 #[inline]
73 fn frobenius(&self) -> Self {
74 let a = &self.value;
75 let m = F::FROBENIUS_MATRIX;
76 let c0 = a[0] + m[0][1] * a[1] + m[0][2] * a[2];
77 let c1 = m[1][1] * a[1] + m[1][2] * a[2];
78 let c2 = m[2][1] * a[1] + m[2][2] * a[2];
79 Self::new([c0, c1, c2])
80 }
81
82 #[inline]
84 fn repeated_frobenius(&self, count: usize) -> Self {
85 match count % 3 {
86 0 => *self,
87 _ => {
88 let mut result = *self;
89 for _ in 0..(count % 3) {
90 result = result.frobenius();
91 }
92 result
93 }
94 }
95 }
96
97 #[inline]
105 fn pseudo_inv(&self) -> Self {
106 if self.is_zero() {
107 return Self::ZERO;
108 }
109 let a_exp_p = self.frobenius();
110 let prod_conj = (*self * a_exp_p).frobenius();
111 let norm = self.compute_norm_with_prod_conj(&prod_conj);
112 debug_assert_eq!(Self::from(norm), *self * prod_conj);
113 prod_conj * norm.inverse()
114 }
115}
116
117impl<F: CubicTrinomialExtendable> CubicTrinomialExtensionField<F> {
118 #[inline]
123 fn compute_norm_with_prod_conj(&self, prod_conj: &Self) -> F {
124 let a = &self.value;
125 let b = &prod_conj.value;
126
127 let c0 = a[0] * b[0];
129 let c3 = F::dot_product::<2>(&[a[1], a[2]], &[b[2], b[1]]);
130
131 c0 + c3
132 }
133}
134
135impl<F, A> PrimeCharacteristicRing for CubicTrinomialExtensionField<F, A>
136where
137 F: CubicTrinomialExtendable,
138 A: ExtensionAlgebra<F, 3, CubicTrinomial> + Copy,
139{
140 type PrimeSubfield = <A as PrimeCharacteristicRing>::PrimeSubfield;
141
142 const ZERO: Self = Self::new([A::ZERO; 3]);
143 const ONE: Self = Self::new(field_to_array(A::ONE));
144 const TWO: Self = Self::new(field_to_array(A::TWO));
145 const NEG_ONE: Self = Self::new(field_to_array(A::NEG_ONE));
146
147 #[inline]
148 fn from_prime_subfield(f: Self::PrimeSubfield) -> Self {
149 <A as PrimeCharacteristicRing>::from_prime_subfield(f).into()
150 }
151
152 #[inline]
153 fn halve(&self) -> Self {
154 Self::new(array::from_fn(|i| self.value[i].halve()))
155 }
156
157 #[inline(always)]
158 fn square(&self) -> Self {
159 let mut res = Self::default();
160 <A as ExtensionAlgebra<F, 3, CubicTrinomial>>::ext_square(&self.value, &mut res.value);
161 res
162 }
163
164 #[inline]
165 fn mul_2exp_u64(&self, exp: u64) -> Self {
166 Self::new(array::from_fn(|i| self.value[i].mul_2exp_u64(exp)))
167 }
168
169 #[inline]
170 fn div_2exp_u64(&self, exp: u64) -> Self {
171 Self::new(array::from_fn(|i| self.value[i].div_2exp_u64(exp)))
172 }
173
174 #[inline]
175 fn zero_vec(len: usize) -> Vec<Self> {
176 unsafe { reconstitute_from_base(A::zero_vec(len * 3)) }
177 }
178}
179
180impl<F: CubicTrinomialExtendable> Algebra<F> for CubicTrinomialExtensionField<F> {}
181
182impl<F: CubicTrinomialExtendable> RawDataSerializable for CubicTrinomialExtensionField<F> {
183 const NUM_BYTES: usize = F::NUM_BYTES * 3;
184
185 #[inline]
186 fn into_bytes(self) -> impl IntoIterator<Item = u8> {
187 self.value.into_iter().flat_map(|x| x.into_bytes())
188 }
189
190 #[inline]
191 fn into_byte_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u8> {
192 F::into_byte_stream(input.into_iter().flat_map(|x| x.value))
193 }
194
195 #[inline]
196 fn into_u32_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u32> {
197 F::into_u32_stream(input.into_iter().flat_map(|x| x.value))
198 }
199
200 #[inline]
201 fn into_u64_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u64> {
202 F::into_u64_stream(input.into_iter().flat_map(|x| x.value))
203 }
204
205 #[inline]
206 fn into_parallel_byte_streams<const N: usize>(
207 input: impl IntoIterator<Item = [Self; N]>,
208 ) -> impl IntoIterator<Item = [u8; N]> {
209 F::into_parallel_byte_streams(
210 input
211 .into_iter()
212 .flat_map(|x| (0..3).map(move |i| array::from_fn(|j| x[j].value[i]))),
213 )
214 }
215
216 #[inline]
217 fn into_parallel_u32_streams<const N: usize>(
218 input: impl IntoIterator<Item = [Self; N]>,
219 ) -> impl IntoIterator<Item = [u32; N]> {
220 F::into_parallel_u32_streams(
221 input
222 .into_iter()
223 .flat_map(|x| (0..3).map(move |i| array::from_fn(|j| x[j].value[i]))),
224 )
225 }
226
227 #[inline]
228 fn into_parallel_u64_streams<const N: usize>(
229 input: impl IntoIterator<Item = [Self; N]>,
230 ) -> impl IntoIterator<Item = [u64; N]> {
231 F::into_parallel_u64_streams(
232 input
233 .into_iter()
234 .flat_map(|x| (0..3).map(move |i| array::from_fn(|j| x[j].value[i]))),
235 )
236 }
237}
238
239impl<F: CubicTrinomialExtendable> Field for CubicTrinomialExtensionField<F> {
240 type Packing = Self;
241
242 const GENERATOR: Self = Self::new(F::EXT_GENERATOR);
243
244 fn try_inverse(&self) -> Option<Self> {
245 if self.is_zero() {
246 return None;
247 }
248 Some(self.pseudo_inv())
249 }
250
251 #[inline]
252 fn add_slices(slice_1: &mut [Self], slice_2: &[Self]) {
253 unsafe {
254 let base_slice_1 = as_base_slice_mut(slice_1);
255 let base_slice_2 = as_base_slice(slice_2);
256 F::add_slices(base_slice_1, base_slice_2);
257 }
258 }
259
260 #[inline]
261 fn order() -> BigUint {
262 F::order().pow(3)
263 }
264}
265
266impl<F: CubicTrinomialExtendable> Display for CubicTrinomialExtensionField<F> {
267 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
268 if self.is_zero() {
269 write!(f, "0")
270 } else {
271 let str = self
272 .value
273 .iter()
274 .enumerate()
275 .filter(|(_, x)| !x.is_zero())
276 .map(|(i, x)| match (i, x.is_one()) {
277 (0, _) => format!("{x}"),
278 (1, true) => "X".to_string(),
279 (1, false) => format!("{x} X"),
280 (_, true) => format!("X^{i}"),
281 (_, false) => format!("{x} X^{i}"),
282 })
283 .join(" + ");
284 write!(f, "{str}")
285 }
286 }
287}
288
289impl<F, A> Neg for CubicTrinomialExtensionField<F, A>
290where
291 F: CubicTrinomialExtendable,
292 A: Algebra<F>,
293{
294 type Output = Self;
295
296 #[inline]
297 fn neg(self) -> Self {
298 Self::new(self.value.map(A::neg))
299 }
300}
301
302impl<F, A> Add for CubicTrinomialExtensionField<F, A>
303where
304 F: CubicTrinomialExtendable,
305 A: ExtensionAlgebra<F, 3, CubicTrinomial>,
306{
307 type Output = Self;
308
309 #[inline]
310 fn add(self, rhs: Self) -> Self {
311 Self::new(<A as ExtensionAlgebra<F, 3, CubicTrinomial>>::ext_add(
312 &self.value,
313 &rhs.value,
314 ))
315 }
316}
317
318impl<F, A> Add<A> for CubicTrinomialExtensionField<F, A>
319where
320 F: CubicTrinomialExtendable,
321 A: Algebra<F>,
322{
323 type Output = Self;
324
325 #[inline]
326 fn add(mut self, rhs: A) -> Self {
327 self.value[0] += rhs;
328 self
329 }
330}
331
332impl<F, A> AddAssign for CubicTrinomialExtensionField<F, A>
333where
334 F: CubicTrinomialExtendable,
335 A: ExtensionAlgebra<F, 3, CubicTrinomial>,
336{
337 #[inline]
338 fn add_assign(&mut self, rhs: Self) {
339 self.value =
340 <A as ExtensionAlgebra<F, 3, CubicTrinomial>>::ext_add(&self.value, &rhs.value);
341 }
342}
343
344impl<F, A> AddAssign<A> for CubicTrinomialExtensionField<F, A>
345where
346 F: CubicTrinomialExtendable,
347 A: Algebra<F>,
348{
349 #[inline]
350 fn add_assign(&mut self, rhs: A) {
351 self.value[0] += rhs;
352 }
353}
354
355impl<F, A> Sum for CubicTrinomialExtensionField<F, A>
356where
357 F: CubicTrinomialExtendable,
358 A: ExtensionAlgebra<F, 3, CubicTrinomial> + Copy,
359{
360 #[inline]
361 fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
362 iter.reduce(|acc, x| acc + x).unwrap_or(Self::ZERO)
363 }
364}
365
366impl<F, A> Sub for CubicTrinomialExtensionField<F, A>
367where
368 F: CubicTrinomialExtendable,
369 A: ExtensionAlgebra<F, 3, CubicTrinomial>,
370{
371 type Output = Self;
372
373 #[inline]
374 fn sub(self, rhs: Self) -> Self {
375 Self::new(<A as ExtensionAlgebra<F, 3, CubicTrinomial>>::ext_sub(
376 &self.value,
377 &rhs.value,
378 ))
379 }
380}
381
382impl<F, A> Sub<A> for CubicTrinomialExtensionField<F, A>
383where
384 F: CubicTrinomialExtendable,
385 A: Algebra<F>,
386{
387 type Output = Self;
388
389 #[inline]
390 fn sub(self, rhs: A) -> Self {
391 let mut res = self.value;
392 res[0] -= rhs;
393 Self::new(res)
394 }
395}
396
397impl<F, A> SubAssign for CubicTrinomialExtensionField<F, A>
398where
399 F: CubicTrinomialExtendable,
400 A: ExtensionAlgebra<F, 3, CubicTrinomial>,
401{
402 #[inline]
403 fn sub_assign(&mut self, rhs: Self) {
404 self.value =
405 <A as ExtensionAlgebra<F, 3, CubicTrinomial>>::ext_sub(&self.value, &rhs.value);
406 }
407}
408
409impl<F, A> SubAssign<A> for CubicTrinomialExtensionField<F, A>
410where
411 F: CubicTrinomialExtendable,
412 A: Algebra<F>,
413{
414 #[inline]
415 fn sub_assign(&mut self, rhs: A) {
416 self.value[0] -= rhs;
417 }
418}
419
420impl<F, A> Mul for CubicTrinomialExtensionField<F, A>
421where
422 F: CubicTrinomialExtendable,
423 A: ExtensionAlgebra<F, 3, CubicTrinomial>,
424{
425 type Output = Self;
426
427 #[inline]
428 fn mul(self, rhs: Self) -> Self {
429 let mut res = Self::default();
430 <A as ExtensionAlgebra<F, 3, CubicTrinomial>>::ext_mul(
431 &self.value,
432 &rhs.value,
433 &mut res.value,
434 );
435 res
436 }
437}
438
439impl<F, A> Mul<A> for CubicTrinomialExtensionField<F, A>
440where
441 F: CubicTrinomialExtendable,
442 A: ExtensionAlgebra<F, 3, CubicTrinomial>,
443{
444 type Output = Self;
445
446 #[inline]
447 fn mul(self, rhs: A) -> Self {
448 Self::new(<A as ExtensionAlgebra<F, 3, CubicTrinomial>>::ext_base_mul(
449 self.value, rhs,
450 ))
451 }
452}
453
454impl<F, A> MulAssign for CubicTrinomialExtensionField<F, A>
455where
456 F: CubicTrinomialExtendable,
457 A: ExtensionAlgebra<F, 3, CubicTrinomial>,
458{
459 #[inline]
460 fn mul_assign(&mut self, rhs: Self) {
461 *self = self.clone() * rhs;
462 }
463}
464
465impl<F, A> MulAssign<A> for CubicTrinomialExtensionField<F, A>
466where
467 F: CubicTrinomialExtendable,
468 A: ExtensionAlgebra<F, 3, CubicTrinomial>,
469{
470 #[inline]
471 fn mul_assign(&mut self, rhs: A) {
472 *self = self.clone() * rhs;
473 }
474}
475
476impl<F, A> Product for CubicTrinomialExtensionField<F, A>
477where
478 F: CubicTrinomialExtendable,
479 A: ExtensionAlgebra<F, 3, CubicTrinomial> + Copy,
480{
481 #[inline]
482 fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
483 iter.reduce(|acc, x| acc * x).unwrap_or(Self::ONE)
484 }
485}
486
487impl<F> Div for CubicTrinomialExtensionField<F>
488where
489 F: CubicTrinomialExtendable,
490{
491 type Output = Self;
492
493 #[allow(clippy::suspicious_arithmetic_impl)]
494 #[inline]
495 fn div(self, rhs: Self) -> Self::Output {
496 self * rhs.inverse()
497 }
498}
499
500impl<F> DivAssign for CubicTrinomialExtensionField<F>
501where
502 F: CubicTrinomialExtendable,
503{
504 #[inline]
505 fn div_assign(&mut self, rhs: Self) {
506 *self = *self / rhs;
507 }
508}
509
510impl<F: CubicTrinomialExtendable + HasTwoAdicCubicExtension> TwoAdicField
511 for CubicTrinomialExtensionField<F>
512{
513 const TWO_ADICITY: usize = F::EXT_TWO_ADICITY;
514
515 #[inline]
516 fn two_adic_generator(bits: usize) -> Self {
517 Self::new(F::ext_two_adic_generator(bits))
518 }
519}
520
521#[inline]
523pub fn trinomial_cubic_mul<R: PrimeCharacteristicRing>(a: &[R; 3], b: &[R; 3], res: &mut [R; 3]) {
524 let b0_plus_b2 = b[0].dup() + b[2].dup();
525 let b1_plus_b2 = b[1].dup() + b[2].dup();
526
527 res[0] = R::dot_product::<3>(a, &[b[0].dup(), b[2].dup(), b[1].dup()]);
528 res[1] = R::dot_product::<3>(a, &[b[1].dup(), b0_plus_b2.dup(), b1_plus_b2]);
529 res[2] = R::dot_product::<3>(a, &[b[2].dup(), b[1].dup(), b0_plus_b2]);
530}
531
532#[inline]
533pub fn cubic_square<R: PrimeCharacteristicRing>(a: &[R; 3], res: &mut [R; 3]) {
534 let a0_plus_a2 = a[0].dup() + a[2].dup();
535 let a1_plus_a2 = a[1].dup() + a[2].dup();
536
537 res[0] = R::dot_product::<3>(a, &[a[0].dup(), a[2].dup(), a[1].dup()]);
538 res[1] = R::dot_product::<3>(a, &[a[1].dup(), a0_plus_a2.dup(), a1_plus_a2]);
539 res[2] = R::dot_product::<3>(a, &[a[2].dup(), a[1].dup(), a0_plus_a2]);
540}