1use alloc::format;
9use alloc::string::ToString;
10use alloc::vec::Vec;
11use core::array;
12use core::fmt::{self, Debug, Display, Formatter};
13use core::iter::{Product, Sum};
14use core::marker::PhantomData;
15use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
16
17use itertools::Itertools;
18use num_bigint::BigUint;
19use p3_util::{as_base_slice, as_base_slice_mut, flatten_to_base, reconstitute_from_base};
20use rand::distr::StandardUniform;
21use rand::prelude::Distribution;
22use serde::{Deserialize, Serialize};
23
24use super::packed_quintic_extension::PackedQuinticTrinomialExtensionField;
25use super::{HasFrobenius, HasTwoAdicQuinticExtension};
26use crate::extension::{QuinticExtendableAlgebra, QuinticTrinomialExtendable};
27use crate::field::Field;
28use crate::{
29 Algebra, BasedVectorSpace, ExtensionField, Packable, PackedFieldExtension,
30 PrimeCharacteristicRing, RawDataSerializable, TwoAdicField, field_to_array,
31};
32
33#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Serialize, Deserialize, PartialOrd, Ord)]
37#[repr(transparent)] #[must_use]
39pub struct QuinticTrinomialExtensionField<F, A = F> {
40 #[serde(
41 with = "p3_util::array_serialization",
42 bound(serialize = "A: Serialize", deserialize = "A: Deserialize<'de>")
43 )]
44 pub(crate) value: [A; 5],
45 _phantom: PhantomData<F>,
46}
47
48impl<F, A> QuinticTrinomialExtensionField<F, A> {
49 #[inline]
53 pub const fn new(value: [A; 5]) -> Self {
54 Self {
55 value,
56 _phantom: PhantomData,
57 }
58 }
59}
60
61impl<F: Copy> QuinticTrinomialExtensionField<F, F> {
62 #[inline]
67 pub const fn new_array<const N: usize>(input: [[F; 5]; N]) -> [Self; N] {
68 const { assert!(N > 0) }
69 let mut output = [Self::new(input[0]); N];
70 let mut i = 1;
71 while i < N {
72 output[i] = Self::new(input[i]);
73 i += 1;
74 }
75 output
76 }
77}
78
79impl<F: Field, A: Algebra<F>> Default for QuinticTrinomialExtensionField<F, A> {
80 fn default() -> Self {
81 Self::new(array::from_fn(|_| A::ZERO))
82 }
83}
84
85impl<F: Field, A: Algebra<F>> From<A> for QuinticTrinomialExtensionField<F, A> {
86 fn from(x: A) -> Self {
87 Self::new(field_to_array(x))
88 }
89}
90
91impl<F, A> From<[A; 5]> for QuinticTrinomialExtensionField<F, A> {
92 #[inline]
93 fn from(x: [A; 5]) -> Self {
94 Self {
95 value: x,
96 _phantom: PhantomData,
97 }
98 }
99}
100
101impl<F: QuinticTrinomialExtendable> Packable for QuinticTrinomialExtensionField<F> {}
102
103impl<F: QuinticTrinomialExtendable, A: Algebra<F>> BasedVectorSpace<A>
104 for QuinticTrinomialExtensionField<F, A>
105{
106 const DIMENSION: usize = 5;
107
108 #[inline]
109 fn as_basis_coefficients_slice(&self) -> &[A] {
110 &self.value
111 }
112
113 #[inline]
114 fn from_basis_coefficients_fn<Fn: FnMut(usize) -> A>(f: Fn) -> Self {
115 Self::new(array::from_fn(f))
116 }
117
118 #[inline]
119 fn from_basis_coefficients_iter<I: ExactSizeIterator<Item = A>>(mut iter: I) -> Option<Self> {
120 (iter.len() == 5).then(|| Self::new(array::from_fn(|_| iter.next().unwrap())))
121 }
122
123 #[inline]
124 fn flatten_to_base(vec: Vec<Self>) -> Vec<A> {
125 unsafe { flatten_to_base::<A, Self>(vec) }
127 }
128
129 #[inline]
130 fn reconstitute_from_base(vec: Vec<A>) -> Vec<Self> {
131 unsafe { reconstitute_from_base::<A, Self>(vec) }
133 }
134}
135
136impl<F: QuinticTrinomialExtendable> ExtensionField<F> for QuinticTrinomialExtensionField<F>
137where
138 PackedQuinticTrinomialExtensionField<F, F::Packing>: PackedFieldExtension<F, Self>,
139{
140 type ExtensionPacking = PackedQuinticTrinomialExtensionField<F, F::Packing>;
141
142 #[inline]
143 fn is_in_basefield(&self) -> bool {
144 self.value[1..].iter().all(F::is_zero)
145 }
146
147 #[inline]
148 fn as_base(&self) -> Option<F> {
149 <Self as ExtensionField<F>>::is_in_basefield(self).then(|| self.value[0])
150 }
151}
152
153impl<F: QuinticTrinomialExtendable> HasFrobenius<F> for QuinticTrinomialExtensionField<F> {
154 #[inline]
155 fn frobenius(&self) -> Self {
156 let a = &self.value;
157 let fc = &F::FROBENIUS_COEFFS;
158
159 let a_tail = &[a[1], a[2], a[3], a[4]];
161 let c0 = a[0] + F::dot_product::<4>(a_tail, &[fc[0][0], fc[1][0], fc[2][0], fc[3][0]]);
162 let c1 = F::dot_product::<4>(a_tail, &[fc[0][1], fc[1][1], fc[2][1], fc[3][1]]);
163 let c2 = F::dot_product::<4>(a_tail, &[fc[0][2], fc[1][2], fc[2][2], fc[3][2]]);
164 let c3 = F::dot_product::<4>(a_tail, &[fc[0][3], fc[1][3], fc[2][3], fc[3][3]]);
165 let c4 = F::dot_product::<4>(a_tail, &[fc[0][4], fc[1][4], fc[2][4], fc[3][4]]);
166
167 Self::new([c0, c1, c2, c3, c4])
168 }
169
170 #[inline]
172 fn repeated_frobenius(&self, count: usize) -> Self {
173 match count % 5 {
174 0 => *self,
175 _ => {
176 let mut result = *self;
177 for _ in 0..(count % 5) {
178 result = result.frobenius();
179 }
180 result
181 }
182 }
183 }
184
185 #[inline]
193 fn pseudo_inv(&self) -> Self {
194 if self.is_zero() {
195 return Self::ZERO;
196 }
197
198 let a_exp_p = self.frobenius();
201 let a_exp_p_plus_p2 = (*self * a_exp_p).frobenius();
202 let prod_conj = a_exp_p_plus_p2 * a_exp_p_plus_p2.repeated_frobenius(2);
203
204 let norm = self.compute_norm_with_prod_conj(&prod_conj);
207 debug_assert_eq!(Self::from(norm), *self * prod_conj);
208
209 prod_conj * norm.inverse()
210 }
211}
212
213impl<F: QuinticTrinomialExtendable> QuinticTrinomialExtensionField<F> {
214 #[inline]
219 fn compute_norm_with_prod_conj(&self, prod_conj: &Self) -> F {
220 let a = &self.value;
221 let b = &prod_conj.value;
222
223 let c0 = a[0] * b[0];
226 let c5 = F::dot_product::<4>(&[a[1], a[2], a[3], a[4]], &[b[4], b[3], b[2], b[1]]);
227 let c8 = a[4] * b[4];
228
229 c0 + c5 - c8
230 }
231}
232
233impl<F, A> PrimeCharacteristicRing for QuinticTrinomialExtensionField<F, A>
234where
235 F: QuinticTrinomialExtendable,
236 A: QuinticExtendableAlgebra<F>,
237{
238 type PrimeSubfield = <A as PrimeCharacteristicRing>::PrimeSubfield;
239
240 const ZERO: Self = Self::new([A::ZERO; 5]);
241 const ONE: Self = Self::new(field_to_array(A::ONE));
242 const TWO: Self = Self::new(field_to_array(A::TWO));
243 const NEG_ONE: Self = Self::new(field_to_array(A::NEG_ONE));
244
245 #[inline]
246 fn from_prime_subfield(f: Self::PrimeSubfield) -> Self {
247 <A as PrimeCharacteristicRing>::from_prime_subfield(f).into()
248 }
249
250 #[inline]
251 fn halve(&self) -> Self {
252 Self::new(self.value.clone().map(|x| x.halve()))
253 }
254
255 #[inline(always)]
256 fn square(&self) -> Self {
257 let mut res = Self::default();
258 A::quintic_square(&self.value, &mut res.value);
259 res
260 }
261
262 #[inline]
263 fn mul_2exp_u64(&self, exp: u64) -> Self {
264 Self::new(self.value.clone().map(|x| x.mul_2exp_u64(exp)))
265 }
266
267 #[inline]
268 fn div_2exp_u64(&self, exp: u64) -> Self {
269 Self::new(self.value.clone().map(|x| x.div_2exp_u64(exp)))
270 }
271
272 #[inline]
273 fn zero_vec(len: usize) -> Vec<Self> {
274 unsafe { reconstitute_from_base(A::zero_vec(len * 5)) }
276 }
277}
278
279impl<F: QuinticTrinomialExtendable> Algebra<F> for QuinticTrinomialExtensionField<F> {}
280
281impl<F: QuinticTrinomialExtendable> RawDataSerializable for QuinticTrinomialExtensionField<F> {
282 const NUM_BYTES: usize = F::NUM_BYTES * 5;
283
284 #[inline]
285 fn into_bytes(self) -> impl IntoIterator<Item = u8> {
286 self.value.into_iter().flat_map(|x| x.into_bytes())
287 }
288
289 #[inline]
290 fn into_byte_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u8> {
291 F::into_byte_stream(input.into_iter().flat_map(|x| x.value))
292 }
293
294 #[inline]
295 fn into_u32_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u32> {
296 F::into_u32_stream(input.into_iter().flat_map(|x| x.value))
297 }
298
299 #[inline]
300 fn into_u64_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u64> {
301 F::into_u64_stream(input.into_iter().flat_map(|x| x.value))
302 }
303
304 #[inline]
305 fn into_parallel_byte_streams<const N: usize>(
306 input: impl IntoIterator<Item = [Self; N]>,
307 ) -> impl IntoIterator<Item = [u8; N]> {
308 F::into_parallel_byte_streams(
309 input
310 .into_iter()
311 .flat_map(|x| (0..5).map(move |i| array::from_fn(|j| x[j].value[i]))),
312 )
313 }
314
315 #[inline]
316 fn into_parallel_u32_streams<const N: usize>(
317 input: impl IntoIterator<Item = [Self; N]>,
318 ) -> impl IntoIterator<Item = [u32; N]> {
319 F::into_parallel_u32_streams(
320 input
321 .into_iter()
322 .flat_map(|x| (0..5).map(move |i| array::from_fn(|j| x[j].value[i]))),
323 )
324 }
325
326 #[inline]
327 fn into_parallel_u64_streams<const N: usize>(
328 input: impl IntoIterator<Item = [Self; N]>,
329 ) -> impl IntoIterator<Item = [u64; N]> {
330 F::into_parallel_u64_streams(
331 input
332 .into_iter()
333 .flat_map(|x| (0..5).map(move |i| array::from_fn(|j| x[j].value[i]))),
334 )
335 }
336}
337
338impl<F: QuinticTrinomialExtendable> Field for QuinticTrinomialExtensionField<F> {
339 type Packing = Self;
340
341 const GENERATOR: Self = Self::new(F::EXT_GENERATOR);
342
343 fn try_inverse(&self) -> Option<Self> {
344 if self.is_zero() {
345 return None;
346 }
347 Some(self.pseudo_inv())
348 }
349
350 #[inline]
351 fn add_slices(slice_1: &mut [Self], slice_2: &[Self]) {
352 unsafe {
355 let base_slice_1 = as_base_slice_mut(slice_1);
356 let base_slice_2 = as_base_slice(slice_2);
357 F::add_slices(base_slice_1, base_slice_2);
358 }
359 }
360
361 #[inline]
362 fn order() -> BigUint {
363 F::order().pow(5)
364 }
365}
366
367impl<F: QuinticTrinomialExtendable> Display for QuinticTrinomialExtensionField<F> {
368 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
369 if self.is_zero() {
370 write!(f, "0")
371 } else {
372 let str = self
373 .value
374 .iter()
375 .enumerate()
376 .filter(|(_, x)| !x.is_zero())
377 .map(|(i, x)| match (i, x.is_one()) {
378 (0, _) => format!("{x}"),
379 (1, true) => "X".to_string(),
380 (1, false) => format!("{x} X"),
381 (_, true) => format!("X^{i}"),
382 (_, false) => format!("{x} X^{i}"),
383 })
384 .join(" + ");
385 write!(f, "{str}")
386 }
387 }
388}
389
390impl<F, A> Neg for QuinticTrinomialExtensionField<F, A>
391where
392 F: QuinticTrinomialExtendable,
393 A: Algebra<F>,
394{
395 type Output = Self;
396
397 #[inline]
398 fn neg(self) -> Self {
399 Self::new(self.value.map(A::neg))
400 }
401}
402
403impl<F, A> Add for QuinticTrinomialExtensionField<F, A>
404where
405 F: QuinticTrinomialExtendable,
406 A: QuinticExtendableAlgebra<F>,
407{
408 type Output = Self;
409
410 #[inline]
411 fn add(self, rhs: Self) -> Self {
412 Self::new(A::quintic_add(&self.value, &rhs.value))
413 }
414}
415
416impl<F, A> Add<A> for QuinticTrinomialExtensionField<F, A>
417where
418 F: QuinticTrinomialExtendable,
419 A: Algebra<F>,
420{
421 type Output = Self;
422
423 #[inline]
424 fn add(mut self, rhs: A) -> Self {
425 self.value[0] += rhs;
426 self
427 }
428}
429
430impl<F, A> AddAssign for QuinticTrinomialExtensionField<F, A>
431where
432 F: QuinticTrinomialExtendable,
433 A: QuinticExtendableAlgebra<F>,
434{
435 #[inline]
436 fn add_assign(&mut self, rhs: Self) {
437 self.value = A::quintic_add(&self.value, &rhs.value);
438 }
439}
440
441impl<F, A> AddAssign<A> for QuinticTrinomialExtensionField<F, A>
442where
443 F: QuinticTrinomialExtendable,
444 A: Algebra<F>,
445{
446 #[inline]
447 fn add_assign(&mut self, rhs: A) {
448 self.value[0] += rhs;
449 }
450}
451
452impl<F, A> Sum for QuinticTrinomialExtensionField<F, A>
453where
454 F: QuinticTrinomialExtendable,
455 A: QuinticExtendableAlgebra<F>,
456{
457 #[inline]
458 fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
459 iter.reduce(|acc, x| acc + x).unwrap_or(Self::ZERO)
460 }
461}
462
463impl<F, A> Sub for QuinticTrinomialExtensionField<F, A>
464where
465 F: QuinticTrinomialExtendable,
466 A: QuinticExtendableAlgebra<F>,
467{
468 type Output = Self;
469
470 #[inline]
471 fn sub(self, rhs: Self) -> Self {
472 Self::new(A::quintic_sub(&self.value, &rhs.value))
473 }
474}
475
476impl<F, A> Sub<A> for QuinticTrinomialExtensionField<F, A>
477where
478 F: QuinticTrinomialExtendable,
479 A: Algebra<F>,
480{
481 type Output = Self;
482
483 #[inline]
484 fn sub(self, rhs: A) -> Self {
485 let mut res = self.value;
486 res[0] -= rhs;
487 Self::new(res)
488 }
489}
490
491impl<F, A> SubAssign for QuinticTrinomialExtensionField<F, A>
492where
493 F: QuinticTrinomialExtendable,
494 A: QuinticExtendableAlgebra<F>,
495{
496 #[inline]
497 fn sub_assign(&mut self, rhs: Self) {
498 self.value = A::quintic_sub(&self.value, &rhs.value);
499 }
500}
501
502impl<F, A> SubAssign<A> for QuinticTrinomialExtensionField<F, A>
503where
504 F: QuinticTrinomialExtendable,
505 A: Algebra<F>,
506{
507 #[inline]
508 fn sub_assign(&mut self, rhs: A) {
509 self.value[0] -= rhs;
510 }
511}
512
513impl<F, A> Mul for QuinticTrinomialExtensionField<F, A>
514where
515 F: QuinticTrinomialExtendable,
516 A: QuinticExtendableAlgebra<F>,
517{
518 type Output = Self;
519
520 #[inline]
521 fn mul(self, rhs: Self) -> Self {
522 let mut res = Self::default();
523 A::quintic_mul(&self.value, &rhs.value, &mut res.value);
524 res
525 }
526}
527
528impl<F, A> Mul<A> for QuinticTrinomialExtensionField<F, A>
529where
530 F: QuinticTrinomialExtendable,
531 A: QuinticExtendableAlgebra<F>,
532{
533 type Output = Self;
534
535 #[inline]
536 fn mul(self, rhs: A) -> Self {
537 Self::new(A::quintic_base_mul(self.value, rhs))
538 }
539}
540
541impl<F, A> MulAssign for QuinticTrinomialExtensionField<F, A>
542where
543 F: QuinticTrinomialExtendable,
544 A: QuinticExtendableAlgebra<F>,
545{
546 #[inline]
547 fn mul_assign(&mut self, rhs: Self) {
548 *self = self.clone() * rhs;
549 }
550}
551
552impl<F, A> MulAssign<A> for QuinticTrinomialExtensionField<F, A>
553where
554 F: QuinticTrinomialExtendable,
555 A: QuinticExtendableAlgebra<F>,
556{
557 #[inline]
558 fn mul_assign(&mut self, rhs: A) {
559 *self = self.clone() * rhs;
560 }
561}
562
563impl<F, A> Product for QuinticTrinomialExtensionField<F, A>
564where
565 F: QuinticTrinomialExtendable,
566 A: QuinticExtendableAlgebra<F>,
567{
568 #[inline]
569 fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
570 iter.reduce(|acc, x| acc * x).unwrap_or(Self::ONE)
571 }
572}
573
574impl<F> Div for QuinticTrinomialExtensionField<F>
575where
576 F: QuinticTrinomialExtendable,
577{
578 type Output = Self;
579
580 #[allow(clippy::suspicious_arithmetic_impl)]
581 #[inline]
582 fn div(self, rhs: Self) -> Self::Output {
583 self * rhs.inverse()
584 }
585}
586
587impl<F> DivAssign for QuinticTrinomialExtensionField<F>
588where
589 F: QuinticTrinomialExtendable,
590{
591 #[inline]
592 fn div_assign(&mut self, rhs: Self) {
593 *self = *self / rhs;
594 }
595}
596
597impl<F: QuinticTrinomialExtendable> Distribution<QuinticTrinomialExtensionField<F>>
598 for StandardUniform
599where
600 Self: Distribution<F>,
601{
602 #[inline]
603 fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> QuinticTrinomialExtensionField<F> {
604 QuinticTrinomialExtensionField::new(array::from_fn(|_| self.sample(rng)))
605 }
606}
607
608impl<F: QuinticTrinomialExtendable + HasTwoAdicQuinticExtension> TwoAdicField
609 for QuinticTrinomialExtensionField<F>
610{
611 const TWO_ADICITY: usize = F::EXT_TWO_ADICITY;
612
613 #[inline]
614 fn two_adic_generator(bits: usize) -> Self {
615 Self::new(F::ext_two_adic_generator(bits))
616 }
617}
618
619#[inline]
621pub fn trinomial_quintic_mul<R: PrimeCharacteristicRing>(a: &[R; 5], b: &[R; 5], res: &mut [R; 5]) {
622 let c0 = a[0].clone() * b[0].clone();
624 let c1 = R::dot_product::<2>(&[a[0].clone(), a[1].clone()], &[b[1].clone(), b[0].clone()]);
625 let c2 = R::dot_product::<3>(
626 &[a[0].clone(), a[1].clone(), a[2].clone()],
627 &[b[2].clone(), b[1].clone(), b[0].clone()],
628 );
629 let c3 = R::dot_product::<4>(
630 &[a[0].clone(), a[1].clone(), a[2].clone(), a[3].clone()],
631 &[b[3].clone(), b[2].clone(), b[1].clone(), b[0].clone()],
632 );
633 let c4 = R::dot_product::<5>(
634 a,
635 &[
636 b[4].clone(),
637 b[3].clone(),
638 b[2].clone(),
639 b[1].clone(),
640 b[0].clone(),
641 ],
642 );
643
644 let c5 = R::dot_product::<4>(
646 &[a[1].clone(), a[2].clone(), a[3].clone(), a[4].clone()],
647 &[b[4].clone(), b[3].clone(), b[2].clone(), b[1].clone()],
648 );
649 let c6 = R::dot_product::<3>(
650 &[a[2].clone(), a[3].clone(), a[4].clone()],
651 &[b[4].clone(), b[3].clone(), b[2].clone()],
652 );
653 let c7 = R::dot_product::<2>(&[a[3].clone(), a[4].clone()], &[b[4].clone(), b[3].clone()]);
654 let c8 = a[4].clone() * b[4].clone();
655
656 let c5_minus_c8 = c5 - c8.clone();
658
659 res[0] = c0 + c5_minus_c8.clone();
661 res[1] = c1 + c6.clone();
662 res[2] = c2 - c5_minus_c8 + c7.clone();
663 res[3] = c3 - c6 + c8;
664 res[4] = c4 - c7;
665}
666
667#[inline]
669pub(super) fn quintic_square<R: PrimeCharacteristicRing>(a: &[R; 5], res: &mut [R; 5]) {
670 let a0_2 = a[0].double();
672 let a1_2 = a[1].double();
673 let a2_2 = a[2].double();
674 let a3_2 = a[3].double();
675
676 let c0 = a[0].square();
679 let c1 = a0_2.clone() * a[1].clone();
681 let c2 = R::dot_product::<2>(&[a0_2.clone(), a[1].clone()], &[a[2].clone(), a[1].clone()]);
683 let c3 = R::dot_product::<2>(&[a0_2.clone(), a1_2.clone()], &[a[3].clone(), a[2].clone()]);
685 let c4 = R::dot_product::<3>(
687 &[a0_2, a1_2.clone(), a[2].clone()],
688 &[a[4].clone(), a[3].clone(), a[2].clone()],
689 );
690 let c5 = R::dot_product::<2>(&[a1_2, a2_2.clone()], &[a[4].clone(), a[3].clone()]);
692 let c6 = R::dot_product::<2>(&[a2_2, a[3].clone()], &[a[4].clone(), a[3].clone()]);
694 let c7 = a3_2 * a[4].clone();
696 let c8 = a[4].square();
698
699 let c5_minus_c8 = c5 - c8.clone();
701
702 res[0] = c0 + c5_minus_c8.clone();
704 res[1] = c1 + c6.clone();
705 res[2] = c2 - c5_minus_c8 + c7.clone();
706 res[3] = c3 - c6 + c8;
707 res[4] = c4 - c7;
708}