sp1_stark/
septic_extension.rs

1//! A septic extension with an irreducible polynomial `z^7 - 2z - 5`.
2use num_bigint::BigUint;
3use num_traits::One;
4use p3_field::{
5    AbstractExtensionField, AbstractField, ExtensionField, Field, Packable, PrimeField32,
6};
7use serde::{Deserialize, Serialize};
8use std::{
9    array,
10    fmt::Display,
11    iter::{Product, Sum},
12    ops::{Add, AddAssign, Div, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign},
13};
14
15use crate::air::{SP1AirBuilder, SepticExtensionAirBuilder};
16
17/// A septic extension with an irreducible polynomial `z^7 - 2z - 5`.
18///
19/// The field can be constructed as `F_{p^7} = F_p[z]/(z^7 - 2z - 5)`.
20#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
21#[repr(C)]
22pub struct SepticExtension<F>(pub [F; 7]);
23
24impl<F: AbstractField> AbstractField for SepticExtension<F> {
25    type F = SepticExtension<F::F>;
26
27    fn zero() -> Self {
28        SepticExtension([
29            F::zero(),
30            F::zero(),
31            F::zero(),
32            F::zero(),
33            F::zero(),
34            F::zero(),
35            F::zero(),
36        ])
37    }
38
39    fn one() -> Self {
40        SepticExtension([
41            F::one(),
42            F::zero(),
43            F::zero(),
44            F::zero(),
45            F::zero(),
46            F::zero(),
47            F::zero(),
48        ])
49    }
50
51    fn two() -> Self {
52        SepticExtension([
53            F::two(),
54            F::zero(),
55            F::zero(),
56            F::zero(),
57            F::zero(),
58            F::zero(),
59            F::zero(),
60        ])
61    }
62
63    fn neg_one() -> Self {
64        SepticExtension([
65            F::neg_one(),
66            F::zero(),
67            F::zero(),
68            F::zero(),
69            F::zero(),
70            F::zero(),
71            F::zero(),
72        ])
73    }
74
75    fn from_f(f: Self::F) -> Self {
76        SepticExtension([
77            F::from_f(f.0[0]),
78            F::from_f(f.0[1]),
79            F::from_f(f.0[2]),
80            F::from_f(f.0[3]),
81            F::from_f(f.0[4]),
82            F::from_f(f.0[5]),
83            F::from_f(f.0[6]),
84        ])
85    }
86
87    fn from_bool(b: bool) -> Self {
88        SepticExtension([
89            F::from_bool(b),
90            F::zero(),
91            F::zero(),
92            F::zero(),
93            F::zero(),
94            F::zero(),
95            F::zero(),
96        ])
97    }
98
99    fn from_canonical_u8(n: u8) -> Self {
100        SepticExtension([
101            F::from_canonical_u8(n),
102            F::zero(),
103            F::zero(),
104            F::zero(),
105            F::zero(),
106            F::zero(),
107            F::zero(),
108        ])
109    }
110
111    fn from_canonical_u16(n: u16) -> Self {
112        SepticExtension([
113            F::from_canonical_u16(n),
114            F::zero(),
115            F::zero(),
116            F::zero(),
117            F::zero(),
118            F::zero(),
119            F::zero(),
120        ])
121    }
122
123    fn from_canonical_u32(n: u32) -> Self {
124        SepticExtension([
125            F::from_canonical_u32(n),
126            F::zero(),
127            F::zero(),
128            F::zero(),
129            F::zero(),
130            F::zero(),
131            F::zero(),
132        ])
133    }
134
135    fn from_canonical_u64(n: u64) -> Self {
136        SepticExtension([
137            F::from_canonical_u64(n),
138            F::zero(),
139            F::zero(),
140            F::zero(),
141            F::zero(),
142            F::zero(),
143            F::zero(),
144        ])
145    }
146
147    fn from_canonical_usize(n: usize) -> Self {
148        SepticExtension([
149            F::from_canonical_usize(n),
150            F::zero(),
151            F::zero(),
152            F::zero(),
153            F::zero(),
154            F::zero(),
155            F::zero(),
156        ])
157    }
158
159    fn from_wrapped_u32(n: u32) -> Self {
160        SepticExtension([
161            F::from_wrapped_u32(n),
162            F::zero(),
163            F::zero(),
164            F::zero(),
165            F::zero(),
166            F::zero(),
167            F::zero(),
168        ])
169    }
170
171    fn from_wrapped_u64(n: u64) -> Self {
172        SepticExtension([
173            F::from_wrapped_u64(n),
174            F::zero(),
175            F::zero(),
176            F::zero(),
177            F::zero(),
178            F::zero(),
179            F::zero(),
180        ])
181    }
182
183    fn generator() -> Self {
184        SepticExtension([F::two(), F::one(), F::zero(), F::zero(), F::zero(), F::zero(), F::zero()])
185    }
186}
187
188impl<F: Field> Field for SepticExtension<F> {
189    type Packing = Self;
190
191    fn try_inverse(&self) -> Option<Self> {
192        if self.is_zero() {
193            return None;
194        }
195        Some(self.inv())
196    }
197
198    fn order() -> BigUint {
199        F::order().pow(7)
200    }
201}
202
203impl<F: AbstractField> AbstractExtensionField<F> for SepticExtension<F> {
204    const D: usize = 7;
205
206    fn from_base(b: F) -> Self {
207        SepticExtension([b, F::zero(), F::zero(), F::zero(), F::zero(), F::zero(), F::zero()])
208    }
209
210    fn from_base_slice(bs: &[F]) -> Self {
211        SepticExtension([
212            bs[0].clone(),
213            bs[1].clone(),
214            bs[2].clone(),
215            bs[3].clone(),
216            bs[4].clone(),
217            bs[5].clone(),
218            bs[6].clone(),
219        ])
220    }
221
222    fn from_base_fn<G: FnMut(usize) -> F>(f: G) -> Self {
223        Self(array::from_fn(f))
224    }
225
226    fn as_base_slice(&self) -> &[F] {
227        self.0.as_slice()
228    }
229}
230
231impl<F: Field> ExtensionField<F> for SepticExtension<F> {
232    type ExtensionPacking = SepticExtension<F::Packing>;
233}
234
235impl<F: Field> Packable for SepticExtension<F> {}
236
237impl<F: AbstractField> Add for SepticExtension<F> {
238    type Output = Self;
239
240    fn add(self, rhs: Self) -> Self::Output {
241        let mut res = self.0;
242        for (r, rhs_val) in res.iter_mut().zip(rhs.0) {
243            *r = (*r).clone() + rhs_val;
244        }
245        Self(res)
246    }
247}
248
249impl<F: AbstractField> AddAssign for SepticExtension<F> {
250    fn add_assign(&mut self, rhs: Self) {
251        self.0[0] += rhs.0[0].clone();
252        self.0[1] += rhs.0[1].clone();
253        self.0[2] += rhs.0[2].clone();
254        self.0[3] += rhs.0[3].clone();
255        self.0[4] += rhs.0[4].clone();
256        self.0[5] += rhs.0[5].clone();
257        self.0[6] += rhs.0[6].clone();
258    }
259}
260
261impl<F: AbstractField> Sub for SepticExtension<F> {
262    type Output = Self;
263
264    fn sub(self, rhs: Self) -> Self::Output {
265        let mut res = self.0;
266        for (r, rhs_val) in res.iter_mut().zip(rhs.0) {
267            *r = (*r).clone() - rhs_val;
268        }
269        Self(res)
270    }
271}
272
273impl<F: AbstractField> SubAssign for SepticExtension<F> {
274    fn sub_assign(&mut self, rhs: Self) {
275        self.0[0] -= rhs.0[0].clone();
276        self.0[1] -= rhs.0[1].clone();
277        self.0[2] -= rhs.0[2].clone();
278        self.0[3] -= rhs.0[3].clone();
279        self.0[4] -= rhs.0[4].clone();
280        self.0[5] -= rhs.0[5].clone();
281        self.0[6] -= rhs.0[6].clone();
282    }
283}
284
285impl<F: AbstractField> Neg for SepticExtension<F> {
286    type Output = Self;
287
288    fn neg(self) -> Self::Output {
289        let mut res = self.0;
290        for r in res.iter_mut() {
291            *r = -r.clone();
292        }
293        Self(res)
294    }
295}
296
297impl<F: AbstractField> Mul for SepticExtension<F> {
298    type Output = Self;
299
300    fn mul(self, rhs: Self) -> Self::Output {
301        let mut res: [F; 13] = core::array::from_fn(|_| F::zero());
302        for i in 0..7 {
303            for j in 0..7 {
304                res[i + j] = res[i + j].clone() + self.0[i].clone() * rhs.0[j].clone();
305            }
306        }
307        let mut ret: [F; 7] = core::array::from_fn(|i| res[i].clone());
308        for i in 7..13 {
309            ret[i - 7] = ret[i - 7].clone() + res[i].clone() * F::from_canonical_u32(5);
310            ret[i - 6] = ret[i - 6].clone() + res[i].clone() * F::from_canonical_u32(2);
311        }
312        Self(ret)
313    }
314}
315
316impl<F: AbstractField> MulAssign for SepticExtension<F> {
317    fn mul_assign(&mut self, rhs: Self) {
318        let res = self.clone() * rhs;
319        *self = res;
320    }
321}
322
323impl<F: AbstractField> Product for SepticExtension<F> {
324    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
325        let one = Self::one();
326        iter.fold(one, |acc, x| acc * x)
327    }
328}
329
330impl<F: AbstractField> Sum for SepticExtension<F> {
331    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
332        let zero = Self::zero();
333        iter.fold(zero, |acc, x| acc + x)
334    }
335}
336
337impl<F: AbstractField> From<F> for SepticExtension<F> {
338    fn from(f: F) -> Self {
339        SepticExtension([f, F::zero(), F::zero(), F::zero(), F::zero(), F::zero(), F::zero()])
340    }
341}
342
343impl<F: AbstractField> Add<F> for SepticExtension<F> {
344    type Output = Self;
345
346    fn add(self, rhs: F) -> Self::Output {
347        SepticExtension([
348            self.0[0].clone() + rhs,
349            self.0[1].clone(),
350            self.0[2].clone(),
351            self.0[3].clone(),
352            self.0[4].clone(),
353            self.0[5].clone(),
354            self.0[6].clone(),
355        ])
356    }
357}
358
359impl<F: AbstractField> AddAssign<F> for SepticExtension<F> {
360    fn add_assign(&mut self, rhs: F) {
361        self.0[0] += rhs;
362    }
363}
364
365impl<F: AbstractField> Sub<F> for SepticExtension<F> {
366    type Output = Self;
367
368    fn sub(self, rhs: F) -> Self::Output {
369        self + (-rhs)
370    }
371}
372
373impl<F: AbstractField> SubAssign<F> for SepticExtension<F> {
374    fn sub_assign(&mut self, rhs: F) {
375        self.0[0] -= rhs;
376    }
377}
378
379impl<F: AbstractField> Mul<F> for SepticExtension<F> {
380    type Output = Self;
381
382    fn mul(self, rhs: F) -> Self::Output {
383        SepticExtension([
384            self.0[0].clone() * rhs.clone(),
385            self.0[1].clone() * rhs.clone(),
386            self.0[2].clone() * rhs.clone(),
387            self.0[3].clone() * rhs.clone(),
388            self.0[4].clone() * rhs.clone(),
389            self.0[5].clone() * rhs.clone(),
390            self.0[6].clone() * rhs.clone(),
391        ])
392    }
393}
394
395impl<F: AbstractField> MulAssign<F> for SepticExtension<F> {
396    fn mul_assign(&mut self, rhs: F) {
397        for i in 0..7 {
398            self.0[i] *= rhs.clone();
399        }
400    }
401}
402
403impl<F: Field> Div for SepticExtension<F> {
404    type Output = Self;
405
406    #[allow(clippy::suspicious_arithmetic_impl)]
407    fn div(self, rhs: Self) -> Self::Output {
408        self * rhs.inverse()
409    }
410}
411
412impl<F: AbstractField> Display for SepticExtension<F> {
413    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
414        write!(f, "{:?}", self.0)
415    }
416}
417
418impl<F: Field> SepticExtension<F> {
419    /// Returns the value of z^{index * p} in the [`SepticExtension`] field.
420    fn z_pow_p(index: u32) -> Self {
421        // The constants written below are specifically for the BabyBear field.
422        debug_assert_eq!(F::order(), BigUint::from(2013265921u32));
423        if index == 0 {
424            return Self::one();
425        }
426        if index == 1 {
427            return SepticExtension([
428                F::from_canonical_u32(954599710),
429                F::from_canonical_u32(1359279693),
430                F::from_canonical_u32(566669999),
431                F::from_canonical_u32(1982781815),
432                F::from_canonical_u32(1735718361),
433                F::from_canonical_u32(1174868538),
434                F::from_canonical_u32(1120871770),
435            ]);
436        }
437        if index == 2 {
438            return SepticExtension([
439                F::from_canonical_u32(862825265),
440                F::from_canonical_u32(597046311),
441                F::from_canonical_u32(978840770),
442                F::from_canonical_u32(1790138282),
443                F::from_canonical_u32(1044777201),
444                F::from_canonical_u32(835869808),
445                F::from_canonical_u32(1342179023),
446            ]);
447        }
448        if index == 3 {
449            return SepticExtension([
450                F::from_canonical_u32(596273169),
451                F::from_canonical_u32(658837454),
452                F::from_canonical_u32(1515468261),
453                F::from_canonical_u32(367059247),
454                F::from_canonical_u32(781278880),
455                F::from_canonical_u32(1544222616),
456                F::from_canonical_u32(155490465),
457            ]);
458        }
459        if index == 4 {
460            return SepticExtension([
461                F::from_canonical_u32(557608863),
462                F::from_canonical_u32(1173670028),
463                F::from_canonical_u32(1749546888),
464                F::from_canonical_u32(1086464137),
465                F::from_canonical_u32(803900099),
466                F::from_canonical_u32(1288818584),
467                F::from_canonical_u32(1184677604),
468            ]);
469        }
470        if index == 5 {
471            return SepticExtension([
472                F::from_canonical_u32(763416381),
473                F::from_canonical_u32(1252567168),
474                F::from_canonical_u32(628856225),
475                F::from_canonical_u32(1771903394),
476                F::from_canonical_u32(650712211),
477                F::from_canonical_u32(19417363),
478                F::from_canonical_u32(57990258),
479            ]);
480        }
481        if index == 6 {
482            return SepticExtension([
483                F::from_canonical_u32(1734711039),
484                F::from_canonical_u32(1749813853),
485                F::from_canonical_u32(1227235221),
486                F::from_canonical_u32(1707730636),
487                F::from_canonical_u32(424560395),
488                F::from_canonical_u32(1007029514),
489                F::from_canonical_u32(498034669),
490            ]);
491        }
492        unreachable!();
493    }
494
495    /// Returns the value of z^{index * p^2} in the [`SepticExtension`] field.
496    fn z_pow_p2(index: u32) -> Self {
497        // The constants written below are specifically for the BabyBear field.
498        debug_assert_eq!(F::order(), BigUint::from(2013265921u32));
499        if index == 0 {
500            return Self::one();
501        }
502        if index == 1 {
503            return SepticExtension([
504                F::from_canonical_u32(1013489358),
505                F::from_canonical_u32(1619071628),
506                F::from_canonical_u32(304593143),
507                F::from_canonical_u32(1949397349),
508                F::from_canonical_u32(1564307636),
509                F::from_canonical_u32(327761151),
510                F::from_canonical_u32(415430835),
511            ]);
512        }
513        if index == 2 {
514            return SepticExtension([
515                F::from_canonical_u32(209824426),
516                F::from_canonical_u32(1313900768),
517                F::from_canonical_u32(38410482),
518                F::from_canonical_u32(256593180),
519                F::from_canonical_u32(1708830551),
520                F::from_canonical_u32(1244995038),
521                F::from_canonical_u32(1555324019),
522            ]);
523        }
524        if index == 3 {
525            return SepticExtension([
526                F::from_canonical_u32(1475628651),
527                F::from_canonical_u32(777565847),
528                F::from_canonical_u32(704492386),
529                F::from_canonical_u32(1218528120),
530                F::from_canonical_u32(1245363405),
531                F::from_canonical_u32(475884575),
532                F::from_canonical_u32(649166061),
533            ]);
534        }
535        if index == 4 {
536            return SepticExtension([
537                F::from_canonical_u32(550038364),
538                F::from_canonical_u32(948935655),
539                F::from_canonical_u32(68722023),
540                F::from_canonical_u32(1251345762),
541                F::from_canonical_u32(1692456177),
542                F::from_canonical_u32(1177958698),
543                F::from_canonical_u32(350232928),
544            ]);
545        }
546        if index == 5 {
547            return SepticExtension([
548                F::from_canonical_u32(882720258),
549                F::from_canonical_u32(821925756),
550                F::from_canonical_u32(199955840),
551                F::from_canonical_u32(812002876),
552                F::from_canonical_u32(1484951277),
553                F::from_canonical_u32(1063138035),
554                F::from_canonical_u32(491712810),
555            ]);
556        }
557        if index == 6 {
558            return SepticExtension([
559                F::from_canonical_u32(738287111),
560                F::from_canonical_u32(1955364991),
561                F::from_canonical_u32(552724293),
562                F::from_canonical_u32(1175775744),
563                F::from_canonical_u32(341623997),
564                F::from_canonical_u32(1454022463),
565                F::from_canonical_u32(408193320),
566            ]);
567        }
568        unreachable!();
569    }
570
571    #[must_use]
572    fn frobenius(&self) -> Self {
573        let mut result = Self::zero();
574        result += self.0[0];
575        result += Self::z_pow_p(1) * self.0[1];
576        result += Self::z_pow_p(2) * self.0[2];
577        result += Self::z_pow_p(3) * self.0[3];
578        result += Self::z_pow_p(4) * self.0[4];
579        result += Self::z_pow_p(5) * self.0[5];
580        result += Self::z_pow_p(6) * self.0[6];
581        result
582    }
583
584    #[must_use]
585    fn double_frobenius(&self) -> Self {
586        let mut result = Self::zero();
587        result += self.0[0];
588        result += Self::z_pow_p2(1) * self.0[1];
589        result += Self::z_pow_p2(2) * self.0[2];
590        result += Self::z_pow_p2(3) * self.0[3];
591        result += Self::z_pow_p2(4) * self.0[4];
592        result += Self::z_pow_p2(5) * self.0[5];
593        result += Self::z_pow_p2(6) * self.0[6];
594        result
595    }
596
597    #[must_use]
598    fn pow_r_1(&self) -> Self {
599        let base = self.frobenius() * self.double_frobenius();
600        let base_p2 = base.double_frobenius();
601        let base_p4 = base_p2.double_frobenius();
602        base * base_p2 * base_p4
603    }
604
605    #[must_use]
606    fn inv(&self) -> Self {
607        let pow_r_1 = self.pow_r_1();
608        let pow_r = pow_r_1 * *self;
609        pow_r_1 * pow_r.0[0].inverse()
610    }
611
612    fn is_square(&self) -> (F, bool) {
613        let pow_r_1 = self.pow_r_1();
614        let pow_r = pow_r_1 * *self;
615        let exp = (F::order() - BigUint::one()) / BigUint::from(2u8);
616        let exp = exp.to_u64_digits()[0];
617
618        (pow_r.0[0], pow_r.0[0].exp_u64(exp) == F::one())
619    }
620
621    /// Computes the square root of the septic field extension element.
622    /// Returns None if the element is not a square, and Some(result) if it is a square.
623    pub fn sqrt(&self) -> Option<Self> {
624        let n = *self;
625
626        if n == Self::zero() || n == Self::one() {
627            return Some(n);
628        }
629
630        let (numerator, is_square) = n.is_square();
631
632        if !is_square {
633            return None;
634        }
635
636        let mut n_iter = n;
637        let mut n_power = n;
638        for i in 1..30 {
639            n_iter *= n_iter;
640            if i >= 26 {
641                n_power *= n_iter;
642            }
643        }
644
645        let mut n_frobenius = n_power.frobenius();
646        let mut denominator = n_frobenius;
647
648        n_frobenius = n_frobenius.double_frobenius();
649        denominator *= n_frobenius;
650        n_frobenius = n_frobenius.double_frobenius();
651        denominator *= n_frobenius;
652        denominator *= n;
653
654        let base = numerator.inverse();
655        let g = F::generator();
656        let mut a = F::one();
657        let mut nonresidue = F::one() - base;
658        let legendre_exp = (F::order() - BigUint::one()) / BigUint::from(2u8);
659
660        while nonresidue.exp_u64(legendre_exp.to_u64_digits()[0]) == F::one() {
661            a *= g;
662            nonresidue = a.square() - base;
663        }
664
665        let order = F::order();
666        let cipolla_pow = (&order + BigUint::one()) / BigUint::from(2u8);
667        let mut x = CipollaExtension::new(a, F::one());
668        x = x.pow(&cipolla_pow, nonresidue);
669
670        Some(denominator * x.real)
671    }
672}
673
674impl<F: PrimeField32> SepticExtension<F> {
675    /// Returns whether the extension field element viewed as an y-coordinate of a digest represents
676    /// a receive interaction.
677    pub fn is_receive(&self) -> bool {
678        1 <= self.0[6].as_canonical_u32() &&
679            self.0[6].as_canonical_u32() <= F::ORDER_U32.div_ceil(2)
680    }
681
682    /// Returns whether the extension field element viewed as an y-coordinate of a digest represents
683    /// a send interaction.
684    pub fn is_send(&self) -> bool {
685        F::ORDER_U32.div_ceil(2) <= self.0[6].as_canonical_u32() &&
686            self.0[6].as_canonical_u32() <= (F::ORDER_U32 - 1)
687    }
688
689    /// Returns whether the extension field element viewed as an y-coordinate of a digest cannot
690    /// represent anything.
691    pub fn is_exception(&self) -> bool {
692        self.0[6].as_canonical_u32() == 0
693    }
694}
695
696/// Extension field for Cipolla's algorithm, taken from <https://github.com/Plonky3/Plonky3/pull/439/files>.
697#[derive(Clone, Copy, Debug)]
698struct CipollaExtension<F: Field> {
699    real: F,
700    imag: F,
701}
702
703impl<F: Field> CipollaExtension<F> {
704    fn new(real: F, imag: F) -> Self {
705        Self { real, imag }
706    }
707
708    fn one() -> Self {
709        Self::new(F::one(), F::zero())
710    }
711
712    fn mul_ext(&self, other: Self, nonresidue: F) -> Self {
713        Self::new(
714            self.real * other.real + nonresidue * self.imag * other.imag,
715            self.real * other.imag + self.imag * other.real,
716        )
717    }
718
719    fn pow(&self, exp: &BigUint, nonresidue: F) -> Self {
720        let mut result = Self::one();
721        let mut base = *self;
722        let bits = exp.bits();
723
724        for i in 0..bits {
725            if exp.bit(i) {
726                result = result.mul_ext(base, nonresidue);
727            }
728            base = base.mul_ext(base, nonresidue);
729        }
730        result
731    }
732}
733
734/// A block of columns for septic extension.
735#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
736#[repr(C)]
737pub struct SepticBlock<T>(pub [T; 7]);
738
739impl<T> SepticBlock<T> {
740    /// Maps a `SepticBlock<T>` to `SepticBlock<U>` based on a map from `T` to `U`.
741    pub fn map<F, U>(self, f: F) -> SepticBlock<U>
742    where
743        F: FnMut(T) -> U,
744    {
745        SepticBlock(self.0.map(f))
746    }
747
748    /// A function similar to `core:array::from_fn`.
749    pub fn from_base_fn<G: FnMut(usize) -> T>(f: G) -> Self {
750        Self(array::from_fn(f))
751    }
752}
753
754impl<T: Clone> SepticBlock<T> {
755    /// Takes a `SepticBlock` into a `SepticExtension` of expressions.
756    pub fn as_extension<AB: SepticExtensionAirBuilder<Var = T>>(
757        &self,
758    ) -> SepticExtension<AB::Expr> {
759        let arr: [AB::Expr; 7] = self.0.clone().map(|x| AB::Expr::zero() + x);
760        SepticExtension(arr)
761    }
762
763    /// Takes a single expression into a `SepticExtension` of expressions.
764    pub fn as_extension_from_base<AB: SP1AirBuilder<Var = T>>(
765        &self,
766        base: AB::Expr,
767    ) -> SepticExtension<AB::Expr> {
768        let mut arr: [AB::Expr; 7] = self.0.clone().map(|_| AB::Expr::zero());
769        arr[0] = base;
770
771        SepticExtension(arr)
772    }
773}
774
775impl<T> From<[T; 7]> for SepticBlock<T> {
776    fn from(arr: [T; 7]) -> Self {
777        Self(arr)
778    }
779}
780
781impl<T: AbstractField> From<T> for SepticBlock<T> {
782    fn from(value: T) -> Self {
783        Self([value, T::zero(), T::zero(), T::zero(), T::zero(), T::zero(), T::zero()])
784    }
785}
786
787impl<T: Copy> From<&[T]> for SepticBlock<T> {
788    fn from(slice: &[T]) -> Self {
789        let arr: [T; 7] = slice.try_into().unwrap();
790        Self(arr)
791    }
792}
793
794impl<T, I> Index<I> for SepticBlock<T>
795where
796    [T]: Index<I>,
797{
798    type Output = <[T] as Index<I>>::Output;
799
800    #[inline]
801    fn index(&self, index: I) -> &Self::Output {
802        Index::index(&self.0, index)
803    }
804}
805
806impl<T, I> IndexMut<I> for SepticBlock<T>
807where
808    [T]: IndexMut<I>,
809{
810    #[inline]
811    fn index_mut(&mut self, index: I) -> &mut Self::Output {
812        IndexMut::index_mut(&mut self.0, index)
813    }
814}
815
816impl<T> IntoIterator for SepticBlock<T> {
817    type Item = T;
818    type IntoIter = std::array::IntoIter<T, 7>;
819
820    fn into_iter(self) -> Self::IntoIter {
821        self.0.into_iter()
822    }
823}
824
825#[cfg(test)]
826mod tests {
827    #![allow(clippy::print_stdout)]
828
829    use p3_baby_bear::BabyBear;
830
831    use super::*;
832
833    #[test]
834    fn test_mul() {
835        let a: SepticExtension<BabyBear> = SepticExtension::from_canonical_u32(1);
836        let b: SepticExtension<BabyBear> = SepticExtension::from_canonical_u32(2);
837        let c = a * b;
838        println!("{c}");
839    }
840
841    #[test]
842    fn test_inv() {
843        for i in 0..256 {
844            let a: SepticExtension<BabyBear> = SepticExtension([
845                BabyBear::from_canonical_u32(i + 3),
846                BabyBear::from_canonical_u32(2 * i + 6),
847                BabyBear::from_canonical_u32(5 * i + 17),
848                BabyBear::from_canonical_u32(6 * i + 91),
849                BabyBear::from_canonical_u32(8 * i + 37),
850                BabyBear::from_canonical_u32(11 * i + 35),
851                BabyBear::from_canonical_u32(14 * i + 33),
852            ]);
853            let b = a.inv();
854            assert_eq!(a * b, SepticExtension::<BabyBear>::one());
855        }
856    }
857
858    #[test]
859    fn test_legendre() {
860        let a: SepticExtension<BabyBear> = SepticExtension::generator();
861        let mut b = SepticExtension::<BabyBear>::one();
862        for i in 1..256 {
863            b *= a;
864            let (_, c) = b.is_square();
865            assert!(c == (i % 2 == 0));
866        }
867    }
868
869    #[test]
870    fn test_sqrt() {
871        for i in 0..256 {
872            let a: SepticExtension<BabyBear> = SepticExtension([
873                BabyBear::from_canonical_u32(i + 3),
874                BabyBear::from_canonical_u32(2 * i + 6),
875                BabyBear::from_canonical_u32(5 * i + 17),
876                BabyBear::from_canonical_u32(6 * i + 91),
877                BabyBear::from_canonical_u32(8 * i + 37),
878                BabyBear::from_canonical_u32(11 * i + 35),
879                BabyBear::from_canonical_u32(14 * i + 33),
880            ]);
881            let b = a * a;
882            let recovered_a = b.sqrt().unwrap();
883            assert_eq!(recovered_a * recovered_a, b);
884        }
885        let mut b = SepticExtension::<BabyBear>::one();
886        for i in 1..256 {
887            let a: SepticExtension<BabyBear> = SepticExtension::generator();
888            b *= a;
889            let c = b.sqrt();
890            if i % 2 == 1 {
891                assert!(c.is_none());
892            } else {
893                let c = c.unwrap();
894                assert_eq!(c * c, b);
895            }
896        }
897    }
898}