p3_field/
array.rs

1use core::array;
2use core::iter::{Product, Sum};
3use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign};
4
5use crate::batch_inverse::batch_multiplicative_inverse_general;
6use crate::{Algebra, Field, PackedValue, PrimeCharacteristicRing};
7
8#[derive(Clone, Copy, Debug, Eq, PartialEq)]
9#[repr(transparent)] // Needed to make `transmute`s safe.
10pub struct FieldArray<F: Field, const N: usize>(pub [F; N]);
11
12impl<F: Field, const N: usize> FieldArray<F, N> {
13    pub(crate) fn inverse(&self) -> Self {
14        let mut result = Self::default();
15        batch_multiplicative_inverse_general(&self.0, &mut result.0, |x| x.inverse());
16        result
17    }
18}
19
20impl<F: Field, const N: usize> Default for FieldArray<F, N> {
21    fn default() -> Self {
22        Self::ZERO
23    }
24}
25
26impl<F: Field, const N: usize> From<F> for FieldArray<F, N> {
27    fn from(val: F) -> Self {
28        [val; N].into()
29    }
30}
31
32impl<F: Field, const N: usize> From<[F; N]> for FieldArray<F, N> {
33    fn from(arr: [F; N]) -> Self {
34        Self(arr)
35    }
36}
37
38impl<F: Field, const N: usize> PrimeCharacteristicRing for FieldArray<F, N> {
39    type PrimeSubfield = F::PrimeSubfield;
40
41    const ZERO: Self = Self([F::ZERO; N]);
42    const ONE: Self = Self([F::ONE; N]);
43    const TWO: Self = Self([F::TWO; N]);
44    const NEG_ONE: Self = Self([F::NEG_ONE; N]);
45
46    #[inline]
47    fn from_prime_subfield(f: Self::PrimeSubfield) -> Self {
48        F::from_prime_subfield(f).into()
49    }
50}
51
52impl<F: Field, const N: usize> Algebra<F> for FieldArray<F, N> {}
53
54unsafe impl<F: Field, const N: usize> PackedValue for FieldArray<F, N> {
55    type Value = F;
56
57    const WIDTH: usize = N;
58
59    fn from_slice(slice: &[Self::Value]) -> &Self {
60        assert_eq!(slice.len(), Self::WIDTH);
61        unsafe { &*slice.as_ptr().cast() }
62    }
63
64    fn from_slice_mut(slice: &mut [Self::Value]) -> &mut Self {
65        assert_eq!(slice.len(), Self::WIDTH);
66        unsafe { &mut *slice.as_mut_ptr().cast() }
67    }
68
69    fn from_fn<Fn>(f: Fn) -> Self
70    where
71        Fn: FnMut(usize) -> Self::Value,
72    {
73        Self(array::from_fn(f))
74    }
75
76    fn as_slice(&self) -> &[Self::Value] {
77        &self.0
78    }
79
80    fn as_slice_mut(&mut self) -> &mut [Self::Value] {
81        &mut self.0
82    }
83}
84
85impl<F: Field, const N: usize> Add for FieldArray<F, N> {
86    type Output = Self;
87
88    #[inline]
89    fn add(self, rhs: Self) -> Self::Output {
90        array::from_fn(|i| self.0[i] + rhs.0[i]).into()
91    }
92}
93
94impl<F: Field, const N: usize> Add<F> for FieldArray<F, N> {
95    type Output = Self;
96
97    #[inline]
98    fn add(self, rhs: F) -> Self::Output {
99        self.0.map(|x| x + rhs).into()
100    }
101}
102
103impl<F: Field, const N: usize> AddAssign for FieldArray<F, N> {
104    #[inline]
105    fn add_assign(&mut self, rhs: Self) {
106        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x += y);
107    }
108}
109
110impl<F: Field, const N: usize> AddAssign<F> for FieldArray<F, N> {
111    #[inline]
112    fn add_assign(&mut self, rhs: F) {
113        self.0.iter_mut().for_each(|x| *x += rhs);
114    }
115}
116
117impl<F: Field, const N: usize> Sub for FieldArray<F, N> {
118    type Output = Self;
119
120    #[inline]
121    fn sub(self, rhs: Self) -> Self::Output {
122        array::from_fn(|i| self.0[i] - rhs.0[i]).into()
123    }
124}
125
126impl<F: Field, const N: usize> Sub<F> for FieldArray<F, N> {
127    type Output = Self;
128
129    #[inline]
130    fn sub(self, rhs: F) -> Self::Output {
131        self.0.map(|x| x - rhs).into()
132    }
133}
134
135impl<F: Field, const N: usize> SubAssign for FieldArray<F, N> {
136    #[inline]
137    fn sub_assign(&mut self, rhs: Self) {
138        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x -= y);
139    }
140}
141
142impl<F: Field, const N: usize> SubAssign<F> for FieldArray<F, N> {
143    #[inline]
144    fn sub_assign(&mut self, rhs: F) {
145        self.0.iter_mut().for_each(|x| *x -= rhs);
146    }
147}
148
149impl<F: Field, const N: usize> Neg for FieldArray<F, N> {
150    type Output = Self;
151
152    #[inline]
153    fn neg(self) -> Self::Output {
154        self.0.map(|x| -x).into()
155    }
156}
157
158impl<F: Field, const N: usize> Mul for FieldArray<F, N> {
159    type Output = Self;
160
161    #[inline]
162    fn mul(self, rhs: Self) -> Self::Output {
163        array::from_fn(|i| self.0[i] * rhs.0[i]).into()
164    }
165}
166
167impl<F: Field, const N: usize> Mul<F> for FieldArray<F, N> {
168    type Output = Self;
169
170    #[inline]
171    fn mul(self, rhs: F) -> Self::Output {
172        self.0.map(|x| x * rhs).into()
173    }
174}
175
176impl<F: Field, const N: usize> MulAssign for FieldArray<F, N> {
177    #[inline]
178    fn mul_assign(&mut self, rhs: Self) {
179        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x *= y);
180    }
181}
182
183impl<F: Field, const N: usize> MulAssign<F> for FieldArray<F, N> {
184    #[inline]
185    fn mul_assign(&mut self, rhs: F) {
186        self.0.iter_mut().for_each(|x| *x *= rhs);
187    }
188}
189
190impl<F: Field, const N: usize> Div<F> for FieldArray<F, N> {
191    type Output = Self;
192
193    #[allow(clippy::suspicious_arithmetic_impl)]
194    #[inline]
195    fn div(self, rhs: F) -> Self::Output {
196        let rhs_inv = rhs.inverse();
197        self * rhs_inv
198    }
199}
200
201impl<F: Field, const N: usize> Sum for FieldArray<F, N> {
202    #[inline]
203    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
204        iter.reduce(|lhs, rhs| lhs + rhs).unwrap_or(Self::ZERO)
205    }
206}
207
208impl<F: Field, const N: usize> Product for FieldArray<F, N> {
209    #[inline]
210    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
211        iter.reduce(|lhs, rhs| lhs * rhs).unwrap_or(Self::ONE)
212    }
213}