sp1_hypercube/air/
extension.rs

1use std::ops::{Add, Div, Mul, Neg, Sub};
2
3use slop_algebra::{
4    extension::{BinomialExtensionField, BinomiallyExtendable},
5    AbstractExtensionField, AbstractField, Field,
6};
7use sp1_derive::AlignedBorrow;
8
9const D: usize = 4;
10
11/// A binomial extension element represented over a generic type `T`.
12#[derive(AlignedBorrow, Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
13#[repr(C)]
14pub struct BinomialExtension<T>(pub [T; D]);
15
16impl<T> BinomialExtension<T> {
17    /// Creates a new binomial extension element from a base element.
18    pub fn from_base(b: T) -> Self
19    where
20        T: AbstractField,
21    {
22        let mut arr: [T; D] = core::array::from_fn(|_| T::zero());
23        arr[0] = b;
24        Self(arr)
25    }
26
27    /// Returns a reference to the underlying slice.
28    pub const fn as_base_slice(&self) -> &[T] {
29        &self.0
30    }
31
32    /// Creates a new binomial extension element from a binomial extension element.
33    #[allow(clippy::needless_pass_by_value)]
34    pub fn from<S: Into<T> + Clone>(from: BinomialExtension<S>) -> Self {
35        BinomialExtension(core::array::from_fn(|i| from.0[i].clone().into()))
36    }
37}
38
39impl<T: Add<Output = T> + Clone> Add for BinomialExtension<T> {
40    type Output = Self;
41
42    fn add(self, rhs: Self) -> Self::Output {
43        Self(core::array::from_fn(|i| self.0[i].clone() + rhs.0[i].clone()))
44    }
45}
46
47impl<T: Sub<Output = T> + Clone> Sub for BinomialExtension<T> {
48    type Output = Self;
49
50    fn sub(self, rhs: Self) -> Self::Output {
51        Self(core::array::from_fn(|i| self.0[i].clone() - rhs.0[i].clone()))
52    }
53}
54
55impl<T: Add<Output = T> + Mul<Output = T> + AbstractField> Mul for BinomialExtension<T> {
56    type Output = Self;
57
58    fn mul(self, rhs: Self) -> Self::Output {
59        let mut result = [T::zero(), T::zero(), T::zero(), T::zero()];
60        // This value is specific for SP1Field prime's extension `F_p[x]/(x^4 - 3)`.
61        let w = T::from_canonical_u32(3);
62
63        for i in 0..D {
64            for j in 0..D {
65                if i + j >= D {
66                    result[i + j - D] = result[i + j - D].clone()
67                        + w.clone() * self.0[i].clone() * rhs.0[j].clone();
68                } else {
69                    result[i + j] = result[i + j].clone() + self.0[i].clone() * rhs.0[j].clone();
70                }
71            }
72        }
73
74        Self(result)
75    }
76}
77
78impl<F> Div for BinomialExtension<F>
79where
80    F: BinomiallyExtendable<D>,
81{
82    type Output = Self;
83
84    fn div(self, rhs: Self) -> Self::Output {
85        let p3_ef_lhs = BinomialExtensionField::from_base_slice(&self.0);
86        let p3_ef_rhs = BinomialExtensionField::from_base_slice(&rhs.0);
87        let p3_ef_result = p3_ef_lhs / p3_ef_rhs;
88        Self(p3_ef_result.as_base_slice().try_into().unwrap())
89    }
90}
91
92impl<F> BinomialExtension<F>
93where
94    F: BinomiallyExtendable<4>,
95{
96    /// Returns the multiplicative inverse of the element.
97    #[must_use]
98    pub fn inverse(&self) -> Self {
99        let p3_ef = BinomialExtensionField::from_base_slice(&self.0);
100        let p3_ef_inverse = p3_ef.inverse();
101        Self(p3_ef_inverse.as_base_slice().try_into().unwrap())
102    }
103
104    /// Returns the multiplicative inverse of the element, if it exists.
105    #[must_use]
106    pub fn try_inverse(&self) -> Option<Self> {
107        let p3_ef = BinomialExtensionField::from_base_slice(&self.0);
108        let p3_ef_inverse = p3_ef.try_inverse()?;
109        Some(Self(p3_ef_inverse.as_base_slice().try_into().unwrap()))
110    }
111}
112
113impl<T: AbstractField + Copy> Neg for BinomialExtension<T> {
114    type Output = Self;
115
116    fn neg(self) -> Self::Output {
117        Self([-self.0[0], -self.0[1], -self.0[2], -self.0[3]])
118    }
119}
120
121impl<AF> From<BinomialExtensionField<AF, D>> for BinomialExtension<AF>
122where
123    AF: AbstractField + Copy,
124    AF::F: BinomiallyExtendable<D>,
125{
126    fn from(value: BinomialExtensionField<AF, D>) -> Self {
127        let arr: [AF; D] = value.as_base_slice().try_into().unwrap();
128        Self(arr)
129    }
130}
131
132impl<AF> From<BinomialExtension<AF>> for BinomialExtensionField<AF, D>
133where
134    AF: AbstractField + Copy,
135    AF::F: BinomiallyExtendable<D>,
136{
137    fn from(value: BinomialExtension<AF>) -> Self {
138        BinomialExtensionField::from_base_slice(&value.0)
139    }
140}
141
142impl<T> IntoIterator for BinomialExtension<T> {
143    type Item = T;
144    type IntoIter = core::array::IntoIter<T, D>;
145
146    fn into_iter(self) -> Self::IntoIter {
147        self.0.into_iter()
148    }
149}
sp1_hypercube/air/extension.rs

sp1_hypercube/air/
extension.rs
