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::{Field, FieldAlgebra, PackedValue};
7
8#[derive(Clone, Copy, Debug, Eq, PartialEq)]
9#[repr(transparent)] // This 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> FieldAlgebra for FieldArray<F, N> {
39    type F = F;
40
41    const ZERO: Self = FieldArray([F::ZERO; N]);
42    const ONE: Self = FieldArray([F::ONE; N]);
43    const TWO: Self = FieldArray([F::TWO; N]);
44    const NEG_ONE: Self = FieldArray([F::NEG_ONE; N]);
45
46    #[inline]
47    fn from_f(f: Self::F) -> Self {
48        f.into()
49    }
50
51    fn from_canonical_u8(n: u8) -> Self {
52        [F::from_canonical_u8(n); N].into()
53    }
54
55    fn from_canonical_u16(n: u16) -> Self {
56        [F::from_canonical_u16(n); N].into()
57    }
58
59    fn from_canonical_u32(n: u32) -> Self {
60        [F::from_canonical_u32(n); N].into()
61    }
62
63    fn from_canonical_u64(n: u64) -> Self {
64        [F::from_canonical_u64(n); N].into()
65    }
66
67    fn from_canonical_usize(n: usize) -> Self {
68        [F::from_canonical_usize(n); N].into()
69    }
70
71    fn from_wrapped_u32(n: u32) -> Self {
72        [F::from_wrapped_u32(n); N].into()
73    }
74
75    fn from_wrapped_u64(n: u64) -> Self {
76        [F::from_wrapped_u64(n); N].into()
77    }
78}
79
80unsafe impl<F: Field, const N: usize> PackedValue for FieldArray<F, N> {
81    type Value = F;
82
83    const WIDTH: usize = N;
84
85    fn from_slice(slice: &[Self::Value]) -> &Self {
86        assert_eq!(slice.len(), Self::WIDTH);
87        unsafe { &*slice.as_ptr().cast() }
88    }
89
90    fn from_slice_mut(slice: &mut [Self::Value]) -> &mut Self {
91        assert_eq!(slice.len(), Self::WIDTH);
92        unsafe { &mut *slice.as_mut_ptr().cast() }
93    }
94
95    fn from_fn<Fn>(f: Fn) -> Self
96    where
97        Fn: FnMut(usize) -> Self::Value,
98    {
99        Self(array::from_fn(f))
100    }
101
102    fn as_slice(&self) -> &[Self::Value] {
103        &self.0
104    }
105
106    fn as_slice_mut(&mut self) -> &mut [Self::Value] {
107        &mut self.0
108    }
109}
110
111impl<F: Field, const N: usize> Add for FieldArray<F, N> {
112    type Output = Self;
113
114    #[inline]
115    fn add(self, rhs: Self) -> Self::Output {
116        array::from_fn(|i| self.0[i] + rhs.0[i]).into()
117    }
118}
119
120impl<F: Field, const N: usize> Add<F> for FieldArray<F, N> {
121    type Output = Self;
122
123    #[inline]
124    fn add(self, rhs: F) -> Self::Output {
125        self.0.map(|x| x + rhs).into()
126    }
127}
128
129impl<F: Field, const N: usize> AddAssign for FieldArray<F, N> {
130    #[inline]
131    fn add_assign(&mut self, rhs: Self) {
132        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x += y);
133    }
134}
135
136impl<F: Field, const N: usize> AddAssign<F> for FieldArray<F, N> {
137    #[inline]
138    fn add_assign(&mut self, rhs: F) {
139        self.0.iter_mut().for_each(|x| *x += rhs);
140    }
141}
142
143impl<F: Field, const N: usize> Sub for FieldArray<F, N> {
144    type Output = Self;
145
146    #[inline]
147    fn sub(self, rhs: Self) -> Self::Output {
148        array::from_fn(|i| self.0[i] - rhs.0[i]).into()
149    }
150}
151
152impl<F: Field, const N: usize> Sub<F> for FieldArray<F, N> {
153    type Output = Self;
154
155    #[inline]
156    fn sub(self, rhs: F) -> Self::Output {
157        self.0.map(|x| x - rhs).into()
158    }
159}
160
161impl<F: Field, const N: usize> SubAssign for FieldArray<F, N> {
162    #[inline]
163    fn sub_assign(&mut self, rhs: Self) {
164        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x -= y);
165    }
166}
167
168impl<F: Field, const N: usize> SubAssign<F> for FieldArray<F, N> {
169    #[inline]
170    fn sub_assign(&mut self, rhs: F) {
171        self.0.iter_mut().for_each(|x| *x -= rhs);
172    }
173}
174
175impl<F: Field, const N: usize> Neg for FieldArray<F, N> {
176    type Output = Self;
177
178    #[inline]
179    fn neg(self) -> Self::Output {
180        self.0.map(|x| -x).into()
181    }
182}
183
184impl<F: Field, const N: usize> Mul for FieldArray<F, N> {
185    type Output = Self;
186
187    #[inline]
188    fn mul(self, rhs: Self) -> Self::Output {
189        array::from_fn(|i| self.0[i] * rhs.0[i]).into()
190    }
191}
192
193impl<F: Field, const N: usize> Mul<F> for FieldArray<F, N> {
194    type Output = Self;
195
196    #[inline]
197    fn mul(self, rhs: F) -> Self::Output {
198        self.0.map(|x| x * rhs).into()
199    }
200}
201
202impl<F: Field, const N: usize> MulAssign for FieldArray<F, N> {
203    #[inline]
204    fn mul_assign(&mut self, rhs: Self) {
205        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x *= y);
206    }
207}
208
209impl<F: Field, const N: usize> MulAssign<F> for FieldArray<F, N> {
210    #[inline]
211    fn mul_assign(&mut self, rhs: F) {
212        self.0.iter_mut().for_each(|x| *x *= rhs);
213    }
214}
215
216impl<F: Field, const N: usize> Div<F> for FieldArray<F, N> {
217    type Output = Self;
218
219    #[allow(clippy::suspicious_arithmetic_impl)]
220    #[inline]
221    fn div(self, rhs: F) -> Self::Output {
222        let rhs_inv = rhs.inverse();
223        self * rhs_inv
224    }
225}
226
227impl<F: Field, const N: usize> Sum for FieldArray<F, N> {
228    #[inline]
229    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
230        iter.reduce(|lhs, rhs| lhs + rhs).unwrap_or(Self::ZERO)
231    }
232}
233
234impl<F: Field, const N: usize> Product for FieldArray<F, N> {
235    #[inline]
236    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
237        iter.reduce(|lhs, rhs| lhs * rhs).unwrap_or(Self::ONE)
238    }
239}