twenty_first/math/
traits.rs

1use std::fmt::Debug;
2use std::fmt::Display;
3use std::hash::Hash;
4use std::ops::Add;
5use std::ops::AddAssign;
6use std::ops::Div;
7use std::ops::Mul;
8use std::ops::MulAssign;
9use std::ops::Neg;
10use std::ops::Sub;
11use std::ops::SubAssign;
12
13use num_traits::ConstOne;
14use num_traits::ConstZero;
15use num_traits::Zero;
16use rayon::prelude::*;
17use serde::Serialize;
18use serde::de::DeserializeOwned;
19
20pub trait CyclicGroupGenerator
21where
22    Self: Sized,
23{
24    fn get_cyclic_group_elements(&self, max: Option<usize>) -> Vec<Self>;
25}
26
27// TODO: Assert if we're risking inverting 0 at any point.
28pub trait Inverse
29where
30    Self: Sized + Zero,
31{
32    /// The multiplicative inverse: `a * a.inverse() == 1`
33    ///
34    /// # Panics
35    ///
36    /// Panics if `self` does not have a multiplicative inverse, for example, when
37    /// `self` is zero. (For fields, this is the only case.)
38    fn inverse(&self) -> Self;
39
40    fn inverse_or_zero(&self) -> Self {
41        if self.is_zero() {
42            Self::zero()
43        } else {
44            self.inverse()
45        }
46    }
47}
48
49pub trait PrimitiveRootOfUnity
50where
51    Self: Sized,
52{
53    fn primitive_root_of_unity(n: u64) -> Option<Self>;
54}
55
56pub trait ModPowU64 {
57    #[must_use]
58    fn mod_pow_u64(&self, pow: u64) -> Self;
59}
60
61pub trait ModPowU32 {
62    #[must_use]
63    fn mod_pow_u32(&self, exp: u32) -> Self;
64}
65
66pub trait FiniteField:
67    Copy
68    + Debug
69    + Display
70    + Eq
71    + Serialize
72    + DeserializeOwned
73    + Hash
74    + ConstZero
75    + ConstOne
76    + Add<Output = Self>
77    + Mul<Output = Self>
78    + Sub<Output = Self>
79    + Div<Output = Self>
80    + Neg<Output = Self>
81    + AddAssign
82    + MulAssign
83    + SubAssign
84    + CyclicGroupGenerator
85    + PrimitiveRootOfUnity
86    + Inverse
87    + ModPowU32
88    + From<u64>
89    + Send
90    + Sync
91{
92    /// If values of this type are laid out in memory as a run of
93    /// [`BFieldElement`]s, the number of base field elements in that run.
94    ///
95    /// [`BFieldElement`]: crate::math::b_field_element::BFieldElement
96    /// Types with any other layout must leave this at the default `None`.
97    ///
98    /// SIMD kernels, like the NTT's, use this to process the base field limbs
99    /// of many elements at once. Setting it for a type whose layout does not
100    /// match the description is unsound.
101    const NUM_BFE_LIMBS: Option<usize> = None;
102
103    /// Montgomery Batch Inversion
104    // Adapted from https://paulmillr.com/posts/noble-secp256k1-fast-ecc/#batch-inversion
105    fn batch_inversion(input: Vec<Self>) -> Vec<Self> {
106        let input_length = input.len();
107        if input_length == 0 {
108            return Vec::<Self>::new();
109        }
110
111        let zero = Self::zero();
112        let one = Self::one();
113        let mut scratch: Vec<Self> = vec![zero; input_length];
114        let mut acc = one;
115        scratch[0] = input[0];
116
117        for i in 0..input_length {
118            assert!(!input[i].is_zero(), "Cannot do batch inversion on zero");
119            scratch[i] = acc;
120            acc *= input[i];
121        }
122
123        acc = acc.inverse();
124
125        let mut res = input;
126        for i in (0..input_length).rev() {
127            let tmp = acc * res[i];
128            res[i] = acc * scratch[i];
129            acc = tmp;
130        }
131
132        res
133    }
134
135    /// Parallel version of [`batch_inversion`](Self::batch_inversion).
136    ///
137    /// # Panics
138    ///
139    /// Panics if any of the elements is zero.
140    fn par_batch_inversion(mut input: Vec<Self>) -> Vec<Self> {
141        // Large enough to amortize the one inversion per chunk, small enough
142        // to keep all threads busy.
143        const CHUNK_SIZE: usize = 1 << 12;
144
145        input.par_chunks_mut(CHUNK_SIZE).for_each(|chunk| {
146            let inverses = Self::batch_inversion(chunk.to_vec());
147            chunk.copy_from_slice(&inverses);
148        });
149        input
150    }
151
152    #[inline(always)]
153    fn square(self) -> Self {
154        self * self
155    }
156}
157
158#[cfg(test)]
159#[cfg_attr(coverage_nightly, coverage(off))]
160mod tests {
161    use super::*;
162    use crate::math::other::random_elements;
163    use crate::prelude::*;
164    use crate::tests::test;
165
166    #[macro_rules_attr::apply(test)]
167    fn parallel_and_serial_batch_inversion_agree() {
168        // lengths around the chunk boundaries of `par_batch_inversion`
169        for len in [
170            0,
171            1,
172            7,
173            (1 << 12) - 1,
174            1 << 12,
175            (1 << 12) + 1,
176            (1 << 14) + 3,
177        ] {
178            let bfes: Vec<BFieldElement> = random_elements(len);
179            let bfes: Vec<_> = bfes.into_iter().filter(|x| !x.is_zero()).collect();
180            let serial = BFieldElement::batch_inversion(bfes.clone());
181            let parallel = BFieldElement::par_batch_inversion(bfes.clone());
182            assert_eq!(serial, parallel, "{len}");
183            for (x, x_inv) in bfes.into_iter().zip(parallel) {
184                assert_eq!(BFieldElement::ONE, x * x_inv);
185            }
186
187            let xfes: Vec<XFieldElement> = random_elements(len);
188            let xfes: Vec<_> = xfes.into_iter().filter(|x| !x.is_zero()).collect();
189            let serial_xfes = XFieldElement::batch_inversion(xfes.clone());
190            let parallel_xfes = XFieldElement::par_batch_inversion(xfes);
191            assert_eq!(serial_xfes, parallel_xfes, "{len}");
192        }
193    }
194}
twenty_first/math/traits.rs

twenty_first/math/
traits.rs
