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}