Skip to main content

p3_field/
field.rs

1use alloc::vec;
2use alloc::vec::Vec;
3use core::fmt::{Debug, Display};
4use core::hash::Hash;
5use core::iter::{Product, Sum, zip};
6use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
7use core::{array, slice};
8
9use num_bigint::BigUint;
10use p3_maybe_rayon::prelude::*;
11use p3_util::{flatten_to_base, iter_array_chunks_padded};
12use serde::Serialize;
13use serde::de::DeserializeOwned;
14
15use crate::exponentiation::bits_u64;
16use crate::integers::{QuotientMap, from_integer_types};
17use crate::packed::PackedField;
18use crate::{Dup, Packable, PackedFieldExtension, PackedValue};
19
20/// A commutative ring, `R`, with prime characteristic, `p`.
21///
22/// This permits elements like:
23/// - A single finite field element.
24/// - A symbolic expression which would evaluate to a field element.
25/// - An array of finite field elements.
26/// - A polynomial with coefficients in a finite field.
27///
28/// ### Mathematical Description
29///
30/// Mathematically, a commutative ring is a set of objects which supports an addition-like
31/// like operation, `+`, and a multiplication-like operation `*`.
32///
33/// Let `x, y, z` denote arbitrary elements of the set.
34///
35/// Then, an operation is addition-like if it satisfies the following properties:
36/// - Commutativity => `x + y = y + x`
37/// - Associativity => `x + (y + z) = (x + y) + z`
38/// - Unit => There exists an identity element `ZERO` satisfying `x + ZERO = x`.
39/// - Inverses => For every `x` there exists a unique inverse `(-x)` satisfying `x + (-x) = ZERO`
40///
41/// Similarly, an operation is multiplication-like if it satisfies the following properties:
42/// - Commutativity => `x * y = y * x`
43/// - Associativity => `x * (y * z) = (x * y) * z`
44/// - Unit => There exists an identity element `ONE` satisfying `x * ONE = x`.
45/// - Distributivity => The two operations `+` and `*` must together satisfy `x * (y + z) = (x * y) + (x * z)`
46///
47/// Unlike in the addition case, we do not require inverses to exist with respect to `*`.
48///
49/// The simplest examples of commutative rings are the integers (`ℤ`), and the integers mod `N` (`ℤ/N`).
50///
51/// The characteristic of a ring is the smallest positive integer `r` such that `0 = r . 1 = 1 + 1 + ... + 1 (r times)`.
52/// For example, the characteristic of the modulo ring `ℤ/N` is `N`.
53///
54/// Rings with prime characteristic are particularly special due to their close relationship with finite fields.
55pub trait PrimeCharacteristicRing:
56    Sized
57    + Default
58    + Dup
59    + Add<Output = Self>
60    + AddAssign
61    + Sub<Output = Self>
62    + SubAssign
63    + Neg<Output = Self>
64    + Mul<Output = Self>
65    + MulAssign
66    + Sum
67    + Product
68    + Debug
69{
70    /// The field `ℤ/p` where the characteristic of this ring is p.
71    type PrimeSubfield: PrimeField;
72
73    /// The additive identity of the ring.
74    ///
75    /// For every element `a` in the ring we require the following properties:
76    ///
77    /// `a + ZERO = ZERO + a = a,`
78    ///
79    /// `a + (-a) = (-a) + a = ZERO.`
80    const ZERO: Self;
81
82    /// The multiplicative identity of the ring.
83    ///
84    /// For every element `a` in the ring we require the following property:
85    ///
86    /// `a*ONE = ONE*a = a.`
87    const ONE: Self;
88
89    /// The element in the ring given by `ONE + ONE`.
90    ///
91    /// This is provided as a convenience as `TWO` occurs regularly in
92    /// the proving system. This also is slightly faster than computing
93    /// it via addition. Note that multiplication by `TWO` is discouraged.
94    /// Instead of `a * TWO` use `a.double()` which will be faster.
95    ///
96    /// If the field has characteristic 2 this is equal to ZERO.
97    const TWO: Self;
98
99    /// The element in the ring given by `-ONE`.
100    ///
101    /// This is provided as a convenience as `NEG_ONE` occurs regularly in
102    /// the proving system. This also is slightly faster than computing
103    /// it via negation. Note that where possible `NEG_ONE` should be absorbed
104    /// into mathematical operations. For example `a - b` will be faster
105    /// than `a + NEG_ONE * b` and similarly `(-b)` is faster than `NEG_ONE * b`.
106    ///
107    /// If the field has characteristic 2 this is equal to ONE.
108    const NEG_ONE: Self;
109
110    /// Embed an element of the prime field `ℤ/p` into the ring `R`.
111    ///
112    /// Given any element `[r] ∈ ℤ/p`, represented by an integer `r` between `0` and `p - 1`
113    /// `from_prime_subfield([r])` will be equal to:
114    ///
115    /// `Self::ONE + ... + Self::ONE (r times)`
116    #[must_use]
117    fn from_prime_subfield(f: Self::PrimeSubfield) -> Self;
118
119    /// Return `Self::ONE` if `b` is `true` and `Self::ZERO` if `b` is `false`.
120    #[must_use]
121    #[inline(always)]
122    fn from_bool(b: bool) -> Self {
123        // Some rings might reimplement this to avoid the branch.
124        if b { Self::ONE } else { Self::ZERO }
125    }
126
127    from_integer_types!(
128        u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize
129    );
130
131    /// The elementary function `double(a) = 2*a`.
132    ///
133    /// This function should be preferred over calling `a + a` or `TWO * a` as a faster implementation may be available for some rings.
134    /// If the field has characteristic 2 then this returns 0.
135    #[must_use]
136    #[inline(always)]
137    fn double(&self) -> Self {
138        self.dup() + self.dup()
139    }
140
141    /// The elementary function `halve(a) = a/2`.
142    ///
143    /// # Panics
144    /// The function will panic if the field has characteristic 2.
145    #[must_use]
146    #[inline]
147    fn halve(&self) -> Self {
148        // This must be overwritten by PrimeField implementations as this definition
149        // is circular when PrimeSubfield = Self. It should also be overwritten by
150        // most rings to avoid the multiplication.
151        let half = Self::from_prime_subfield(Self::PrimeSubfield::ONE.halve());
152        self.dup() * half
153    }
154
155    /// The elementary function `square(a) = a^2`.
156    ///
157    /// This function should be preferred over calling `a * a`, as a faster implementation may be available for some rings.
158    #[must_use]
159    #[inline(always)]
160    fn square(&self) -> Self {
161        self.dup() * self.dup()
162    }
163
164    /// The elementary function `cube(a) = a^3`.
165    ///
166    /// This function should be preferred over calling `a * a * a`, as a faster implementation may be available for some rings.
167    #[must_use]
168    #[inline(always)]
169    fn cube(&self) -> Self {
170        self.square() * self.dup()
171    }
172
173    /// Computes the arithmetic generalization of boolean `xor`.
174    ///
175    /// For boolean inputs, `x ^ y = x + y - 2 xy`.
176    #[must_use]
177    #[inline(always)]
178    fn xor(&self, y: &Self) -> Self {
179        self.dup() + y.dup() - self.dup() * y.dup().double()
180    }
181
182    /// Computes the arithmetic generalization of a triple `xor`.
183    ///
184    /// For boolean inputs `x ^ y ^ z = x + y + z - 2(xy + xz + yz) + 4xyz`.
185    #[must_use]
186    #[inline(always)]
187    fn xor3(&self, y: &Self, z: &Self) -> Self {
188        self.xor(y).xor(z)
189    }
190
191    /// Computes the arithmetic generalization of `andnot`.
192    ///
193    /// For boolean inputs `(!x) & y = (1 - x)y`.
194    #[must_use]
195    #[inline(always)]
196    fn andn(&self, y: &Self) -> Self {
197        (Self::ONE - self.dup()) * y.dup()
198    }
199
200    /// The vanishing polynomial for boolean values: `x * (x - 1)`.
201    ///
202    /// This is a polynomial of degree `2` that evaluates to `0` if the input is `0` or `1`.
203    /// If our space is a field, then this will be nonzero on all other inputs.
204    #[must_use]
205    #[inline(always)]
206    fn bool_check(&self) -> Self {
207        // Note: We could delegate to `andn`, but to maintain backwards
208        // compatible AIR definitions, we stick with `x * (x - 1)` here.
209        self.dup() * (self.dup() - Self::ONE)
210    }
211
212    /// Exponentiation by a `u64` power.
213    ///
214    /// This uses the standard square and multiply approach.
215    /// For specific powers regularly used and known in advance,
216    /// this will be slower than custom addition chain exponentiation.
217    #[must_use]
218    #[inline]
219    fn exp_u64(&self, power: u64) -> Self {
220        let mut current = self.dup();
221        let mut product = Self::ONE;
222
223        for j in 0..bits_u64(power) {
224            if (power >> j) & 1 != 0 {
225                product *= current.dup();
226            }
227            current = current.square();
228        }
229        product
230    }
231
232    /// Exponentiation by a small constant power.
233    ///
234    /// For a collection of small values we implement custom multiplication chain circuits which can be faster than the
235    /// simpler square and multiply approach.
236    ///
237    /// For large values this defaults back to `self.exp_u64`.
238    #[must_use]
239    #[inline(always)]
240    fn exp_const_u64<const POWER: u64>(&self) -> Self {
241        match POWER {
242            0 => Self::ONE,
243            1 => self.dup(),
244            2 => self.square(),
245            3 => self.cube(),
246            4 => self.square().square(),
247            5 => self.square().square() * self.dup(),
248            6 => self.square().cube(),
249            7 => {
250                let x2 = self.square();
251                let x3 = x2.dup() * self.dup();
252                let x4 = x2.square();
253                x3 * x4
254            }
255            _ => self.exp_u64(POWER),
256        }
257    }
258
259    /// The elementary function `exp_power_of_2(a, power_log) = a^{2^power_log}`.
260    ///
261    /// Computed via repeated squaring.
262    #[must_use]
263    #[inline]
264    fn exp_power_of_2(&self, power_log: usize) -> Self {
265        let mut res = self.dup();
266        for _ in 0..power_log {
267            res = res.square();
268        }
269        res
270    }
271
272    /// The elementary function `mul_2exp_u64(a, exp) = a * 2^{exp}`.
273    ///
274    /// Here `2^{exp}` is computed using the square and multiply approach.
275    #[must_use]
276    #[inline]
277    fn mul_2exp_u64(&self, exp: u64) -> Self {
278        // Some rings might want to reimplement this to avoid the
279        // exponentiations (and potentially even the multiplication).
280        self.dup() * Self::TWO.exp_u64(exp)
281    }
282
283    /// Divide by a given power of two. `div_2exp_u64(a, exp) = a/2^exp`
284    ///
285    /// # Panics
286    /// The function will panic if the field has characteristic 2.
287    #[must_use]
288    #[inline]
289    fn div_2exp_u64(&self, exp: u64) -> Self {
290        // Some rings might want to reimplement this to avoid the
291        // exponentiations (and potentially even the multiplication).
292        self.dup() * Self::from_prime_subfield(Self::PrimeSubfield::ONE.halve().exp_u64(exp))
293    }
294
295    /// Construct an iterator which returns powers of `self`: `self^0, self^1, self^2, ...`.
296    #[must_use]
297    #[inline]
298    fn powers(&self) -> Powers<Self> {
299        self.shifted_powers(Self::ONE)
300    }
301
302    /// Construct an iterator which returns powers of `self` shifted by `start`: `start, start*self^1, start*self^2, ...`.
303    #[must_use]
304    #[inline]
305    fn shifted_powers(&self, start: Self) -> Powers<Self> {
306        Powers {
307            base: self.dup(),
308            current: start,
309        }
310    }
311
312    /// Compute the dot product of two vectors.
313    ///
314    /// ```text
315    ///     result = u[0]*v[0] + u[1]*v[1] + ... + u[N-1]*v[N-1]
316    /// ```
317    ///
318    /// The products are combined with a balanced tree rather than a running accumulator.
319    /// A running accumulator makes every addition wait for the previous one to retire.
320    /// The tree keeps several partial sums in flight, so the latency chain is shorter.
321    ///
322    /// Rings whose modular reduction is linear over the accumulated representation
323    /// should override this to accumulate all `N` products unreduced and reduce once.
324    #[must_use]
325    #[inline]
326    fn dot_product<const N: usize>(u: &[Self; N], v: &[Self; N]) -> Self {
327        // Materialise the `N` products first, so none of the multiplies waits on a sum.
328        let products: [Self; N] = array::from_fn(|i| u[i].dup() * v[i].dup());
329
330        // Balanced tree of depth log2(N) instead of a linear chain of N - 1 adds.
331        Self::sum_array::<N>(&products)
332    }
333
334    /// Compute the sum of a slice of elements whose length is a compile time constant.
335    ///
336    /// The rust compiler doesn't realize that add is associative
337    /// so we help it out and minimize the dependency chains by hand.
338    /// Thus while this function has the same throughput as `input.iter().sum()`,
339    /// it will usually have much lower latency.
340    ///
341    /// # Panics
342    ///
343    /// May panic if the length of the input slice is not equal to `N`.
344    #[must_use]
345    #[inline]
346    fn sum_array<const N: usize>(input: &[Self]) -> Self {
347        // It looks a little strange but using a const parameter and an assert_eq! instead of
348        // using input.len() leads to a significant performance improvement.
349        // We could make this input &[Self; N] but that would require sticking .try_into().unwrap() everywhere.
350        // Checking godbolt, the compiler seems to unroll everything anyway.
351        assert_eq!(N, input.len());
352
353        // For `N <= 8` we implement a tree sum structure and for `N > 8` we break the input into
354        // chunks of `8`, perform a tree sum on each chunk and sum the results. The parameter `8`
355        // was determined experimentally by testing the speed of the poseidon2 internal layer computations.
356        // This is a useful benchmark as we have a mix of summations of size 15, 23 with other work in between.
357        // I only tested this on `AVX2` though so there might be a better value for other architectures.
358        match N {
359            0 => Self::ZERO,
360            1 => input[0].dup(),
361            2 => input[0].dup() + input[1].dup(),
362            3 => input[0].dup() + input[1].dup() + input[2].dup(),
363            4 => (input[0].dup() + input[1].dup()) + (input[2].dup() + input[3].dup()),
364            5 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<1>(&input[4..]),
365            6 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<2>(&input[4..]),
366            7 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<3>(&input[4..]),
367            8 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<4>(&input[4..]),
368            _ => {
369                // We know that N > 8 here so this saves an add over the usual
370                // initialisation of acc to Self::ZERO.
371                let mut acc = Self::sum_array::<8>(&input[..8]);
372                for i in (16..=N).step_by(8) {
373                    acc += Self::sum_array::<8>(&input[(i - 8)..i]);
374                }
375                // This would be much cleaner if we could use const generic expressions but
376                // this will do for now.
377                match N & 7 {
378                    0 => acc,
379                    1 => acc + Self::sum_array::<1>(&input[(8 * (N / 8))..]),
380                    2 => acc + Self::sum_array::<2>(&input[(8 * (N / 8))..]),
381                    3 => acc + Self::sum_array::<3>(&input[(8 * (N / 8))..]),
382                    4 => acc + Self::sum_array::<4>(&input[(8 * (N / 8))..]),
383                    5 => acc + Self::sum_array::<5>(&input[(8 * (N / 8))..]),
384                    6 => acc + Self::sum_array::<6>(&input[(8 * (N / 8))..]),
385                    7 => acc + Self::sum_array::<7>(&input[(8 * (N / 8))..]),
386                    _ => unreachable!(),
387                }
388            }
389        }
390    }
391
392    /// Allocates a vector of zero elements of length `len`. Many operating systems zero pages
393    /// before assigning them to a userspace process. In that case, our process should not need to
394    /// write zeros, which would be redundant. However, the compiler may not always recognize this.
395    ///
396    /// In particular, `vec![Self::ZERO; len]` appears to result in redundant userspace zeroing.
397    /// This is the default implementation, but implementers may wish to provide their own
398    /// implementation which transmutes something like `vec![0u32; len]`.
399    #[must_use]
400    #[inline]
401    fn zero_vec(len: usize) -> Vec<Self> {
402        vec![Self::ZERO; len]
403    }
404}
405
406/// A vector space `V` over `F` with a fixed basis. Fixing the basis allows elements of `V` to be
407/// converted to and from `DIMENSION` many elements of `F` which are interpreted as basis coefficients.
408///
409/// We usually expect `F` to be a field but do not enforce this and so allow it to be just a ring.
410/// This lets every ring implement `BasedVectorSpace<Self>` and is useful in a couple of other cases.
411///
412/// ## Safety
413/// We make no guarantees about consistency of the choice of basis across different versions of Plonky3.
414/// If this choice of basis changes, the behaviour of `BasedVectorSpace` will also change. Due to this,
415/// we recommend avoiding using this trait unless absolutely necessary.
416///
417/// ### Mathematical Description
418/// Given a vector space, `A` over `F`, a basis is a set of elements `B = {b_0, ..., b_{n-1}}`
419/// in `A` such that, given any element `a`, we can find a unique set of `n` elements of `F`,
420/// `f_0, ..., f_{n - 1}` satisfying `a = f_0 b_0 + ... + f_{n - 1} b_{n - 1}`. Thus the choice
421/// of `B` gives rise to a natural linear map between the vector space `A` and the canonical
422/// `n` dimensional vector space `F^n`.
423///
424/// This allows us to map between elements of `A` and arrays of `n` elements of `F`.
425/// Clearly this map depends entirely on the choice of basis `B` which may change
426/// across versions of Plonky3.
427///
428/// The situation is slightly more complicated in cases where `F` is not a field but boils down
429/// to an identical description once we enforce that `A` is a free module over `F`.
430pub trait BasedVectorSpace<F: PrimeCharacteristicRing>: Sized {
431    /// The dimension of the vector space, i.e. the number of elements in
432    /// its basis.
433    const DIMENSION: usize;
434
435    /// Fixes a basis for the algebra `A` and uses this to
436    /// map an element of `A` to a slice of `DIMENSION` `F` elements.
437    ///
438    /// # Safety
439    ///
440    /// The value produced by this function fundamentally depends
441    /// on the choice of basis. Care must be taken
442    /// to ensure portability if these values might ever be passed to
443    /// (or rederived within) another compilation environment where a
444    /// different basis might have been used.
445    #[must_use]
446    fn as_basis_coefficients_slice(&self) -> &[F];
447
448    /// Fixes a basis for the algebra `A` and uses this to
449    /// map `DIMENSION` `F` elements to an element of `A`.
450    ///
451    /// # Safety
452    ///
453    /// The value produced by this function fundamentally depends
454    /// on the choice of basis. Care must be taken
455    /// to ensure portability if these values might ever be passed to
456    /// (or rederived within) another compilation environment where a
457    /// different basis might have been used.
458    ///
459    /// Returns `None` if the length of the slice is different to `DIMENSION`.
460    #[must_use]
461    #[inline]
462    fn from_basis_coefficients_slice(slice: &[F]) -> Option<Self> {
463        Self::from_basis_coefficients_iter(slice.iter().cloned())
464    }
465
466    /// Fixes a basis for the algebra `A` and uses this to
467    /// map `DIMENSION` `F` elements to an element of `A`. Similar
468    /// to `core:array::from_fn`, the `DIMENSION` `F` elements are
469    /// given by `Fn(0), ..., Fn(DIMENSION - 1)` called in that order.
470    ///
471    /// # Safety
472    ///
473    /// The value produced by this function fundamentally depends
474    /// on the choice of basis. Care must be taken
475    /// to ensure portability if these values might ever be passed to
476    /// (or rederived within) another compilation environment where a
477    /// different basis might have been used.
478    #[must_use]
479    fn from_basis_coefficients_fn<Fn: FnMut(usize) -> F>(f: Fn) -> Self;
480
481    /// Fixes a basis for the algebra `A` and uses this to
482    /// map `DIMENSION` `F` elements to an element of `A`.
483    ///
484    /// # Safety
485    ///
486    /// The value produced by this function fundamentally depends
487    /// on the choice of basis. Care must be taken
488    /// to ensure portability if these values might ever be passed to
489    /// (or rederived within) another compilation environment where a
490    /// different basis might have been used.
491    ///
492    /// Returns `None` if the length of the iterator is different to `DIMENSION`.
493    #[must_use]
494    fn from_basis_coefficients_iter<I: ExactSizeIterator<Item = F>>(iter: I) -> Option<Self>;
495
496    /// Given a basis for the Algebra `A`, return the i'th basis element.
497    ///
498    /// # Safety
499    ///
500    /// The value produced by this function fundamentally depends
501    /// on the choice of basis. Care must be taken
502    /// to ensure portability if these values might ever be passed to
503    /// (or rederived within) another compilation environment where a
504    /// different basis might have been used.
505    ///
506    /// Returns `None` if `i` is greater than or equal to `DIMENSION`.
507    #[must_use]
508    #[inline]
509    fn ith_basis_element(i: usize) -> Option<Self> {
510        (i < Self::DIMENSION).then(|| Self::from_basis_coefficients_fn(|j| F::from_bool(i == j)))
511    }
512
513    /// Convert from a vector of `Self` to a vector of `F` by flattening the basis coefficients.
514    ///
515    /// Depending on the `BasedVectorSpace` this may be essentially a no-op and should certainly
516    /// be reimplemented in those cases.
517    ///
518    /// # Safety
519    ///
520    /// The value produced by this function fundamentally depends
521    /// on the choice of basis. Care must be taken
522    /// to ensure portability if these values might ever be passed to
523    /// (or rederived within) another compilation environment where a
524    /// different basis might have been used.
525    #[must_use]
526    #[inline]
527    fn flatten_to_base(vec: Vec<Self>) -> Vec<F> {
528        vec.into_iter()
529            .flat_map(|x| x.as_basis_coefficients_slice().to_vec())
530            .collect()
531    }
532
533    /// Convert from a vector of `F` to a vector of `Self` by combining the basis coefficients.
534    ///
535    /// Depending on the `BasedVectorSpace` this may be essentially a no-op and should certainly
536    /// be reimplemented in those cases.
537    ///
538    /// # Panics
539    /// This will panic if the length of `vec` is not a multiple of `Self::DIMENSION`.
540    ///
541    /// # Safety
542    ///
543    /// The value produced by this function fundamentally depends
544    /// on the choice of basis. Care must be taken
545    /// to ensure portability if these values might ever be passed to
546    /// (or rederived within) another compilation environment where a
547    /// different basis might have been used.
548    #[must_use]
549    #[inline]
550    fn reconstitute_from_base(vec: Vec<F>) -> Vec<Self>
551    where
552        F: Sync,
553        Self: Send,
554    {
555        assert_eq!(vec.len() % Self::DIMENSION, 0);
556
557        vec.par_chunks_exact(Self::DIMENSION)
558            .map(|chunk| {
559                Self::from_basis_coefficients_slice(chunk)
560                    .expect("Chunk length not equal to dimension")
561            })
562            .collect()
563    }
564}
565
566/// Compiler-independent identity of an algebra's ordered coefficient basis over `F`.
567///
568/// Together with the coefficient field identity and dimension, these bytes must
569/// uniquely identify the defining relations and the ordered basis used by
570/// [`BasedVectorSpace`]. Different polynomials or basis orderings require different
571/// identifiers, even for isomorphic fields. Use an unambiguous, versioned encoding;
572/// Rust type names, memory layouts, and compiler-dependent encodings are forbidden.
573/// Changing this identity changes Fiat-Shamir transcripts.
574pub trait AlgebraIdentity<F: PrimeCharacteristicRing>: BasedVectorSpace<F> {
575    /// Return the stable, unambiguous identifier of the defining relations and basis.
576    fn algebra_id() -> Vec<u8>;
577}
578
579impl<F: PrimeCharacteristicRing> AlgebraIdentity<F> for F {
580    fn algebra_id() -> Vec<u8> {
581        b"p3-scalar-basis-v1".to_vec()
582    }
583}
584
585/// Values that can act as sponge lanes for delimiter padding.
586///
587/// This is used by symmetric sponge adapters that need canonical `0` and `1` symbols while
588/// supporting both field/ring-based lanes and `u64`-based Keccak lanes behind one API.
589pub trait SpongePaddingValue: Copy {
590    /// The empty-lane value.
591    const PAD_ZERO: Self;
592
593    /// The delimiter value injected after the final absorbed element.
594    const PAD_ONE: Self;
595}
596
597impl<T: PrimeCharacteristicRing + Copy> SpongePaddingValue for T {
598    const PAD_ZERO: Self = Self::ZERO;
599    const PAD_ONE: Self = Self::ONE;
600}
601
602impl SpongePaddingValue for u64 {
603    const PAD_ZERO: Self = 0;
604    const PAD_ONE: Self = 1;
605}
606
607impl<const N: usize> SpongePaddingValue for [u64; N] {
608    const PAD_ZERO: Self = [0; N];
609    const PAD_ONE: Self = [1; N];
610}
611
612/// Trait for fields that support uniform bit sampling optimizations.
613pub trait UniformSamplingField {
614    /// Maximum number of bits we can sample at negligible (~1/field prime) probability of
615    /// triggering an error / requiring a resample.
616    const MAX_SINGLE_SAMPLE_BITS: usize;
617    /// An array storing the largest value `m_k` for each `k` in [0, 31], such that `m_k`
618    /// is a multiple of `2^k` and less than P. `m_k` is defined as:
619    ///
620    /// \( m_k = ⌊P / 2^k⌋ · 2^k \)
621    ///
622    /// This is used as a rejection sampling threshold (or error trigger), when sampling
623    /// random bits from uniformly sampled field elements. As long as we sample up to the `k`
624    /// least significant bits in the range [0, m_k), we sample from exactly `m_k` elements. As
625    /// `m_k` is divisible by 2^k, each of the least significant `k` bits has exactly the same
626    /// number of zeroes and ones, leading to a uniform sampling.
627    const SAMPLING_BITS_M: [u64; 64];
628}
629
630impl<F: PrimeCharacteristicRing> BasedVectorSpace<F> for F {
631    const DIMENSION: usize = 1;
632
633    #[inline]
634    fn as_basis_coefficients_slice(&self) -> &[F] {
635        slice::from_ref(self)
636    }
637
638    #[inline]
639    fn from_basis_coefficients_fn<Fn: FnMut(usize) -> F>(mut f: Fn) -> Self {
640        f(0)
641    }
642
643    #[inline]
644    fn from_basis_coefficients_iter<I: ExactSizeIterator<Item = F>>(mut iter: I) -> Option<Self> {
645        (iter.len() == 1).then(|| iter.next().unwrap()) // Unwrap will not panic as we know the length is 1.
646    }
647
648    #[inline]
649    fn flatten_to_base(vec: Vec<Self>) -> Vec<F> {
650        vec
651    }
652
653    #[inline]
654    fn reconstitute_from_base(vec: Vec<F>) -> Vec<Self> {
655        vec
656    }
657}
658
659/// A ring implements `InjectiveMonomial<N>` if the algebraic function
660/// `f(x) = x^N` is an injective map on elements of the ring.
661///
662/// We do not enforce that this map be invertible as there are useful
663/// cases such as polynomials or symbolic expressions where no inverse exists.
664///
665/// However, if the ring is a field with order `q` or an array of such field elements,
666/// then `f(x) = x^N` will be injective if and only if it is invertible and so in
667/// such cases this monomial acts as a permutation. Moreover, this will occur
668/// exactly when `N` and `q - 1` are relatively prime i.e. `gcd(N, q - 1) = 1`.
669pub trait InjectiveMonomial<const N: u64>: PrimeCharacteristicRing {
670    /// Compute `x -> x^n` for a given `n > 1` such that this
671    /// map is injective.
672    #[must_use]
673    #[inline]
674    fn injective_exp_n(&self) -> Self {
675        self.exp_const_u64::<N>()
676    }
677}
678
679/// A ring implements `PermutationMonomial<N>` if the algebraic function
680/// `f(x) = x^N` is invertible and thus acts as a permutation on elements of the ring.
681///
682/// In all cases we care about, this means that we can find another integer `K` such
683/// that `x = x^{NK}` for all elements of our ring.
684pub trait PermutationMonomial<const N: u64>: InjectiveMonomial<N> {
685    /// Compute `x -> x^K` for a given `K > 1` such that
686    /// `x^{NK} = x` for all elements `x`.
687    #[must_use]
688    fn injective_exp_root_n(&self) -> Self;
689}
690
691/// A ring `R` implements `Algebra<F>` if there is an injective homomorphism
692///  from `F` into `R`; in particular only `F::ZERO` maps to `R::ZERO`.
693///
694/// For the most part, we will usually expect `F` to be a field but there
695/// are a few cases where it is handy to allow it to just be a ring. In
696/// particular, every ring naturally implements `Algebra<Self>`.
697///
698/// ### Mathematical Description
699///
700/// Let `x` and `y` denote arbitrary elements of `F`. Then
701/// we require that our map `from` has the properties:
702/// - Preserves Identity: `from(F::ONE) = R::ONE`
703/// - Commutes with Addition: `from(x + y) = from(x) + from(y)`
704/// - Commutes with Multiplication: `from(x * y) = from(x) * from(y)`
705///
706/// Such maps are known as ring homomorphisms and are injective if the
707/// only element which maps to `R::ZERO` is `F::ZERO`.
708///
709/// The existence of this map makes `R` into an `F`-module and hence an `F`-algebra.
710/// If, additionally, `R` is a field, then this makes `R` a field extension of `F`.
711pub trait Algebra<F>:
712    PrimeCharacteristicRing
713    + From<F>
714    + Add<F, Output = Self>
715    + AddAssign<F>
716    + Sub<F, Output = Self>
717    + SubAssign<F>
718    + Mul<F, Output = Self>
719    + MulAssign<F>
720{
721    /// Square `a[0] + a[1] X` modulo `X^2 - w`.
722    ///
723    /// Returns `[a[0]^2 + w * a[1]^2, 2 * a[0] * a[1]]` for arbitrary `w`.
724    /// The default uses a dot product to share reduction work; packed algebras can
725    /// override it when dedicated squaring is cheaper than the two-product reduction.
726    #[must_use]
727    #[inline]
728    fn quadratic_extension_square(a: &[Self; 2], w: F) -> [Self; 2] {
729        let a1_w = a[1].dup() * w;
730        [
731            Self::dot_product(a, &[a[0].dup(), a1_w]),
732            a[0].dup() * a[1].double(),
733        ]
734    }
735
736    /// Dot product between algebra elements and base field scalars.
737    ///
738    /// Given arrays `a` (algebra) and `f` (scalars), computes:
739    ///
740    /// ```text
741    ///   result = a[0]*f[0] + a[1]*f[1] + ... + a[N-1]*f[N-1]
742    /// ```
743    ///
744    /// Uses a tree-structured summation to minimize dependency chains and
745    /// maximize throughput on pipelined architectures.
746    #[must_use]
747    #[inline]
748    fn mixed_dot_product<const N: usize>(a: &[Self; N], f: &[F; N]) -> Self
749    where
750        F: Dup,
751    {
752        let products: [Self; N] = core::array::from_fn(|i| a[i].dup() * f[i].dup());
753        Self::sum_array::<N>(&products)
754    }
755
756    /// Optimal chunk size for [`batched_linear_combination`](Self::batched_linear_combination).
757    ///
758    /// Override in implementations where a different chunk size is faster.
759    /// Must be one of 1, 2, 4, 8, 16, 32, or 64; other values cause a compile error.
760    const BATCHED_LC_CHUNK: usize = 8;
761
762    /// Runtime-length linear combination: `Σ values[i] * coeffs[i]`.
763    ///
764    /// Like [`mixed_dot_product`](Self::mixed_dot_product) but for slices whose
765    /// length is not known at compile time. Processes in chunks of
766    /// [`BATCHED_LC_CHUNK`](Self::BATCHED_LC_CHUNK), delegating each chunk to
767    /// `mixed_dot_product` to leverage SIMD-specialized overrides.
768    #[must_use]
769    #[inline]
770    fn batched_linear_combination(values: &[Self], coeffs: &[F]) -> Self
771    where
772        F: Dup,
773    {
774        const {
775            assert!(
776                matches!(Self::BATCHED_LC_CHUNK, 1 | 2 | 4 | 8 | 16 | 32 | 64),
777                "BATCHED_LC_CHUNK must be one of 1, 2, 4, 8, 16, 32, or 64"
778            );
779        }
780        match Self::BATCHED_LC_CHUNK {
781            1 => chunked_linear_combination::<1, Self, F>(values, coeffs),
782            2 => chunked_linear_combination::<2, Self, F>(values, coeffs),
783            4 => chunked_linear_combination::<4, Self, F>(values, coeffs),
784            8 => chunked_linear_combination::<8, Self, F>(values, coeffs),
785            16 => chunked_linear_combination::<16, Self, F>(values, coeffs),
786            32 => chunked_linear_combination::<32, Self, F>(values, coeffs),
787            64 => chunked_linear_combination::<64, Self, F>(values, coeffs),
788            _ => unreachable!(),
789        }
790    }
791}
792
793/// Compute `Σ values[i] * coeffs[i]` over `N` pairs.
794///
795/// A single long sum forces every add to wait for the previous one. Instead,
796/// we split the pairs into groups of `CHUNK`, sum each group on its own, and
797/// add up the group totals. Several partial sums run in parallel on the CPU,
798/// so the total latency is shorter than one straight chain.
799///
800/// The result is the same for every valid `CHUNK` — only the speed changes.
801///
802/// # Layout
803///
804/// For `N = q * CHUNK + r` with `0 <= r < CHUNK`:
805///
806/// ```text
807///     ┌── group 0 ──┬── group 1 ──┬─ ... ─┬── tail (r) ──┐
808///     │   CHUNK     │   CHUNK     │       │   r pairs    │
809///     └──────┬──────┴──────┬──────┴───────┴──────┬───────┘
810///            ▼             ▼                     ▼
811///       tree-sum      tree-sum             scalar adds
812///            └──► acc ◄────┴──────► acc ◄────────┘
813/// ```
814///
815/// # Panics
816///
817/// Compile-time panic if `CHUNK` is zero.
818#[must_use]
819#[inline]
820pub fn chunked_mixed_dot_product<
821    const CHUNK: usize,
822    A: Algebra<F> + Dup,
823    F: Dup,
824    const N: usize,
825>(
826    values: &[A; N],
827    coeffs: &[F; N],
828) -> A {
829    // CHUNK = 0 would make the group count undefined.
830    const { assert!(CHUNK != 0, "chunked_mixed_dot_product requires CHUNK > 0") }
831
832    // Fast path: N fits in one group → single balanced tree, no outer loop.
833    if N <= CHUNK {
834        let products: [A; N] = core::array::from_fn(|i| values[i].dup() * coeffs[i].dup());
835        return A::sum_array::<N>(&products);
836    }
837
838    // Split off q complete groups; r leftover pairs go to the tail.
839    let (val_chunks, val_rem) = values.as_slice().as_chunks::<CHUNK>();
840    let (coeff_chunks, coeff_rem) = coeffs.as_slice().as_chunks::<CHUNK>();
841    debug_assert_eq!(val_chunks.len(), coeff_chunks.len());
842
843    // One add per group; runs in parallel with the next group's multiplies.
844    let mut acc = A::ZERO;
845    for (vc, cc) in zip(val_chunks, coeff_chunks) {
846        let products: [A; CHUNK] = core::array::from_fn(|i| vc[i].dup() * cc[i].dup());
847        // Balanced tree of depth log2(CHUNK), folded into acc.
848        acc += A::sum_array::<CHUNK>(&products);
849    }
850
851    // Tail: at most CHUNK - 1 pairs as a serial multiply-add chain.
852    debug_assert_eq!(val_rem.len(), coeff_rem.len());
853    for (v, c) in zip(val_rem, coeff_rem) {
854        acc += v.dup() * c.dup();
855    }
856    acc
857}
858
859/// Lower a runtime chunk size into a const-generic call to the fixed-chunk dot product.
860///
861/// Each backend picks its preferred chunk size at runtime; the inner routine
862/// needs it as a const for unrolling. This wrapper bridges the gap.
863///
864/// Supported sizes: `1, 2, 4, 8, 16, 32, 64` — powers of two only, so the
865/// inner balanced tree stays balanced.
866///
867/// # Panics
868///
869/// Runtime panic if `chunk` is outside the supported set.
870#[must_use]
871#[inline]
872pub fn dispatch_chunked_mixed_dot_product<A: Algebra<F> + Dup, F: Dup, const N: usize>(
873    values: &[A; N],
874    coeffs: &[F; N],
875    chunk: usize,
876) -> A {
877    match chunk {
878        1 => chunked_mixed_dot_product::<1, A, F, N>(values, coeffs),
879        2 => chunked_mixed_dot_product::<2, A, F, N>(values, coeffs),
880        4 => chunked_mixed_dot_product::<4, A, F, N>(values, coeffs),
881        8 => chunked_mixed_dot_product::<8, A, F, N>(values, coeffs),
882        16 => chunked_mixed_dot_product::<16, A, F, N>(values, coeffs),
883        32 => chunked_mixed_dot_product::<32, A, F, N>(values, coeffs),
884        64 => chunked_mixed_dot_product::<64, A, F, N>(values, coeffs),
885        // Unsupported chunk = configuration bug in a backend.
886        _ => panic!("mixed_dot_product chunk must be one of 1, 2, 4, 8, 16, 32, or 64"),
887    }
888}
889
890/// Linear combination over runtime-length slices, processing in chunks of `CHUNK`.
891///
892/// Computes `Σ values[i] * coeffs[i]` by batching into fixed-size chunks and
893/// delegating each to [`Algebra::mixed_dot_product`], which SIMD implementations
894/// override with fused multiply-accumulate intrinsics.
895///
896/// This is the implementation backing [`Algebra::batched_linear_combination`].
897/// Use it directly when overriding that method with a different chunk size.
898#[must_use]
899#[inline]
900pub fn chunked_linear_combination<const CHUNK: usize, A: Algebra<F> + Dup, F: Dup>(
901    values: &[A],
902    coeffs: &[F],
903) -> A {
904    const { assert!(CHUNK != 0, "chunked_linear_combination requires CHUNK > 0") }
905    assert_eq!(values.len(), coeffs.len());
906
907    let (val_chunks, val_rem) = values.as_chunks::<CHUNK>();
908    let (coeff_chunks, coeff_rem) = coeffs.as_chunks::<CHUNK>();
909
910    debug_assert_eq!(val_chunks.len(), coeff_chunks.len());
911    let mut acc = A::ZERO;
912    for (vc, cc) in zip(val_chunks, coeff_chunks) {
913        acc += A::mixed_dot_product::<CHUNK>(vc, cc);
914    }
915
916    debug_assert_eq!(val_rem.len(), coeff_rem.len());
917    for (v, c) in zip(val_rem, coeff_rem) {
918        acc += v.dup() * c.dup();
919    }
920    acc
921}
922
923// Every ring is an algebra over itself.
924impl<R: PrimeCharacteristicRing> Algebra<R> for R {
925    #[inline]
926    fn mixed_dot_product<const N: usize>(a: &[Self; N], f: &[R; N]) -> Self {
927        // Scalars and algebra elements are the same type here, so the ring's own
928        // dot product accepts both sides unchanged.
929        //
930        //     mixed dot product over (R, R)  ==  R's own dot product
931        //
932        // That primitive is where delayed reduction lives: a ring that can accumulate
933        // `N` products in an unreduced representation reduces once instead of `N` times.
934        // Without this delegation the generic tile kernels would never see it.
935        Self::dot_product::<N>(a, f)
936    }
937}
938
939/// A collection of methods designed to help hash field elements.
940///
941/// Most fields will want to reimplement many/all of these methods as the default implementations
942/// are slow and involve converting to/from byte representations.
943pub trait RawDataSerializable: Sized {
944    /// The number of bytes which this field element occupies in memory.
945    /// Must be equal to the length of self.into_bytes().
946    const NUM_BYTES: usize;
947
948    /// Convert a field element into a collection of bytes.
949    #[must_use]
950    fn into_bytes(self) -> impl IntoIterator<Item = u8>;
951
952    /// Convert an iterator of field elements into an iterator of bytes.
953    #[must_use]
954    fn into_byte_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u8> {
955        input.into_iter().flat_map(|elem| elem.into_bytes())
956    }
957
958    /// Convert an iterator of field elements into an iterator of u32s.
959    ///
960    /// If `NUM_BYTES` does not divide `4`, multiple `F`s may be packed together to make a single `u32`. Furthermore,
961    /// if `NUM_BYTES * input.len()` does not divide `4`, the final `u32` will involve padding bytes which are set to `0`.
962    #[must_use]
963    fn into_u32_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u32> {
964        let bytes = Self::into_byte_stream(input);
965        iter_array_chunks_padded(bytes, 0).map(u32::from_le_bytes)
966    }
967
968    /// Convert an iterator of field elements into an iterator of u64s.
969    ///
970    /// If `NUM_BYTES` does not divide `8`, multiple `F`s may be packed together to make a single `u64`. Furthermore,
971    /// if `NUM_BYTES * input.len()` does not divide `8`, the final `u64` will involve padding bytes which are set to `0`.
972    #[must_use]
973    fn into_u64_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u64> {
974        let bytes = Self::into_byte_stream(input);
975        iter_array_chunks_padded(bytes, 0).map(u64::from_le_bytes)
976    }
977
978    /// Convert an iterator of field element arrays into an iterator of byte arrays.
979    ///
980    /// Converts an element `[F; N]` into the byte array `[[u8; N]; NUM_BYTES]`. This is
981    /// intended for use with vectorized hash functions which use vector operations
982    /// to compute several hashes in parallel.
983    #[must_use]
984    fn into_parallel_byte_streams<const N: usize>(
985        input: impl IntoIterator<Item = [Self; N]>,
986    ) -> impl IntoIterator<Item = [u8; N]> {
987        input.into_iter().flat_map(|vector| {
988            let bytes = vector.map(|elem| elem.into_bytes().into_iter().collect::<Vec<_>>());
989            (0..Self::NUM_BYTES).map(move |i| array::from_fn(|j| bytes[j][i]))
990        })
991    }
992
993    /// Convert an iterator of field element arrays into an iterator of u32 arrays.
994    ///
995    /// Converts an element `[F; N]` into the u32 array `[[u32; N]; NUM_BYTES/4]`. This is
996    /// intended for use with vectorized hash functions which use vector operations
997    /// to compute several hashes in parallel.
998    ///
999    /// This function is guaranteed to be equivalent to starting with `Iterator<[F; N]>` performing a transpose
1000    /// operation to get `[Iterator<F>; N]`, calling `into_u32_stream` on each element to get `[Iterator<u32>; N]` and then
1001    /// performing another transpose operation to get `Iterator<[u32; N]>`.
1002    ///
1003    /// If `NUM_BYTES` does not divide `4`, multiple `[F; N]`s may be packed together to make a single `[u32; N]`. Furthermore,
1004    /// if `NUM_BYTES * input.len()` does not divide `4`, the final `[u32; N]` will involve padding bytes which are set to `0`.
1005    #[must_use]
1006    fn into_parallel_u32_streams<const N: usize>(
1007        input: impl IntoIterator<Item = [Self; N]>,
1008    ) -> impl IntoIterator<Item = [u32; N]> {
1009        let bytes = Self::into_parallel_byte_streams(input);
1010        iter_array_chunks_padded(bytes, [0; N]).map(|byte_array: [[u8; N]; 4]| {
1011            array::from_fn(|i| u32::from_le_bytes(array::from_fn(|j| byte_array[j][i])))
1012        })
1013    }
1014
1015    /// Convert an iterator of field element arrays into an iterator of u64 arrays.
1016    ///
1017    /// Converts an element `[F; N]` into the u64 array `[[u64; N]; NUM_BYTES/8]`. This is
1018    /// intended for use with vectorized hash functions which use vector operations
1019    /// to compute several hashes in parallel.
1020    ///
1021    /// This function is guaranteed to be equivalent to starting with `Iterator<[F; N]>` performing a transpose
1022    /// operation to get `[Iterator<F>; N]`, calling `into_u64_stream` on each element to get `[Iterator<u64>; N]` and then
1023    /// performing another transpose operation to get `Iterator<[u64; N]>`.
1024    ///
1025    /// If `NUM_BYTES` does not divide `8`, multiple `[F; N]`s may be packed together to make a single `[u64; N]`. Furthermore,
1026    /// if `NUM_BYTES * input.len()` does not divide `8`, the final `[u64; N]` will involve padding bytes which are set to `0`.
1027    #[must_use]
1028    fn into_parallel_u64_streams<const N: usize>(
1029        input: impl IntoIterator<Item = [Self; N]>,
1030    ) -> impl IntoIterator<Item = [u64; N]> {
1031        let bytes = Self::into_parallel_byte_streams(input);
1032        iter_array_chunks_padded(bytes, [0; N]).map(|byte_array: [[u8; N]; 8]| {
1033            array::from_fn(|i| u64::from_le_bytes(array::from_fn(|j| byte_array[j][i])))
1034        })
1035    }
1036}
1037
1038/// A field `F`. This permits both modular fields `ℤ/p` along with their field extensions.
1039///
1040/// A ring is a field if every element `x` has a unique multiplicative inverse `x^{-1}`
1041/// which satisfies `x * x^{-1} = F::ONE`.
1042pub trait Field:
1043    Algebra<Self>
1044    + RawDataSerializable
1045    + Packable
1046    + 'static
1047    + Copy
1048    + Div<Self, Output = Self>
1049    + DivAssign
1050    + Add<Self::Packing, Output = Self::Packing>
1051    + Sub<Self::Packing, Output = Self::Packing>
1052    + Mul<Self::Packing, Output = Self::Packing>
1053    + Eq
1054    + Hash
1055    + Send
1056    + Sync
1057    + Display
1058    + Serialize
1059    + DeserializeOwned
1060{
1061    type Packing: PackedField<Scalar = Self>;
1062
1063    /// A generator of this field's multiplicative group.
1064    const GENERATOR: Self;
1065
1066    /// Whether evaluating multiple packed vectors of this field in lockstep (to overlap
1067    /// independent dependency chains and hide packed-multiplication latency) is expected
1068    /// to help throughput for this field.
1069    ///
1070    /// Only [`p3_uni_stark::quotient_values`](https://docs.rs/p3-uni-stark)'s `aarch64`
1071    /// (`neon`)-gated path reads this constant; on every other target it has no effect,
1072    /// so leaving it at the default is always safe there.
1073    ///
1074    /// Defaults to `false`, so fields fail safe into the plain (non-lockstep) path unless
1075    /// explicitly measured to benefit. Override to `true` only once benchmarks confirm the
1076    /// field's packed multiplication is latency-bound enough for lockstep evaluation to help.
1077    const BENEFITS_FROM_LOCKSTEP_EVALUATION: bool = false;
1078
1079    /// Check if the given field element is equal to the unique additive identity (ZERO).
1080    #[must_use]
1081    #[inline]
1082    fn is_zero(&self) -> bool {
1083        *self == Self::ZERO
1084    }
1085
1086    /// Check if the given field element is equal to the unique multiplicative identity (ONE).
1087    #[must_use]
1088    #[inline]
1089    fn is_one(&self) -> bool {
1090        *self == Self::ONE
1091    }
1092
1093    /// The multiplicative inverse of this field element, if it exists.
1094    ///
1095    /// NOTE: The inverse of `0` is undefined and will return `None`.
1096    #[must_use]
1097    fn try_inverse(&self) -> Option<Self>;
1098
1099    /// The multiplicative inverse of this field element.
1100    ///
1101    /// # Panics
1102    /// The function will panic if the field element is `0`.
1103    /// Use try_inverse if you want to handle this case.
1104    #[must_use]
1105    fn inverse(&self) -> Self {
1106        self.try_inverse().expect("Tried to invert zero")
1107    }
1108
1109    /// A square root of this field element, if one exists.
1110    ///
1111    /// Returns `Some(r)` with `r * r == *self` when this element is a quadratic
1112    /// residue, and `None` when it is a quadratic non-residue. `ZERO` returns
1113    /// `Some(ZERO)`. When two square roots exist, which one is returned is
1114    /// unspecified.
1115    ///
1116    /// The default implementation uses the Tonelli–Shanks algorithm. Fields with
1117    /// a more direct formula (e.g. those with `|F| ≡ 3 mod 4`) may override it.
1118    #[must_use]
1119    fn try_sqrt(&self) -> Option<Self> {
1120        crate::sqrt::tonelli_shanks(*self)
1121    }
1122
1123    /// The `i`-th element of a fixed injective enumeration of `Self`, used as an
1124    /// interpolation node. Must satisfy `interpolation_node(0) == ZERO` and
1125    /// `interpolation_node(1) == ONE`, and be injective for every `i` below the size
1126    /// of the field — no enumeration can do better, and a field smaller than the
1127    /// degree of the polynomial being interpolated is unusable for that protocol
1128    /// anyway. Round-polynomial degrees are tiny, so `0..min(64, |Self|)` is the
1129    /// tested range.
1130    ///
1131    /// The default maps `i` through the prime subfield and is injective only while
1132    /// `i` is below the characteristic. Fields of characteristic below `2^32` must
1133    /// override it.
1134    #[must_use]
1135    fn interpolation_node(i: usize) -> Self {
1136        Self::from_usize(i)
1137    }
1138
1139    /// Add two slices of field elements together, returning the result in the first slice.
1140    ///
1141    /// Makes use of packing to speed up the addition.
1142    ///
1143    /// This is optimal for cases where the two slices are small to medium length. E.g. between
1144    /// `F::Packing::WIDTH` and roughly however many elements fit in a cache line.
1145    ///
1146    /// For larger slices, it's likely worthwhile to use parallelization before calling this.
1147    /// Similarly if you need to add a large number of slices together, it's best to
1148    /// break them into small chunks and call this on the smaller chunks.
1149    ///
1150    /// # Panics
1151    /// The function will panic if the lengths of the two slices are not equal.
1152    #[inline]
1153    fn add_slices(slice_1: &mut [Self], slice_2: &[Self]) {
1154        let (shorts_1, suffix_1) = Self::Packing::pack_slice_with_suffix_mut(slice_1);
1155        let (shorts_2, suffix_2) = Self::Packing::pack_slice_with_suffix(slice_2);
1156        debug_assert_eq!(shorts_1.len(), shorts_2.len());
1157        debug_assert_eq!(suffix_1.len(), suffix_2.len());
1158        for (x_1, &x_2) in shorts_1.iter_mut().zip(shorts_2) {
1159            *x_1 += x_2;
1160        }
1161        for (x_1, &x_2) in suffix_1.iter_mut().zip(suffix_2) {
1162            *x_1 += x_2;
1163        }
1164    }
1165
1166    /// Accumulate `acc[c * N + j] += scales[j] * row[c]` over a stream of packed rows.
1167    ///
1168    /// Each item provides one matrix row as `acc.len() / N` packed base-field words,
1169    /// together with the row's `N` extension-field weights. `acc` is laid out with the
1170    /// `N` weights of each word group adjacent, and its length must be a multiple of `N`.
1171    ///
1172    /// This is the inner kernel of batched columnwise (weighted-sum-of-rows) dot
1173    /// products. Fields may override it to defer modular reductions across rows.
1174    fn batched_columnwise_dot_product<EF, R, I, const N: usize>(
1175        acc: &mut [EF::ExtensionPacking],
1176        items: I,
1177    ) where
1178        EF: ExtensionField<Self>,
1179        R: Iterator<Item = Self::Packing>,
1180        I: Iterator<Item = (R, [EF; N])>,
1181    {
1182        generic_batched_columnwise_dot_product::<Self, EF, R, I, N>(acc, items);
1183    }
1184
1185    /// The number of elements in the field.
1186    ///
1187    /// This will either be prime if the field is a PrimeField or a power of a
1188    /// prime if the field is an extension field.
1189    #[must_use]
1190    fn order() -> BigUint;
1191
1192    /// The number of bits required to define an element of this field.
1193    ///
1194    /// Usually due to storage and practical reasons the memory size of
1195    /// a field element will be a little larger than bits().
1196    #[must_use]
1197    #[inline]
1198    fn bits() -> usize {
1199        Self::order().bits() as usize
1200    }
1201}
1202
1203/// The generic accumulation behind [`Field::batched_columnwise_dot_product`]:
1204/// `acc[c * N + j] += scales[j] * row[c]` over a stream of packed rows.
1205///
1206/// Kept as a free function so that specialized `Field` implementations can fall back
1207/// to it for extension degrees their kernels do not cover.
1208pub fn generic_batched_columnwise_dot_product<F, EF, R, I, const N: usize>(
1209    acc: &mut [EF::ExtensionPacking],
1210    items: I,
1211) where
1212    F: Field,
1213    EF: ExtensionField<F>,
1214    R: Iterator<Item = F::Packing>,
1215    I: Iterator<Item = (R, [EF; N])>,
1216{
1217    for (row, scales) in items {
1218        let packed_scales = scales.map(EF::ExtensionPacking::from);
1219        for (acc_c, r) in acc.as_chunks_mut::<N>().0.iter_mut().zip(row) {
1220            for (a, &s) in acc_c.iter_mut().zip(&packed_scales) {
1221                *a += s * r;
1222            }
1223        }
1224    }
1225}
1226
1227/// A field isomorphic to `ℤ/p` for some prime `p`.
1228///
1229/// There is a natural map from `ℤ` to `ℤ/p` which sends an integer `r` to its conjugacy class `[r]`.
1230/// Canonically, each conjugacy class `[r]` can be represented by the unique integer `s` in `[0, p - 1)`
1231/// satisfying `s = r mod p`. This however is often not the most convenient computational representation
1232/// and so internal representations of field elements might differ from this and may change over time.
1233pub trait PrimeField:
1234    Field
1235    + Ord
1236    + QuotientMap<u8>
1237    + QuotientMap<u16>
1238    + QuotientMap<u32>
1239    + QuotientMap<u64>
1240    + QuotientMap<u128>
1241    + QuotientMap<usize>
1242    + QuotientMap<i8>
1243    + QuotientMap<i16>
1244    + QuotientMap<i32>
1245    + QuotientMap<i64>
1246    + QuotientMap<i128>
1247    + QuotientMap<isize>
1248{
1249    /// Return the representative of `value` in canonical form
1250    /// which lies in the range `0 <= x < self.order()`.
1251    #[must_use]
1252    fn as_canonical_biguint(&self) -> BigUint;
1253}
1254
1255/// A prime field `ℤ/p` with order, `p < 2^64`.
1256pub trait PrimeField64: PrimeField {
1257    const ORDER_U64: u64;
1258
1259    /// Return the representative of `value` in canonical form
1260    /// which lies in the range `0 <= x < ORDER_U64`.
1261    #[must_use]
1262    fn as_canonical_u64(&self) -> u64;
1263
1264    /// Convert a field element to a `u64` such that any two field elements
1265    /// are converted to the same `u64` if and only if they represent the same value.
1266    ///
1267    /// This will be the fastest way to convert a field element to a `u64` and
1268    /// is intended for use in hashing. It will also be consistent across different targets.
1269    #[must_use]
1270    #[inline(always)]
1271    fn to_unique_u64(&self) -> u64 {
1272        // A simple default which is optimal for some fields.
1273        self.as_canonical_u64()
1274    }
1275}
1276
1277/// A prime field `ℤ/p` with order `p < 2^32`.
1278pub trait PrimeField32: PrimeField64 {
1279    const ORDER_U32: u32;
1280
1281    /// Return the representative of `value` in canonical form
1282    /// which lies in the range `0 <= x < ORDER_U64`.
1283    #[must_use]
1284    fn as_canonical_u32(&self) -> u32;
1285
1286    /// Convert a field element to a `u32` such that any two field elements
1287    /// are converted to the same `u32` if and only if they represent the same value.
1288    ///
1289    /// This will be the fastest way to convert a field element to a `u32` and
1290    /// is intended for use in hashing. It will also be consistent across different targets.
1291    #[must_use]
1292    #[inline(always)]
1293    fn to_unique_u32(&self) -> u32 {
1294        // A simple default which is optimal for some fields.
1295        self.as_canonical_u32()
1296    }
1297}
1298
1299/// A field `EF` which is also an algebra over a field `F`.
1300///
1301/// This provides a couple of convenience methods on top of the
1302/// standard methods provided by `Field`, `Algebra<F>` and `BasedVectorSpace<F>`.
1303///
1304/// It also provides a type which handles packed vectors of extension field elements.
1305pub trait ExtensionField<Base: Field>:
1306    Field + Algebra<Base> + BasedVectorSpace<Base> + AlgebraIdentity<Base>
1307{
1308    type ExtensionPacking: PackedFieldExtension<Base, Self> + 'static + Copy + Send + Sync;
1309
1310    /// Determine if the given element lies in the base field.
1311    #[must_use]
1312    fn is_in_basefield(&self) -> bool;
1313
1314    /// If the element lies in the base field project it down.
1315    /// Otherwise return None.
1316    #[must_use]
1317    fn as_base(&self) -> Option<Base>;
1318
1319    /// Reassemble an element of `Self` from `D = DIMENSION` coefficients in `Self`
1320    /// via `Σⱼ basisⱼ · coeffsⱼ`. Returns `None` if `coeffs.len() != Self::DIMENSION`.
1321    ///
1322    /// This is the `Self`-coefficient counterpart to
1323    /// [`BasedVectorSpace::from_basis_coefficients_slice`], which takes coefficients
1324    /// in `Base`. It is the natural "lifting" operation in commit-and-open protocols:
1325    /// if an extension polynomial decomposes as `f(X) = Σⱼ basisⱼ · fⱼ(X)` with
1326    /// `fⱼ` over `Base`, then `f(z) = Σⱼ basisⱼ · fⱼ(z)` for any `z ∈ Self`.
1327    #[inline]
1328    #[must_use]
1329    fn from_ext_basis_coefficients(coeffs: &[Self]) -> Option<Self> {
1330        (coeffs.len() == Self::DIMENSION).then(|| {
1331            (0..Self::DIMENSION)
1332                .map(|j| Self::ith_basis_element(j).unwrap() * coeffs[j])
1333                .sum()
1334        })
1335    }
1336}
1337
1338// Every field is trivially a one dimensional extension over itself.
1339impl<F: Field> ExtensionField<F> for F {
1340    type ExtensionPacking = F::Packing;
1341
1342    #[inline]
1343    fn is_in_basefield(&self) -> bool {
1344        true
1345    }
1346
1347    #[inline]
1348    fn as_base(&self) -> Option<F> {
1349        Some(*self)
1350    }
1351
1352    #[inline]
1353    fn from_ext_basis_coefficients(coeffs: &[Self]) -> Option<Self> {
1354        (coeffs.len() == 1).then(|| coeffs[0])
1355    }
1356}
1357
1358/// A field which supplies information like the two-adicity of its multiplicative group, and methods
1359/// for obtaining two-adic generators.
1360pub trait TwoAdicField: Field {
1361    /// The number of factors of two in this field's multiplicative group.
1362    const TWO_ADICITY: usize;
1363
1364    /// Returns a generator of the multiplicative group of order `2^bits`.
1365    /// Assumes `bits <= TWO_ADICITY`, otherwise the result is undefined.
1366    #[must_use]
1367    fn two_adic_generator(bits: usize) -> Self;
1368}
1369
1370/// An iterator which returns the powers of a base element `b` shifted by current `c`: `c, c * b, c * b^2, ...`.
1371#[derive(Clone, Debug)]
1372pub struct Powers<R: PrimeCharacteristicRing> {
1373    pub base: R,
1374    pub current: R,
1375}
1376
1377impl<R: PrimeCharacteristicRing> Iterator for Powers<R> {
1378    type Item = R;
1379
1380    fn next(&mut self) -> Option<R> {
1381        let result = self.current.dup();
1382        self.current *= self.base.dup();
1383        Some(result)
1384    }
1385}
1386
1387impl<R: PrimeCharacteristicRing> Powers<R> {
1388    /// Returns an iterator yielding the first `n` powers.
1389    #[inline]
1390    #[must_use]
1391    pub const fn take(self, n: usize) -> BoundedPowers<R> {
1392        BoundedPowers { iter: self, n }
1393    }
1394
1395    /// Fills `slice` with the next `slice.len()` powers yielded by the iterator.
1396    #[inline]
1397    pub fn fill(self, slice: &mut [R]) {
1398        slice
1399            .iter_mut()
1400            .zip(self)
1401            .for_each(|(out, next)| *out = next);
1402    }
1403}
1404
1405impl<F: Field> Powers<F> {
1406    /// Wrapper for `self.take(n).collect()`.
1407    ///
1408    /// Bounded to `F: Field` on purpose: the body resolves `.collect()` to the inherent
1409    /// [`BoundedPowers::collect`] SIMD fast path, which only exists under `F: Field`.
1410    /// Defining this method under a wider bound (e.g. `PrimeCharacteristicRing`) would
1411    /// silently fall back to `Iterator::collect` and bypass packed-field acceleration.
1412    #[inline]
1413    #[must_use]
1414    pub fn collect_n(self, n: usize) -> Vec<F> {
1415        self.take(n).collect()
1416    }
1417}
1418
1419impl<F: Field> BoundedPowers<F> {
1420    /// Collect exactly `num_powers` ascending powers of `self.base`, starting at `self.current`.
1421    ///
1422    /// # Details
1423    ///
1424    /// Each chunk is computed using packed fields.
1425    ///
1426    /// The shared task-size policy picks the chunk length.
1427    ///
1428    /// A request too short to pay for a dispatch is filled on the calling thread.
1429    ///
1430    /// # Performance
1431    ///
1432    /// Enable the `parallel` feature to enable parallelization.
1433    #[must_use]
1434    pub fn collect(self) -> Vec<F> {
1435        let num_powers = self.n;
1436
1437        // When num_powers is small, fallback to serial computation
1438        if num_powers < 16 {
1439            return self.take(num_powers).collect();
1440        }
1441
1442        // Allocate buffer storing packed powers, containing at least `num_powers` scalars.
1443        let width = F::Packing::WIDTH;
1444        let num_packed = num_powers.div_ceil(width);
1445        let mut points_packed = F::Packing::zero_vec(num_packed);
1446
1447        let base = self.iter.base;
1448        let shift = self.iter.current;
1449
1450        // One item writes a packed element and reads none.
1451        // It therefore moves exactly that element's width.
1452        //
1453        // One number answers both questions this fill has:
1454        //
1455        //     chunk >= total  ->  fill the whole buffer on this thread
1456        //     chunk <  total  ->  packed elements one parallel chunk holds
1457        let chunk_size = min_task_len(num_packed, size_of::<F::Packing>());
1458
1459        if chunk_size >= num_packed {
1460            fill_packed_shifted_powers(base, shift, &mut points_packed);
1461        } else {
1462            // Precompute base for each chunk.
1463            let chunk_base = base.exp_u64((chunk_size * width) as u64);
1464
1465            points_packed
1466                .par_chunks_mut(chunk_size)
1467                .enumerate()
1468                .for_each(|(chunk_idx, chunk_slice)| {
1469                    // First power in this chunk
1470                    let chunk_start = shift * chunk_base.exp_u64(chunk_idx as u64);
1471
1472                    // Fill the chunk with packed powers.
1473                    fill_packed_shifted_powers(base, chunk_start, chunk_slice);
1474                });
1475        }
1476
1477        // return the number of requested points, discarding the unused packed powers
1478        // SAFETY: size_of::<F::Packing> always divides size_of::<F::Packing>.
1479        let mut points = unsafe { flatten_to_base(points_packed) };
1480        points.truncate(num_powers);
1481        points
1482    }
1483}
1484
1485/// Number of independent multiplication chains advanced together by [`fill_packed_shifted_powers`].
1486const NUM_POWER_CHAINS: usize = 8;
1487
1488/// Fill `out` with `start, start * base, start * base^2, ...` packed into `P`.
1489///
1490/// Packed output `i` comes from chain `i % NUM_POWER_CHAINS`, and each chain steps by
1491/// `base^(NUM_POWER_CHAINS * P::WIDTH)`. The chains are independent of each other, so
1492/// their multiplications can overlap in the pipeline. Outputs shorter than two rounds
1493/// are filled by a single chain, as setting up the others would cost more than it saves.
1494fn fill_packed_shifted_powers<P: PackedField>(base: P::Scalar, start: P::Scalar, out: &mut [P]) {
1495    let mut powers = P::packed_shifted_powers(base, start);
1496    if out.len() < 2 * NUM_POWER_CHAINS {
1497        powers.fill(out);
1498        return;
1499    }
1500
1501    let mut chains: [P; NUM_POWER_CHAINS] = array::from_fn(|_| powers.next().unwrap());
1502    let step: P = base.exp_u64((NUM_POWER_CHAINS * P::WIDTH) as u64).into();
1503
1504    let (rounds, tail) = out.as_chunks_mut::<NUM_POWER_CHAINS>();
1505    for round in rounds {
1506        *round = chains;
1507        for chain in &mut chains {
1508            *chain *= step;
1509        }
1510    }
1511    tail.copy_from_slice(&chains[..tail.len()]);
1512}
1513
1514/// Same as [`Powers`], but returns a bounded number of powers.
1515#[derive(Clone, Debug)]
1516pub struct BoundedPowers<R: PrimeCharacteristicRing> {
1517    iter: Powers<R>,
1518    n: usize,
1519}
1520
1521impl<R: PrimeCharacteristicRing> Iterator for BoundedPowers<R> {
1522    type Item = R;
1523
1524    fn next(&mut self) -> Option<R> {
1525        (self.n != 0).then(|| {
1526            self.n -= 1;
1527            self.iter.next().unwrap()
1528        })
1529    }
1530
1531    #[inline]
1532    fn size_hint(&self) -> (usize, Option<usize>) {
1533        (self.n, Some(self.n))
1534    }
1535}
1536
1537impl<R: PrimeCharacteristicRing> ExactSizeIterator for BoundedPowers<R> {
1538    #[inline]
1539    fn len(&self) -> usize {
1540        self.n
1541    }
1542}