p3_field/
array.rs

1use core::array;
2use core::iter::{Product, Sum};
3use core::ops::{Add, AddAssign, Div, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign};
4
5use p3_util::{as_base_slice, as_base_slice_mut};
6
7use crate::batch_inverse::batch_multiplicative_inverse_general;
8use crate::{Algebra, Field, PackedValue, PrimeCharacteristicRing};
9
10#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11#[repr(transparent)] // Needed to make `transmute`s safe.
12#[must_use]
13pub struct FieldArray<F: Field, const N: usize>(pub [F; N]);
14
15impl<F: Field, const N: usize> FieldArray<F, N> {
16    /// Compute the element-wise multiplicative inverse using batched inversion.
17    ///
18    /// Uses Montgomery's batch inversion trick to compute all inverses with a
19    /// single field inversion, improving performance when N > 1.
20    #[inline]
21    pub(crate) fn inverse(&self) -> Self {
22        let mut result = Self::default();
23        batch_multiplicative_inverse_general(&self.0, &mut result.0, |x| x.inverse());
24        result
25    }
26
27    /// Apply a function to each element, returning a new `FieldArray<U, N>`.
28    #[inline]
29    pub fn map<Func, U: Field>(self, f: Func) -> FieldArray<U, N>
30    where
31        Func: FnMut(F) -> U,
32    {
33        FieldArray(self.map_into_array(f))
34    }
35
36    /// Apply a function to each element, returning a raw array `[U; N]`.
37    ///
38    /// Unlike [`map`](Self::map), this does not require `U: Field`.
39    #[inline]
40    pub fn map_into_array<Func, U>(self, f: Func) -> [U; N]
41    where
42        Func: FnMut(F) -> U,
43    {
44        self.0.map(f)
45    }
46
47    /// View as a slice of raw `[F; N]` arrays.
48    ///
49    /// This is a zero-cost transmute enabled by the `#[repr(transparent)]` layout.
50    #[inline]
51    pub const fn as_raw_slice(s: &[Self]) -> &[[F; N]] {
52        // SAFETY: `FieldArray<F, N>` is `#[repr(transparent)]` over `[F; N]`,
53        // so `&[FieldArray<F, N>]` and `&[[F; N]]` have identical layouts.
54        unsafe { as_base_slice(s) }
55    }
56
57    /// View as a mutable slice of raw `[F; N]` arrays.
58    ///
59    /// This is a zero-cost transmute enabled by the `#[repr(transparent)]` layout.
60    #[inline]
61    pub const fn as_raw_slice_mut(s: &mut [Self]) -> &mut [[F; N]] {
62        // SAFETY: `FieldArray<F, N>` is `#[repr(transparent)]` over `[F; N]`,
63        // so `&mut [FieldArray<F, N>]` and `&mut [[F; N]]` have identical layouts.
64        unsafe { as_base_slice_mut(s) }
65    }
66
67    /// Reinterpret a slice of `[F; N]` as a slice of `FieldArray<F, N>`.
68    ///
69    /// This is a zero-cost transmute enabled by the `#[repr(transparent)]` layout.
70    #[inline]
71    pub const fn from_raw_slice(s: &[[F; N]]) -> &[Self] {
72        // SAFETY: `FieldArray<F, N>` is `#[repr(transparent)]` over `[F; N]`,
73        // so `&[[F; N]]` and `&[FieldArray<F, N>]` have identical layouts.
74        unsafe { as_base_slice(s) }
75    }
76
77    /// Reinterpret a mutable slice of `[F; N]` as a mutable slice of `FieldArray<F, N>`.
78    ///
79    /// This is a zero-cost transmute enabled by the `#[repr(transparent)]` layout.
80    #[inline]
81    pub const fn from_raw_slice_mut(s: &mut [[F; N]]) -> &mut [Self] {
82        // SAFETY: `FieldArray<F, N>` is `#[repr(transparent)]` over `[F; N]`,
83        // so `&mut [[F; N]]` and `&mut [FieldArray<F, N>]` have identical layouts.
84        unsafe { as_base_slice_mut(s) }
85    }
86}
87
88impl<F: Field, const N: usize> IndexMut<usize> for FieldArray<F, N> {
89    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
90        self.0.index_mut(index)
91    }
92}
93
94impl<F: Field, const N: usize> Index<usize> for FieldArray<F, N> {
95    type Output = F;
96
97    fn index(&self, index: usize) -> &Self::Output {
98        self.0.index(index)
99    }
100}
101
102impl<F: Field, const N: usize> Default for FieldArray<F, N> {
103    fn default() -> Self {
104        Self::ZERO
105    }
106}
107
108impl<F: Field, const N: usize> From<F> for FieldArray<F, N> {
109    fn from(val: F) -> Self {
110        [val; N].into()
111    }
112}
113
114impl<F: Field, const N: usize> From<[F; N]> for FieldArray<F, N> {
115    fn from(arr: [F; N]) -> Self {
116        Self(arr)
117    }
118}
119
120impl<F: Field, const N: usize> PrimeCharacteristicRing for FieldArray<F, N> {
121    type PrimeSubfield = F::PrimeSubfield;
122
123    const ZERO: Self = Self([F::ZERO; N]);
124    const ONE: Self = Self([F::ONE; N]);
125    const TWO: Self = Self([F::TWO; N]);
126    const NEG_ONE: Self = Self([F::NEG_ONE; N]);
127
128    #[inline]
129    fn from_prime_subfield(f: Self::PrimeSubfield) -> Self {
130        F::from_prime_subfield(f).into()
131    }
132
133    #[inline]
134    fn halve(&self) -> Self {
135        Self(self.0.map(|x| x.halve()))
136    }
137}
138
139impl<F: Field, const N: usize> Algebra<F> for FieldArray<F, N> {}
140
141unsafe impl<F: Field, const N: usize> PackedValue for FieldArray<F, N> {
142    type Value = F;
143
144    const WIDTH: usize = N;
145
146    fn from_slice(slice: &[Self::Value]) -> &Self {
147        assert_eq!(slice.len(), Self::WIDTH);
148        unsafe { &*slice.as_ptr().cast() }
149    }
150
151    fn from_slice_mut(slice: &mut [Self::Value]) -> &mut Self {
152        assert_eq!(slice.len(), Self::WIDTH);
153        unsafe { &mut *slice.as_mut_ptr().cast() }
154    }
155
156    fn from_fn<Fn>(f: Fn) -> Self
157    where
158        Fn: FnMut(usize) -> Self::Value,
159    {
160        Self(array::from_fn(f))
161    }
162
163    fn as_slice(&self) -> &[Self::Value] {
164        &self.0
165    }
166
167    fn as_slice_mut(&mut self) -> &mut [Self::Value] {
168        &mut self.0
169    }
170}
171
172impl<F: Field, const N: usize> Add for FieldArray<F, N> {
173    type Output = Self;
174
175    #[inline]
176    fn add(self, rhs: Self) -> Self::Output {
177        array::from_fn(|i| self.0[i] + rhs.0[i]).into()
178    }
179}
180
181impl<F: Field, const N: usize> Add<F> for FieldArray<F, N> {
182    type Output = Self;
183
184    #[inline]
185    fn add(self, rhs: F) -> Self::Output {
186        self.0.map(|x| x + rhs).into()
187    }
188}
189
190impl<F: Field, const N: usize> AddAssign for FieldArray<F, N> {
191    #[inline]
192    fn add_assign(&mut self, rhs: Self) {
193        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x += y);
194    }
195}
196
197impl<F: Field, const N: usize> AddAssign<F> for FieldArray<F, N> {
198    #[inline]
199    fn add_assign(&mut self, rhs: F) {
200        self.0.iter_mut().for_each(|x| *x += rhs);
201    }
202}
203
204impl<F: Field, const N: usize> Sub for FieldArray<F, N> {
205    type Output = Self;
206
207    #[inline]
208    fn sub(self, rhs: Self) -> Self::Output {
209        array::from_fn(|i| self.0[i] - rhs.0[i]).into()
210    }
211}
212
213impl<F: Field, const N: usize> Sub<F> for FieldArray<F, N> {
214    type Output = Self;
215
216    #[inline]
217    fn sub(self, rhs: F) -> Self::Output {
218        self.0.map(|x| x - rhs).into()
219    }
220}
221
222impl<F: Field, const N: usize> SubAssign for FieldArray<F, N> {
223    #[inline]
224    fn sub_assign(&mut self, rhs: Self) {
225        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x -= y);
226    }
227}
228
229impl<F: Field, const N: usize> SubAssign<F> for FieldArray<F, N> {
230    #[inline]
231    fn sub_assign(&mut self, rhs: F) {
232        self.0.iter_mut().for_each(|x| *x -= rhs);
233    }
234}
235
236impl<F: Field, const N: usize> Neg for FieldArray<F, N> {
237    type Output = Self;
238
239    #[inline]
240    fn neg(self) -> Self::Output {
241        self.0.map(|x| -x).into()
242    }
243}
244
245impl<F: Field, const N: usize> Mul for FieldArray<F, N> {
246    type Output = Self;
247
248    #[inline]
249    fn mul(self, rhs: Self) -> Self::Output {
250        array::from_fn(|i| self.0[i] * rhs.0[i]).into()
251    }
252}
253
254impl<F: Field, const N: usize> Mul<F> for FieldArray<F, N> {
255    type Output = Self;
256
257    #[inline]
258    fn mul(self, rhs: F) -> Self::Output {
259        self.0.map(|x| x * rhs).into()
260    }
261}
262
263impl<F: Field, const N: usize> MulAssign for FieldArray<F, N> {
264    #[inline]
265    fn mul_assign(&mut self, rhs: Self) {
266        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x *= y);
267    }
268}
269
270impl<F: Field, const N: usize> MulAssign<F> for FieldArray<F, N> {
271    #[inline]
272    fn mul_assign(&mut self, rhs: F) {
273        self.0.iter_mut().for_each(|x| *x *= rhs);
274    }
275}
276
277impl<F: Field, const N: usize> Div<F> for FieldArray<F, N> {
278    type Output = Self;
279
280    #[allow(clippy::suspicious_arithmetic_impl)]
281    #[inline]
282    fn div(self, rhs: F) -> Self::Output {
283        let rhs_inv = rhs.inverse();
284        self * rhs_inv
285    }
286}
287
288impl<F: Field, const N: usize> Sum for FieldArray<F, N> {
289    #[inline]
290    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
291        iter.reduce(|lhs, rhs| lhs + rhs).unwrap_or(Self::ZERO)
292    }
293}
294
295impl<F: Field, const N: usize> Product for FieldArray<F, N> {
296    #[inline]
297    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
298        iter.reduce(|lhs, rhs| lhs * rhs).unwrap_or(Self::ONE)
299    }
300}
p3_field/array.rs

p3_field/
array.rs
