twenty_first/math/
x_field_element.rs

1use std::fmt::Display;
2use std::iter::Sum;
3use std::ops::Add;
4use std::ops::AddAssign;
5use std::ops::Div;
6use std::ops::Mul;
7use std::ops::MulAssign;
8use std::ops::Neg;
9use std::ops::Sub;
10use std::ops::SubAssign;
11
12use arbitrary::Arbitrary;
13use bfieldcodec_derive::BFieldCodec;
14use num_traits::ConstOne;
15use num_traits::ConstZero;
16use num_traits::One;
17use num_traits::Zero;
18use rand::distr::Distribution;
19use rand::distr::StandardUniform;
20use rand::Rng;
21use serde::Deserialize;
22use serde::Serialize;
23
24use super::digest::Digest;
25use crate::bfe_vec;
26use crate::error::TryFromXFieldElementError;
27use crate::math::b_field_element::BFieldElement;
28use crate::math::polynomial::Polynomial;
29use crate::math::traits::CyclicGroupGenerator;
30use crate::math::traits::FiniteField;
31use crate::math::traits::Inverse;
32use crate::math::traits::ModPowU32;
33use crate::math::traits::ModPowU64;
34use crate::math::traits::PrimitiveRootOfUnity;
35
36pub const EXTENSION_DEGREE: usize = 3;
37
38#[derive(
39    Debug, PartialEq, Eq, Copy, Clone, Hash, Serialize, Deserialize, BFieldCodec, Arbitrary,
40)]
41pub struct XFieldElement {
42    pub coefficients: [BFieldElement; EXTENSION_DEGREE],
43}
44
45/// Simplifies constructing [extension field element](XFieldElement)s.
46///
47/// The type [`XFieldElement`] must be in scope for this macro to work.
48/// See [`XFieldElement::from`] for supported types.
49///
50/// # Examples
51///
52/// ```
53/// # use twenty_first::prelude::*;
54/// let a = xfe!(1);
55/// let b = xfe!([2, 0, 5]);
56/// let c = xfe!([3, 0, 2 + 3]);
57/// assert_eq!(a + b, c);
58/// ```
59#[macro_export]
60macro_rules! xfe {
61    ($value:expr) => {
62        XFieldElement::from($value)
63    };
64}
65
66/// Simplifies constructing vectors of [extension field element][XFieldElement]s.
67///
68/// The type [`XFieldElement`] must be in scope for this macro to work. See also [`xfe!`].
69///
70/// # Examples
71///
72/// Vector of [constants](XFieldElement::new_const).
73///
74/// ```
75/// # use twenty_first::prelude::*;
76/// let a = xfe_vec![1, 2, 3];
77/// let b = vec![xfe!(1), xfe!(2), xfe!(3)];
78/// assert_eq!(a, b);
79/// ```
80///
81/// Vector of general [extension field element](XFieldElement)s.
82///
83/// ```
84/// # use twenty_first::prelude::*;
85/// let a = xfe_vec![[1, 0, 0], [0, 2, 0], [0, 0, 3]];
86/// let b = vec![xfe!([1, 0, 0]), xfe!([0, 2, 0]), xfe!([0, 0, 3])];
87/// assert_eq!(a, b);
88/// ```
89///
90/// Vector with the same [constant](XFieldElement::new_const) for every entry.
91///
92/// ```
93/// # use twenty_first::prelude::*;
94/// let a = xfe_vec![42; 15];
95/// let b = vec![xfe!(42); 15];
96/// assert_eq!(a, b);
97/// ```
98///
99/// Vector with the same general [extension field element](XFieldElement) for every entry.
100///
101/// ```
102/// # use twenty_first::prelude::*;
103/// let a = xfe_vec![[42, 43, 44]; 15];
104/// let b = vec![xfe!([42, 43, 44]); 15];
105/// assert_eq!(a, b);
106/// ```
107#[macro_export]
108macro_rules! xfe_vec {
109    ($x:expr; $n:expr) => {
110        vec![XFieldElement::from($x); $n]
111    };
112    ([$c0:expr, $c1:expr, $c2:expr]; $n:expr) => {
113        vec![XFieldElement::from([$c0, $c1, $c2]); $n]
114    };
115    ($($x:expr),* $(,)?) => {
116        vec![$(XFieldElement::from($x)),*]
117    };
118    ($([$c0:expr, $c1:expr, $c2:expr]),* $(,)?) => {
119        vec![$(XFieldElement::from([$c0, $c1, $c2])),*]
120    };
121}
122
123/// Simplifies constructing arrays of [extension field element][XFieldElement]s.
124///
125/// The type [`XFieldElement`] must be in scope for this macro to work. See also [`xfe!`].
126///
127/// # Examples
128///
129/// Array of [constants](XFieldElement::new_const).
130///
131/// ```
132/// # use twenty_first::prelude::*;
133/// let a = xfe_array![1, 2, 3];
134/// let b = [xfe!(1), xfe!(2), xfe!(3)];
135/// assert_eq!(a, b);
136/// ```
137///
138/// Array of general [extension field element](XFieldElement)s.
139///
140/// ```
141/// # use twenty_first::prelude::*;
142/// let a = xfe_array![[1, 0, 0], [0, 2, 0], [0, 0, 3]];
143/// let b = [xfe!([1, 0, 0]), xfe!([0, 2, 0]), xfe!([0, 0, 3])];
144/// assert_eq!(a, b);
145/// ```
146///
147/// Array with the same [constant](XFieldElement::new_const) for every entry.
148///
149/// ```
150/// # use twenty_first::prelude::*;
151/// let a = xfe_array![42; 15];
152/// let b = [xfe!(42); 15];
153/// assert_eq!(a, b);
154/// ```
155///
156/// Array with the same general [extension field element](XFieldElement) for every entry.
157///
158/// ```
159/// # use twenty_first::prelude::*;
160/// let a = xfe_array![[42, 43, 44]; 15];
161/// let b = [xfe!([42, 43, 44]); 15];
162/// assert_eq!(a, b);
163/// ```
164#[macro_export]
165macro_rules! xfe_array {
166    ($x:expr; $n:expr) => {
167        [XFieldElement::from($x); $n]
168    };
169    ([$c0:expr, $c1:expr, $c2:expr]; $n:expr) => {
170        [XFieldElement::from([$c0, $c1, $c2]); $n]
171    };
172    ($($x:expr),* $(,)?) => {
173        [$(XFieldElement::from($x)),*]
174    };
175    ($([$c0:expr, $c1:expr, $c2:expr]),* $(,)?) => {
176        [$(XFieldElement::from([$c0, $c1, $c2])),*]
177    };
178}
179
180impl From<XFieldElement> for Digest {
181    /// Interpret the `XFieldElement` as a [`Digest`]. No hashing is performed. This
182    /// interpretation can be useful for [`Tip5`](crate::math::tip5::Tip5) and,
183    /// by extension, allows building
184    /// [`MerkleTree`](crate::util_types::merkle_tree::MerkleTree)s directly from
185    /// `XFieldElement`s.
186    fn from(xfe: XFieldElement) -> Self {
187        let [c0, c1, c2] = xfe.coefficients;
188        Digest::new([c0, c1, c2, BFieldElement::ZERO, BFieldElement::ZERO])
189    }
190}
191
192impl TryFrom<Digest> for XFieldElement {
193    type Error = TryFromXFieldElementError;
194
195    fn try_from(digest: Digest) -> Result<Self, Self::Error> {
196        let [c0, c1, c2, zero_0, zero_1] = digest.values();
197        if zero_0 != BFieldElement::ZERO || zero_1 != BFieldElement::ZERO {
198            return Err(TryFromXFieldElementError::InvalidDigest);
199        }
200
201        Ok(Self::new([c0, c1, c2]))
202    }
203}
204
205impl Sum for XFieldElement {
206    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
207        iter.reduce(|a, b| a + b).unwrap_or(XFieldElement::ZERO)
208    }
209}
210
211impl<T> From<T> for XFieldElement
212where
213    T: Into<BFieldElement>,
214{
215    fn from(value: T) -> Self {
216        Self::new_const(value.into())
217    }
218}
219
220impl<T> From<[T; EXTENSION_DEGREE]> for XFieldElement
221where
222    T: Into<BFieldElement>,
223{
224    fn from(value: [T; EXTENSION_DEGREE]) -> Self {
225        Self::new(value.map(Into::into))
226    }
227}
228
229impl From<Polynomial<'_, BFieldElement>> for XFieldElement {
230    fn from(poly: Polynomial<'_, BFieldElement>) -> Self {
231        let (_, rem) = poly.naive_divide(&Self::shah_polynomial());
232        let mut xfe = [BFieldElement::ZERO; EXTENSION_DEGREE];
233
234        let Ok(rem_degree) = usize::try_from(rem.degree()) else {
235            return Self::ZERO;
236        };
237        xfe[..=rem_degree].copy_from_slice(&rem.coefficients()[..=rem_degree]);
238
239        XFieldElement::new(xfe)
240    }
241}
242
243impl TryFrom<&[BFieldElement]> for XFieldElement {
244    type Error = TryFromXFieldElementError;
245
246    fn try_from(value: &[BFieldElement]) -> Result<Self, Self::Error> {
247        value
248            .try_into()
249            .map(XFieldElement::new)
250            .map_err(|_| Self::Error::InvalidLength(value.len()))
251    }
252}
253
254impl TryFrom<Vec<BFieldElement>> for XFieldElement {
255    type Error = TryFromXFieldElementError;
256
257    fn try_from(value: Vec<BFieldElement>) -> Result<Self, Self::Error> {
258        XFieldElement::try_from(value.as_ref())
259    }
260}
261
262impl XFieldElement {
263    /// The quotient defining the [field extension](XFieldElement) over the
264    /// [base field](BFieldElement), namely x³ - x + 1.
265    #[inline]
266    pub fn shah_polynomial() -> Polynomial<'static, BFieldElement> {
267        Polynomial::new(bfe_vec![1, -1, 0, 1])
268    }
269
270    #[inline]
271    pub const fn new(coefficients: [BFieldElement; EXTENSION_DEGREE]) -> Self {
272        Self { coefficients }
273    }
274
275    #[inline]
276    pub const fn new_const(element: BFieldElement) -> Self {
277        let zero = BFieldElement::ZERO;
278        Self::new([element, zero, zero])
279    }
280
281    #[must_use]
282    pub fn inverse(&self) -> Self {
283        assert!(
284            !self.is_zero(),
285            "Cannot invert the zero element in the extension field."
286        );
287        let self_as_poly: Polynomial<BFieldElement> = self.to_owned().into();
288        let (_, a, _) = Polynomial::<BFieldElement>::xgcd(self_as_poly, Self::shah_polynomial());
289        a.into()
290    }
291
292    pub fn unlift(&self) -> Option<BFieldElement> {
293        if self.coefficients[1].is_zero() && self.coefficients[2].is_zero() {
294            Some(self.coefficients[0])
295        } else {
296            None
297        }
298    }
299
300    // `increment` and `decrement` are mainly used for testing purposes
301    pub fn increment(&mut self, index: usize) {
302        self.coefficients[index].increment();
303    }
304
305    pub fn decrement(&mut self, index: usize) {
306        self.coefficients[index].decrement();
307    }
308}
309
310impl Inverse for XFieldElement {
311    #[must_use]
312    fn inverse(&self) -> Self {
313        self.inverse()
314    }
315}
316
317impl PrimitiveRootOfUnity for XFieldElement {
318    fn primitive_root_of_unity(n: u64) -> Option<XFieldElement> {
319        let b_root = BFieldElement::primitive_root_of_unity(n);
320        b_root.map(XFieldElement::new_const)
321    }
322}
323
324impl Distribution<XFieldElement> for StandardUniform {
325    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> XFieldElement {
326        XFieldElement::new(rng.random())
327    }
328}
329
330impl CyclicGroupGenerator for XFieldElement {
331    fn get_cyclic_group_elements(&self, max: Option<usize>) -> Vec<Self> {
332        let mut val = *self;
333        let mut ret: Vec<Self> = vec![Self::one()];
334
335        loop {
336            ret.push(val);
337            val *= *self;
338            if val.is_one() || max.is_some() && ret.len() >= max.unwrap() {
339                break;
340            }
341        }
342        ret
343    }
344}
345
346impl Display for XFieldElement {
347    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
348        if let Some(bfe) = self.unlift() {
349            return write!(f, "{bfe}_xfe");
350        }
351
352        let [c0, c1, c2] = self.coefficients;
353        write!(f, "({c2:>020}·x² + {c1:>020}·x + {c0:>020})")
354    }
355}
356
357impl Zero for XFieldElement {
358    fn zero() -> Self {
359        Self::ZERO
360    }
361
362    fn is_zero(&self) -> bool {
363        self == &Self::ZERO
364    }
365}
366
367impl ConstZero for XFieldElement {
368    const ZERO: Self = Self::new([BFieldElement::ZERO; EXTENSION_DEGREE]);
369}
370
371impl One for XFieldElement {
372    fn one() -> Self {
373        Self::ONE
374    }
375
376    fn is_one(&self) -> bool {
377        self == &Self::ONE
378    }
379}
380
381impl ConstOne for XFieldElement {
382    const ONE: Self = Self::new([BFieldElement::ONE, BFieldElement::ZERO, BFieldElement::ZERO]);
383}
384
385impl FiniteField for XFieldElement {}
386
387impl Add<XFieldElement> for XFieldElement {
388    type Output = Self;
389
390    #[inline]
391    fn add(self, other: Self) -> Self {
392        let [s0, s1, s2] = self.coefficients;
393        let [o0, o1, o2] = other.coefficients;
394        let coefficients = [s0 + o0, s1 + o1, s2 + o2];
395        Self { coefficients }
396    }
397}
398
399impl Add<BFieldElement> for XFieldElement {
400    type Output = Self;
401
402    #[inline]
403    fn add(mut self, other: BFieldElement) -> Self {
404        self.coefficients[0] += other;
405        self
406    }
407}
408
409/// The `bfe + xfe -> xfe` instance belongs to BFieldElement.
410impl Add<XFieldElement> for BFieldElement {
411    type Output = XFieldElement;
412
413    #[inline]
414    fn add(self, mut other: XFieldElement) -> XFieldElement {
415        other.coefficients[0] += self;
416        other
417    }
418}
419
420impl Mul<XFieldElement> for XFieldElement {
421    type Output = Self;
422
423    #[inline]
424    fn mul(self, other: Self) -> Self {
425        // XField * XField means:
426        //
427        // (ax^2 + bx + c) * (dx^2 + ex + f)   (mod x^3 - x + 1)
428        //
429        // =   adx^4 + aex^3 + afx^2
430        //   + bdx^3 + bex^2 + bfx
431        //   + cdx^2 + cex   + cf
432        //
433        // = adx^4 + (ae + bd)x^3 + (af + be + cd)x^2 + (bf + ce)x + cf   (mod x^3 - x + 1)
434
435        let [c, b, a] = self.coefficients;
436        let [f, e, d] = other.coefficients;
437
438        let r0 = c * f - a * e - b * d;
439        let r1 = b * f + c * e - a * d + a * e + b * d;
440        let r2 = a * f + b * e + c * d + a * d;
441
442        Self::new([r0, r1, r2])
443    }
444}
445
446/// XField * BField means scalar multiplication of the
447/// BFieldElement onto each coefficient of the XField.
448impl Mul<BFieldElement> for XFieldElement {
449    type Output = Self;
450
451    #[inline]
452    fn mul(self, other: BFieldElement) -> Self {
453        let coefficients = self.coefficients.map(|c| c * other);
454        Self { coefficients }
455    }
456}
457
458impl Mul<XFieldElement> for BFieldElement {
459    type Output = XFieldElement;
460
461    #[inline]
462    fn mul(self, other: XFieldElement) -> XFieldElement {
463        let coefficients = other.coefficients.map(|c| c * self);
464        XFieldElement { coefficients }
465    }
466}
467
468impl Neg for XFieldElement {
469    type Output = Self;
470
471    #[inline]
472    fn neg(self) -> Self {
473        let coefficients = self.coefficients.map(Neg::neg);
474        Self { coefficients }
475    }
476}
477
478impl Sub<XFieldElement> for XFieldElement {
479    type Output = Self;
480
481    #[inline]
482    fn sub(self, other: Self) -> Self {
483        self + (-other)
484    }
485}
486
487impl Sub<BFieldElement> for XFieldElement {
488    type Output = Self;
489
490    #[inline]
491    fn sub(self, other: BFieldElement) -> Self {
492        self + (-other)
493    }
494}
495
496impl Sub<XFieldElement> for BFieldElement {
497    type Output = XFieldElement;
498
499    #[inline]
500    fn sub(self, other: XFieldElement) -> XFieldElement {
501        self + (-other)
502    }
503}
504
505impl AddAssign<XFieldElement> for XFieldElement {
506    #[inline]
507    fn add_assign(&mut self, rhs: Self) {
508        self.coefficients[0] += rhs.coefficients[0];
509        self.coefficients[1] += rhs.coefficients[1];
510        self.coefficients[2] += rhs.coefficients[2];
511    }
512}
513
514impl AddAssign<BFieldElement> for XFieldElement {
515    #[inline]
516    fn add_assign(&mut self, rhs: BFieldElement) {
517        self.coefficients[0] += rhs;
518    }
519}
520
521impl MulAssign<XFieldElement> for XFieldElement {
522    #[inline]
523    fn mul_assign(&mut self, rhs: Self) {
524        *self = *self * rhs;
525    }
526}
527
528impl MulAssign<BFieldElement> for XFieldElement {
529    #[inline]
530    fn mul_assign(&mut self, rhs: BFieldElement) {
531        *self = *self * rhs;
532    }
533}
534
535impl SubAssign<XFieldElement> for XFieldElement {
536    #[inline]
537    fn sub_assign(&mut self, rhs: Self) {
538        self.coefficients[0] -= rhs.coefficients[0];
539        self.coefficients[1] -= rhs.coefficients[1];
540        self.coefficients[2] -= rhs.coefficients[2];
541    }
542}
543
544impl SubAssign<BFieldElement> for XFieldElement {
545    #[inline]
546    fn sub_assign(&mut self, rhs: BFieldElement) {
547        self.coefficients[0] -= rhs;
548    }
549}
550
551impl Div for XFieldElement {
552    type Output = Self;
553
554    #[allow(clippy::suspicious_arithmetic_impl)]
555    fn div(self, other: Self) -> Self {
556        self * other.inverse()
557    }
558}
559
560impl ModPowU64 for XFieldElement {
561    #[inline]
562    fn mod_pow_u64(&self, exponent: u64) -> Self {
563        let mut x = *self;
564        let mut result = Self::one();
565        let mut i = exponent;
566
567        while i > 0 {
568            if i & 1 == 1 {
569                result *= x;
570            }
571
572            x *= x;
573            i >>= 1;
574        }
575
576        result
577    }
578}
579
580impl ModPowU32 for XFieldElement {
581    #[inline]
582    fn mod_pow_u32(&self, exp: u32) -> Self {
583        self.mod_pow_u64(exp as u64)
584    }
585}
586
587#[cfg(test)]
588mod tests {
589    use itertools::izip;
590    use itertools::Itertools;
591    use num_traits::ConstOne;
592    use proptest::collection::vec;
593    use proptest::prelude::*;
594    use proptest_arbitrary_interop::arb;
595    use test_strategy::proptest;
596
597    use crate::bfe;
598    use crate::math::b_field_element::*;
599    use crate::math::ntt::intt;
600    use crate::math::ntt::ntt;
601    use crate::math::other::random_elements;
602    use crate::math::x_field_element::*;
603
604    impl proptest::arbitrary::Arbitrary for XFieldElement {
605        type Parameters = ();
606
607        fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
608            arb().boxed()
609        }
610
611        type Strategy = BoxedStrategy<Self>;
612    }
613
614    #[test]
615    fn one_zero_test() {
616        let one = XFieldElement::one();
617        assert!(one.is_one());
618        assert!(one.coefficients[0].is_one());
619        assert!(one.coefficients[1].is_zero());
620        assert!(one.coefficients[2].is_zero());
621        assert_eq!(one, XFieldElement::ONE);
622        let zero = XFieldElement::zero();
623        assert!(zero.is_zero());
624        assert!(zero.coefficients[0].is_zero());
625        assert!(zero.coefficients[1].is_zero());
626        assert!(zero.coefficients[2].is_zero());
627        assert_eq!(zero, XFieldElement::ZERO);
628        let two = XFieldElement::new([
629            BFieldElement::new(2),
630            BFieldElement::ZERO,
631            BFieldElement::ZERO,
632        ]);
633        assert!(!two.is_one());
634        assert!(!zero.is_one());
635        let one_as_constant_term_0 = XFieldElement::new([
636            BFieldElement::new(1),
637            BFieldElement::ONE,
638            BFieldElement::ZERO,
639        ]);
640        let one_as_constant_term_1 = XFieldElement::new([
641            BFieldElement::new(1),
642            BFieldElement::ZERO,
643            BFieldElement::ONE,
644        ]);
645        assert!(!one_as_constant_term_0.is_one());
646        assert!(!one_as_constant_term_1.is_one());
647        assert!(!one_as_constant_term_0.is_zero());
648        assert!(!one_as_constant_term_1.is_zero());
649    }
650
651    #[test]
652    fn x_field_random_element_generation_test() {
653        let rand_xs: Vec<XFieldElement> = random_elements(14);
654        assert_eq!(14, rand_xs.len());
655
656        // TODO: Consider doing a statistical test.
657        assert!(rand_xs.into_iter().all_unique());
658    }
659
660    #[test]
661    fn incr_decr_test() {
662        let one_const = XFieldElement::new([1, 0, 0].map(BFieldElement::new));
663        let two_const = XFieldElement::new([2, 0, 0].map(BFieldElement::new));
664        let one_x = XFieldElement::new([0, 1, 0].map(BFieldElement::new));
665        let two_x = XFieldElement::new([0, 2, 0].map(BFieldElement::new));
666        let one_x_squared = XFieldElement::new([0, 0, 1].map(BFieldElement::new));
667        let two_x_squared = XFieldElement::new([0, 0, 2].map(BFieldElement::new));
668        let max_const = XFieldElement::new([BFieldElement::MAX, 0, 0].map(BFieldElement::new));
669        let max_x = XFieldElement::new([0, BFieldElement::MAX, 0].map(BFieldElement::new));
670        let max_x_squared = XFieldElement::new([0, 0, BFieldElement::MAX].map(BFieldElement::new));
671        let mut val = XFieldElement::ZERO;
672        val.increment(0);
673        assert!(val.is_one());
674        val.increment(0);
675        assert_eq!(two_const, val);
676        val.decrement(0);
677        assert!(val.is_one());
678        val.decrement(0);
679        assert!(val.is_zero());
680        val.decrement(0);
681        assert_eq!(max_const, val);
682        val.decrement(0);
683        assert_eq!(max_const - XFieldElement::ONE, val);
684        val.decrement(0);
685        assert_eq!(max_const - XFieldElement::ONE - XFieldElement::ONE, val);
686        val.increment(0);
687        val.increment(0);
688        val.increment(0);
689        assert!(val.is_zero());
690        val.increment(1);
691        assert_eq!(one_x, val);
692        val.increment(1);
693        assert_eq!(two_x, val);
694        val.decrement(1);
695        val.decrement(1);
696        assert!(val.is_zero());
697        val.decrement(1);
698        assert_eq!(max_x, val);
699        val.increment(1);
700        val.increment(2);
701        assert_eq!(one_x_squared, val);
702        val.increment(2);
703        assert_eq!(two_x_squared, val);
704        val.decrement(2);
705        val.decrement(2);
706        assert!(val.is_zero());
707        val.decrement(2);
708        assert_eq!(max_x_squared, val);
709        val.decrement(1);
710        val.decrement(0);
711        assert_eq!(max_x_squared + max_x + max_const, val);
712        val.decrement(2);
713        val.decrement(1);
714        val.decrement(0);
715        assert_eq!(
716            max_x_squared + max_x + max_const - one_const - one_x - one_x_squared,
717            val
718        );
719    }
720
721    #[test]
722    fn x_field_add_test() {
723        let poly1 = XFieldElement::new([2, 0, 0].map(BFieldElement::new));
724        let poly2 = XFieldElement::new([3, 0, 0].map(BFieldElement::new));
725
726        let mut poly_sum = XFieldElement::new([5, 0, 0].map(BFieldElement::new));
727        assert_eq!(poly_sum, poly1 + poly2);
728
729        let poly3 = XFieldElement::new([0, 5, 0].map(BFieldElement::new));
730        let poly4 = XFieldElement::new([0, 7, 0].map(BFieldElement::new));
731
732        poly_sum = XFieldElement::new([0, 12, 0].map(BFieldElement::new));
733        assert_eq!(poly_sum, poly3 + poly4);
734
735        let poly5 = XFieldElement::new([0, 0, 14].map(BFieldElement::new));
736        let poly6 = XFieldElement::new([0, 0, 23].map(BFieldElement::new));
737
738        poly_sum = XFieldElement::new([0, 0, 37].map(BFieldElement::new));
739        assert_eq!(poly_sum, poly5 + poly6);
740
741        let poly7 = XFieldElement::new([0, 0, BFieldElement::MAX].map(BFieldElement::new));
742        let poly8 = XFieldElement::new([0, 0, 23].map(BFieldElement::new));
743
744        poly_sum = XFieldElement::new([0, 0, 22].map(BFieldElement::new));
745        assert_eq!(poly_sum, poly7 + poly8);
746
747        let poly9 = XFieldElement::new([BFieldElement::MAX - 2, 12, 4].map(BFieldElement::new));
748        let poly10 = XFieldElement::new([2, 45000, BFieldElement::MAX - 3].map(BFieldElement::new));
749
750        poly_sum = XFieldElement::new([BFieldElement::MAX, 45012, 0].map(BFieldElement::new));
751        assert_eq!(poly_sum, poly9 + poly10);
752    }
753
754    #[test]
755    fn x_field_sub_test() {
756        let poly1 = XFieldElement::new([2, 0, 0].map(BFieldElement::new));
757        let poly2 = XFieldElement::new([3, 0, 0].map(BFieldElement::new));
758
759        let mut poly_diff = XFieldElement::new([1, 0, 0].map(BFieldElement::new));
760        assert_eq!(poly_diff, poly2 - poly1);
761
762        let poly3 = XFieldElement::new([0, 5, 0].map(BFieldElement::new));
763        let poly4 = XFieldElement::new([0, 7, 0].map(BFieldElement::new));
764
765        poly_diff = XFieldElement::new([0, 2, 0].map(BFieldElement::new));
766        assert_eq!(poly_diff, poly4 - poly3);
767
768        let poly5 = XFieldElement::new([0, 0, 14].map(BFieldElement::new));
769        let poly6 = XFieldElement::new([0, 0, 23].map(BFieldElement::new));
770
771        poly_diff = XFieldElement::new([0, 0, 9].map(BFieldElement::new));
772        assert_eq!(poly_diff, poly6 - poly5);
773
774        let poly7 = XFieldElement::new([0, 0, BFieldElement::MAX].map(BFieldElement::new));
775        let poly8 = XFieldElement::new([0, 0, 23].map(BFieldElement::new));
776
777        poly_diff = XFieldElement::new([0, 0, 24].map(BFieldElement::new));
778        assert_eq!(poly_diff, poly8 - poly7);
779
780        let poly9 = XFieldElement::new([BFieldElement::MAX - 2, 12, 4].map(BFieldElement::new));
781        let poly10 = XFieldElement::new([2, 45000, BFieldElement::MAX - 3].map(BFieldElement::new));
782
783        poly_diff = XFieldElement::new([5, 44988, BFieldElement::MAX - 7].map(BFieldElement::new));
784        assert_eq!(poly_diff, poly10 - poly9);
785    }
786
787    #[test]
788    fn x_field_mul_test() {
789        let poly1 = XFieldElement::new([2, 0, 0].map(BFieldElement::new));
790        let poly2 = XFieldElement::new([3, 0, 0].map(BFieldElement::new));
791
792        let poly12_product = XFieldElement::new([6, 0, 0].map(BFieldElement::new));
793        assert_eq!(poly12_product, poly1 * poly2);
794
795        let poly3 = XFieldElement::new([0, 3, 0].map(BFieldElement::new));
796        let poly4 = XFieldElement::new([0, 3, 0].map(BFieldElement::new));
797
798        let poly34_product = XFieldElement::new([0, 0, 9].map(BFieldElement::new));
799        assert_eq!(poly34_product, poly3 * poly4);
800
801        let poly5 = XFieldElement::new([125, 0, 0].map(BFieldElement::new));
802        let poly6 = XFieldElement::new([0, 0, 5].map(BFieldElement::new));
803
804        let poly56_product = XFieldElement::new([0, 0, 625].map(BFieldElement::new));
805        assert_eq!(poly56_product, poly5 * poly6);
806
807        // x^2 * x^2 = x^4 = x^2 - x mod (x^3 - x + 1)
808        let poly7 = XFieldElement::new([0, 0, 1].map(BFieldElement::new));
809        let poly8 = XFieldElement::new([0, 0, 1].map(BFieldElement::new));
810
811        let poly78_product = XFieldElement::new([0, BFieldElement::MAX, 1].map(BFieldElement::new));
812        assert_eq!(poly78_product, poly7 * poly8);
813
814        // x^2 * x = x^3 = x - 1 mod (x^3 - x + 1)
815        let poly9 = XFieldElement::new([0, 1, 0].map(BFieldElement::new));
816        let poly10 = XFieldElement::new([0, 0, 1].map(BFieldElement::new));
817
818        let poly910_product =
819            XFieldElement::new([BFieldElement::MAX, 1, 0].map(BFieldElement::new));
820        assert_eq!(poly910_product, poly9 * poly10);
821
822        // (13+2x+3x2)(19+5x2) = 247+122x^2+38x+10x^3+15x^4
823        let poly11 = XFieldElement::new([13, 2, 3].map(BFieldElement::new));
824        let poly12 = XFieldElement::new([19, 0, 5].map(BFieldElement::new));
825
826        let poly1112_product = XFieldElement::new([237, 33, 137].map(BFieldElement::new));
827        assert_eq!(poly1112_product, poly11 * poly12);
828    }
829
830    #[test]
831    fn x_field_overloaded_arithmetic_test() {
832        let mut rng = rand::rng();
833        for _ in 0..100 {
834            let xfe = rng.random::<XFieldElement>();
835            let bfe = rng.random::<BFieldElement>();
836
837            // 1. xfe + bfe.lift() = bfe.lift() + xfe
838            // 2. xfe + bfe = xfe + bfe.lift()
839            // 3. bfe + xfe = xfe + bfe.lift()
840            let expected_add = xfe + bfe.lift();
841            assert_eq!(expected_add, bfe.lift() + xfe);
842            assert_eq!(expected_add, xfe + bfe);
843            assert_eq!(expected_add, bfe + xfe);
844
845            // 4. xfe * bfe.lift() = bfe.lift() * xfe
846            // 5. xfe * bfe = xfe * bfe.lift()
847            // 6. bfe * xfe = xfe * bfe.lift()
848            let expected_mul = xfe * bfe.lift();
849            assert_eq!(expected_mul, bfe.lift() * xfe);
850            assert_eq!(expected_mul, xfe * bfe);
851            assert_eq!(expected_mul, bfe * xfe);
852
853            // 7. xfe - bfe = xfe - bfe.lift()
854            // 8. bfe - xfe = xfe - bfe.lift()
855            assert_eq!(xfe - bfe.lift(), xfe - bfe);
856            assert_eq!(bfe.lift() - xfe, bfe - xfe);
857        }
858    }
859
860    #[test]
861    fn x_field_into_test() {
862        let zero_poly: XFieldElement = Polynomial::<BFieldElement>::new(vec![]).into();
863        assert!(zero_poly.is_zero());
864
865        let shah_zero: XFieldElement = XFieldElement::shah_polynomial().into();
866        assert!(shah_zero.is_zero());
867
868        let neg_shah_zero: XFieldElement =
869            XFieldElement::shah_polynomial().scalar_mul(bfe!(-1)).into();
870        assert!(neg_shah_zero.is_zero());
871    }
872
873    #[test]
874    fn x_field_xgcp_test() {
875        // Verify expected properties of XGCP: symmetry and that gcd is always
876        // one. gcd is always one for all field elements.
877        let one = XFieldElement::new([1, 0, 0].map(BFieldElement::new));
878        let two = XFieldElement::new([2, 0, 0].map(BFieldElement::new));
879        let hundred = XFieldElement::new([100, 0, 0].map(BFieldElement::new));
880        let x = XFieldElement::new([0, 1, 0].map(BFieldElement::new));
881        let x_squared = XFieldElement::new([0, 0, 1].map(BFieldElement::new));
882        let one_one_one = XFieldElement::new([1, 1, 1].map(BFieldElement::new));
883        let complex0 = XFieldElement::new([450, 967, 21444444201].map(BFieldElement::new));
884        let complex1 = XFieldElement::new([456230, 0, 4563210789].map(BFieldElement::new));
885        let complex2 = XFieldElement::new([0, 96701, 456703214].map(BFieldElement::new));
886        let complex3 = XFieldElement::new([124504, 9654677, 0].map(BFieldElement::new));
887        let complex4 = XFieldElement::new(
888            [BFieldElement::MAX, BFieldElement::MAX, BFieldElement::MAX].map(BFieldElement::new),
889        );
890        let complex5 =
891            XFieldElement::new([0, BFieldElement::MAX, BFieldElement::MAX].map(BFieldElement::new));
892        let complex6 =
893            XFieldElement::new([BFieldElement::MAX, 0, BFieldElement::MAX].map(BFieldElement::new));
894        let complex7 =
895            XFieldElement::new([BFieldElement::MAX, BFieldElement::MAX, 0].map(BFieldElement::new));
896
897        let x_field_elements = vec![
898            one,
899            two,
900            hundred,
901            x,
902            x_squared,
903            one_one_one,
904            complex0,
905            complex1,
906            complex2,
907            complex3,
908            complex4,
909            complex5,
910            complex6,
911            complex7,
912        ];
913        for x_field_element in x_field_elements.iter() {
914            let x_field_element_poly: Polynomial<BFieldElement> = (*x_field_element).into();
915            // XGCP for x
916            let (gcd_0, a_0, b_0) = Polynomial::xgcd(
917                x_field_element_poly.clone(),
918                XFieldElement::shah_polynomial(),
919            );
920            let (gcd_1, b_1, a_1) =
921                Polynomial::xgcd(XFieldElement::shah_polynomial(), (*x_field_element).into());
922
923            // Verify symmetry, and that all elements are mutual primes, meaning that
924            // they form a field
925            assert!(gcd_0.is_one());
926            assert!(gcd_1.is_one());
927            assert_eq!(a_0, a_1);
928            assert_eq!(b_0, b_1);
929
930            // Verify Bezout relations: ax + by = gcd
931            assert_eq!(
932                gcd_0,
933                a_0 * x_field_element_poly + b_0 * XFieldElement::shah_polynomial()
934            );
935        }
936    }
937
938    #[test]
939    fn x_field_inv_test() {
940        let one = XFieldElement::new([1, 0, 0].map(BFieldElement::new));
941        let one_inv = one.inverse();
942        assert!((one_inv * one).is_one());
943        assert!((one * one_inv).is_one());
944
945        let two = XFieldElement::new([2, 0, 0].map(BFieldElement::new));
946        let two_inv = two.inverse();
947        assert!((two_inv * two).is_one());
948        assert!((two * two_inv).is_one());
949
950        let three = XFieldElement::new([3, 0, 0].map(BFieldElement::new));
951        let three_inv = three.inverse();
952        assert!((three_inv * three).is_one());
953        assert!((three * three_inv).is_one());
954
955        let hundred = XFieldElement::new([100, 0, 0].map(BFieldElement::new));
956        let hundred_inv = hundred.inverse();
957        assert!((hundred_inv * hundred).is_one());
958        assert!((hundred * hundred_inv).is_one());
959
960        let x = XFieldElement::new([0, 1, 0].map(BFieldElement::new));
961        let x_inv = x.inverse();
962        assert!((x_inv * x).is_one());
963        assert!((x * x_inv).is_one());
964
965        // Test batch inversion
966        let mut inverses = XFieldElement::batch_inversion(vec![]);
967        assert!(inverses.is_empty());
968        inverses = XFieldElement::batch_inversion(vec![one]);
969        assert_eq!(1, inverses.len());
970        assert!(inverses[0].is_one());
971        inverses = XFieldElement::batch_inversion(vec![two]);
972        assert_eq!(1, inverses.len());
973        assert_eq!(two_inv, inverses[0]);
974        inverses = XFieldElement::batch_inversion(vec![x]);
975        assert_eq!(1, inverses.len());
976        assert_eq!(x_inv, inverses[0]);
977        inverses = XFieldElement::batch_inversion(vec![two, x]);
978        assert_eq!(2, inverses.len());
979        assert_eq!(two_inv, inverses[0]);
980        assert_eq!(x_inv, inverses[1]);
981
982        let input = vec![one, two, three, hundred, x];
983        inverses = XFieldElement::batch_inversion(input.clone());
984        let inverses_inverses = XFieldElement::batch_inversion(inverses.clone());
985        assert_eq!(input.len(), inverses.len());
986        for i in 0..input.len() {
987            assert!((inverses[i] * input[i]).is_one());
988            assert_eq!(input[i], inverses_inverses[i]);
989        }
990    }
991
992    #[proptest]
993    fn field_element_inversion(
994        #[filter(!#x.is_zero())] x: XFieldElement,
995        #[filter(!#disturbance.is_zero())]
996        #[filter(#x != #disturbance)]
997        disturbance: XFieldElement,
998    ) {
999        let not_x = x - disturbance;
1000        prop_assert_eq!(XFieldElement::ONE, x * x.inverse());
1001        prop_assert_eq!(XFieldElement::ONE, not_x * not_x.inverse());
1002        prop_assert_ne!(XFieldElement::ONE, x * not_x.inverse());
1003    }
1004
1005    #[proptest]
1006    fn field_element_batch_inversion(
1007        #[filter(!#xs.iter().any(|x| x.is_zero()))] xs: Vec<XFieldElement>,
1008    ) {
1009        let inverses = XFieldElement::batch_inversion(xs.clone());
1010        for (x, inv) in xs.into_iter().zip(inverses) {
1011            prop_assert_eq!(XFieldElement::ONE, x * inv);
1012        }
1013    }
1014
1015    #[test]
1016    fn mul_xfe_with_bfe_pbt() {
1017        let test_iterations = 100;
1018        let rands_x: Vec<XFieldElement> = random_elements(test_iterations);
1019        let rands_b: Vec<BFieldElement> = random_elements(test_iterations);
1020        for (mut x, b) in izip!(rands_x, rands_b) {
1021            let res_mul = x * b;
1022            assert_eq!(res_mul.coefficients[0], x.coefficients[0] * b);
1023            assert_eq!(res_mul.coefficients[1], x.coefficients[1] * b);
1024            assert_eq!(res_mul.coefficients[2], x.coefficients[2] * b);
1025
1026            // Also verify that the `MulAssign` implementation agrees with the `Mul` implementation
1027            x *= b;
1028            let res_mul_assign = x;
1029            assert_eq!(res_mul, res_mul_assign);
1030        }
1031    }
1032
1033    #[proptest(cases = 1_000)]
1034    fn x_field_division_mul_pbt(a: XFieldElement, b: XFieldElement) {
1035        let a_b = a * b;
1036        let b_a = b * a;
1037        prop_assert_eq!(a_b, b_a);
1038        prop_assert_eq!(a_b / b, a);
1039        prop_assert_eq!(a_b / a, b);
1040        prop_assert_eq!(a * a, a.square());
1041
1042        // Test the add/sub/mul assign operators
1043        let mut a_minus_b = a;
1044        a_minus_b -= b;
1045        prop_assert_eq!(a - b, a_minus_b);
1046
1047        let mut a_plus_b = a;
1048        a_plus_b += b;
1049        prop_assert_eq!(a + b, a_plus_b);
1050
1051        let mut a_mul_b = a;
1052        a_mul_b *= b;
1053        prop_assert_eq!(a * b, a_mul_b);
1054
1055        // Test the add/sub/mul assign operators, when the higher coefficients are zero.
1056        // Also tests add/sub/mul operators and add/sub/mul assign operators when RHS has
1057        // the type of B field element. And add/sub/mul operators when LHS is a B-field
1058        // element and RHS is an X-field element.
1059        // mul-assign `*=`
1060        let b_field_b = XFieldElement::new_const(b.coefficients[0]);
1061        let mut a_mul_b_field_b_as_x = a;
1062        a_mul_b_field_b_as_x *= b_field_b;
1063        prop_assert_eq!(a * b_field_b, a_mul_b_field_b_as_x);
1064        prop_assert_eq!(a, a_mul_b_field_b_as_x / b_field_b);
1065        prop_assert_eq!(b_field_b, a_mul_b_field_b_as_x / a);
1066        prop_assert_eq!(a_mul_b_field_b_as_x, a * b.coefficients[0]);
1067        prop_assert_eq!(a_mul_b_field_b_as_x, b.coefficients[0] * a);
1068        let mut a_mul_b_field_b_as_b = a;
1069        a_mul_b_field_b_as_b *= b.coefficients[0];
1070        prop_assert_eq!(a_mul_b_field_b_as_b, a_mul_b_field_b_as_x);
1071
1072        // `+=`
1073        let mut a_plus_b_field_b_as_x = a;
1074        a_plus_b_field_b_as_x += b_field_b;
1075        prop_assert_eq!(a + b_field_b, a_plus_b_field_b_as_x);
1076        prop_assert_eq!(a, a_plus_b_field_b_as_x - b_field_b);
1077        prop_assert_eq!(b_field_b, a_plus_b_field_b_as_x - a);
1078        prop_assert_eq!(a_plus_b_field_b_as_x, a + b.coefficients[0]);
1079        prop_assert_eq!(a_plus_b_field_b_as_x, b.coefficients[0] + a);
1080        let mut a_plus_b_field_b_as_b = a;
1081        a_plus_b_field_b_as_b += b.coefficients[0];
1082        prop_assert_eq!(a_plus_b_field_b_as_b, a_plus_b_field_b_as_x);
1083
1084        // `-=`
1085        let mut a_minus_b_field_b_as_x = a;
1086        a_minus_b_field_b_as_x -= b_field_b;
1087        prop_assert_eq!(a - b_field_b, a_minus_b_field_b_as_x);
1088        prop_assert_eq!(a, a_minus_b_field_b_as_x + b_field_b);
1089        prop_assert_eq!(-b_field_b, a_minus_b_field_b_as_x - a);
1090        prop_assert_eq!(a_minus_b_field_b_as_x, a - b.coefficients[0]);
1091        prop_assert_eq!(-a_minus_b_field_b_as_x, b.coefficients[0] - a);
1092        let mut a_minus_b_field_b_as_b = a;
1093        a_minus_b_field_b_as_b -= b.coefficients[0];
1094        prop_assert_eq!(a_minus_b_field_b_as_b, a_minus_b_field_b_as_x);
1095    }
1096
1097    #[test]
1098    fn xfe_mod_pow_zero() {
1099        assert!(XFieldElement::ZERO.mod_pow_u32(0).is_one());
1100        assert!(XFieldElement::ZERO.mod_pow_u64(0).is_one());
1101        assert!(XFieldElement::ONE.mod_pow_u32(0).is_one());
1102        assert!(XFieldElement::ONE.mod_pow_u64(0).is_one());
1103    }
1104
1105    #[proptest]
1106    fn xfe_mod_pow(base: XFieldElement, #[strategy(0_u32..200)] exponent: u32) {
1107        let mut acc = XFieldElement::ONE;
1108        for i in 0..exponent {
1109            assert_eq!(acc, base.mod_pow_u32(i));
1110            acc *= base;
1111        }
1112    }
1113
1114    #[test]
1115    fn xfe_mod_pow_static() {
1116        let three_to_the_n = |n| xfe!(3).mod_pow_u64(n);
1117        let actual = [0, 1, 2, 3, 4, 5].map(three_to_the_n);
1118        let expected = xfe_array![1, 3, 9, 27, 81, 243];
1119        assert_eq!(expected, actual);
1120    }
1121
1122    #[proptest(cases = 100)]
1123    fn xfe_intt_is_inverse_of_xfe_ntt(
1124        #[strategy(1..=11)]
1125        #[map(|log| 1_usize << log)]
1126        _num_inputs: usize,
1127        #[strategy(vec(arb(), #_num_inputs))] inputs: Vec<XFieldElement>,
1128    ) {
1129        let mut rv = inputs.clone();
1130        ntt::<XFieldElement>(&mut rv);
1131        intt::<XFieldElement>(&mut rv);
1132        prop_assert_eq!(inputs, rv);
1133    }
1134
1135    #[proptest(cases = 40)]
1136    fn xfe_ntt_corresponds_to_polynomial_evaluation(
1137        #[strategy(1..=11)]
1138        #[map(|log_2| 1_u64 << log_2)]
1139        root_order: u64,
1140        #[strategy(vec(arb(), #root_order as usize))] inputs: Vec<XFieldElement>,
1141    ) {
1142        let root = XFieldElement::primitive_root_of_unity(root_order).unwrap();
1143        let mut rv = inputs.clone();
1144        ntt::<XFieldElement>(&mut rv);
1145
1146        let poly = Polynomial::new(inputs);
1147        let domain = root.get_cyclic_group_elements(None);
1148        let evaluations = poly.batch_evaluate(&domain);
1149        prop_assert_eq!(evaluations, rv);
1150    }
1151
1152    #[test]
1153    fn inverse_or_zero_of_zero_is_zero() {
1154        let zero = XFieldElement::ZERO;
1155        assert_eq!(zero, zero.inverse_or_zero());
1156    }
1157
1158    #[proptest]
1159    fn inverse_or_zero_of_non_zero_is_inverse(#[filter(!#xfe.is_zero())] xfe: XFieldElement) {
1160        let inv = xfe.inverse_or_zero();
1161        prop_assert_ne!(XFieldElement::ZERO, inv);
1162        prop_assert_eq!(XFieldElement::ONE, xfe * inv);
1163    }
1164
1165    #[test]
1166    #[should_panic(expected = "Cannot invert the zero element in the extension field.")]
1167    fn multiplicative_inverse_of_zero() {
1168        let zero = XFieldElement::ZERO;
1169        let _ = zero.inverse();
1170    }
1171
1172    #[proptest]
1173    fn xfe_to_digest_to_xfe_is_invariant(xfe: XFieldElement) {
1174        let digest: Digest = xfe.into();
1175        let xfe2: XFieldElement = digest.try_into().unwrap();
1176        assert_eq!(xfe, xfe2);
1177    }
1178
1179    #[proptest]
1180    fn converting_random_digest_to_xfield_element_fails(digest: Digest) {
1181        if XFieldElement::try_from(digest).is_ok() {
1182            let reason = "Should not be able to convert random `Digest` to an `XFieldElement`.";
1183            return Err(TestCaseError::Fail(reason.into()));
1184        }
1185    }
1186
1187    #[test]
1188    fn xfe_macro_can_be_used() {
1189        let x = xfe!(42);
1190        let _ = xfe!(42u32);
1191        let _ = xfe!(-1);
1192        let _ = xfe!(x);
1193        let _ = xfe!([x.coefficients[0], x.coefficients[1], x.coefficients[2]]);
1194        let y = xfe!(bfe!(42));
1195        assert_eq!(x, y);
1196
1197        let a = xfe!([bfe!(42), bfe!(43), bfe!(44)]);
1198        let b = xfe!([42, 43, 44]);
1199        assert_eq!(a, b);
1200
1201        let m: [XFieldElement; 3] = xfe_array![[1, 2, 3], [4, 5, 6], [7, 8, 9]];
1202        let n: Vec<XFieldElement> = xfe_vec![[1, 2, 3], [4, 5, 6], [7, 8, 9]];
1203        assert_eq!(m.to_vec(), n);
1204    }
1205
1206    #[proptest]
1207    fn xfe_macro_produces_same_result_as_calling_new(coeffs: [BFieldElement; EXTENSION_DEGREE]) {
1208        let xfe = XFieldElement::new(coeffs);
1209        prop_assert_eq!(xfe, xfe!(coeffs));
1210    }
1211
1212    #[proptest]
1213    fn xfe_macro_produces_same_result_as_calling_new_const(scalar: BFieldElement) {
1214        let xfe = XFieldElement::new_const(scalar);
1215        prop_assert_eq!(xfe, xfe!(scalar));
1216    }
1217}