fp/field/
smallfq.rs

1use std::sync::LazyLock;
2
3use dashmap::DashMap as HashMap;
4
5use super::{
6    Field, Fp,
7    element::{FieldElement, FieldElementContainer},
8    field_internal::FieldInternal,
9};
10use crate::{
11    PRIME_TO_INDEX_MAP,
12    constants::BITS_PER_LIMB,
13    limb::Limb,
14    prime::{Prime, ValidPrime, log2},
15    vector::inner::FqVector,
16};
17
18static SMALL_CONWAY_POLYS: [[[u32; 17]; 15]; 54] = include!("small_conway_polys.txt");
19
20type ZechTable = &'static [SmallFqElement];
21type Polynomial<P> = FqVector<Fp<P>>;
22
23/// A table of lazily initialized [Zech logarithms][zech_logs].
24///
25/// Key is the field, value is a fully initialized table of Zech logarithms.
26///
27/// [zech_logs]: https://en.wikipedia.org/wiki/Zech%27s_logarithm
28static ZECH_LOGS: LazyLock<HashMap<(ValidPrime, u32), ZechTable>> = LazyLock::new(HashMap::new);
29
30fn mul_by_a<P: Prime>(conway_poly: &Polynomial<P>, poly: Polynomial<P>) -> Polynomial<P> {
31    let prime_field = conway_poly.fq();
32
33    // Shift all entries up by one. We're assuming that `poly` is a polynomial representing an
34    // element of the field, so the top coefficient is zero, and there is no overflow.
35    // TODO: If we ever implement a way to shift the coefficients of a vector in-place, this would
36    // be a good place to use it.
37    let mut next = Polynomial::from_slice(
38        prime_field,
39        &std::iter::once(prime_field.zero())
40            .chain(poly.iter())
41            .take(poly.len())
42            .collect::<Vec<_>>(),
43    );
44    let leading_coeff = next.entry(next.len() - 1);
45    next.add(conway_poly, -leading_coeff);
46    next
47}
48
49fn make_zech_log_table<P: Prime>(p: P, d: u32) -> ZechTable {
50    let prime_field = Fp::new(p);
51    let conway_poly = {
52        let v = SMALL_CONWAY_POLYS[PRIME_TO_INDEX_MAP[p.as_usize()]][d as usize - 2]
53            .iter()
54            .take(d as usize + 1)
55            .map(|c| prime_field.el(*c))
56            .collect::<Vec<_>>();
57        Polynomial::from_slice(prime_field, &v)
58    };
59    let mul_by_a = |current| mul_by_a(&conway_poly, current);
60
61    // Generate a lookup table. For every element represented as a polynomial, we store the
62    // power of `a` that corresponds to it.
63    let poly_to_power: HashMap<Polynomial<P>, u32> = HashMap::new();
64    let mut current = Polynomial::new(prime_field, conway_poly.len());
65    current.set_entry(0, prime_field.one());
66    poly_to_power.insert(current.clone(), 0);
67
68    let q = p.pow(d);
69    for i in 1..q - 1 {
70        current = mul_by_a(current);
71        poly_to_power.insert(current.clone(), i);
72    }
73
74    // Loop over all elements again, but now recording logarithms.
75    let mut table = vec![SmallFqElement(None); q as usize - 1];
76    let mut current = Polynomial::new(prime_field, conway_poly.len());
77    current.set_entry(0, prime_field.one());
78    for i in 0..q - 1 {
79        let mut current_plus_1 = current.clone();
80        current_plus_1.add_basis_element(0, prime_field.one());
81        table[i as usize] = SmallFqElement(poly_to_power.get(&current_plus_1).as_deref().copied());
82
83        current = mul_by_a(current);
84    }
85    // From the documentation for `Box::leak`: "This function is mainly useful for data that lives
86    // for the remainder of the program's life". Zech tables are initialized once and then never
87    // mutated, so even storing an immutable reference is fine.
88    Box::leak(table.into_boxed_slice())
89}
90
91/// Return the Zech logarithm table for the given field. If it does not exist yet, initialize it.
92/// The initialization might be fairly expensive (several ms).
93fn zech_logs<P: Prime>(p: P, d: u32) -> ZechTable {
94    let table = ZECH_LOGS
95        .entry((p.to_dyn(), d))
96        .or_insert_with(|| make_zech_log_table(p, d));
97    *table
98}
99
100/// A field of order `q = p^d`, where `q < 2^16` and `d > 1`. Fields of that size are small enough
101/// that we can cache their Zech logarithms.
102///
103/// Note: This populates the Zech logarithm table eagerly, which can be rather expensive (several
104/// milliseconds). Only construct these fields if you're going to use them.
105#[derive(Copy, Clone)]
106pub struct SmallFq<P> {
107    p: P,
108    d: u32,
109    q: u32,
110    table: ZechTable,
111}
112
113impl<P: Prime> SmallFq<P> {
114    pub fn new(p: P, d: u32) -> Self {
115        assert!(d > 1, "Use Fp for prime fields");
116        assert!(log2(p.pow(d) as usize) < 16, "Field too large");
117
118        Self {
119            p,
120            d,
121            q: p.pow(d),
122            table: zech_logs(p, d),
123        }
124    }
125
126    /// Return the element `-1`. If `p = 2`, this is `a^0 = 1`. Otherwise, it is `a^((q - 1) / 2)`.
127    pub fn negative_one(self) -> FieldElement<Self> {
128        let e = if self.p == 2 { 0 } else { (self.q() - 1) / 2 };
129        self.el(SmallFqElement(Some(e)))
130    }
131
132    /// The distinguished primitive element that generates the multiplicative group of the field.
133    pub fn a(self) -> FieldElement<Self> {
134        self.el(SmallFqElement(Some(1)))
135    }
136}
137
138// Custom debug implementation to avoid printing `q` and the Zech table
139impl<P: Prime> std::fmt::Debug for SmallFq<P> {
140    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
141        f.debug_struct("SmallFq")
142            .field("p", &self.p)
143            .field("d", &self.d)
144            .finish()
145    }
146}
147
148impl<P: Prime> PartialEq for SmallFq<P> {
149    fn eq(&self, other: &Self) -> bool {
150        self.p == other.p && self.d == other.d
151    }
152}
153
154impl<P: Prime> Eq for SmallFq<P> {}
155
156impl<P: Prime> std::hash::Hash for SmallFq<P> {
157    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
158        self.p.hash(state);
159        self.d.hash(state);
160    }
161}
162
163impl<P: Prime> std::fmt::Display for SmallFq<P> {
164    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
165        write!(f, "F_({}^{})", self.p, self.d)
166    }
167}
168
169impl<P: Prime> Field for SmallFq<P> {
170    type Characteristic = P;
171
172    fn characteristic(self) -> Self::Characteristic {
173        self.p
174    }
175
176    fn degree(self) -> u32 {
177        self.d
178    }
179
180    // We cache the value of `q` because it's the bottleneck for `FieldInternal::el`
181    fn q(self) -> u32 {
182        self.q
183    }
184
185    fn zero(self) -> FieldElement<Self> {
186        self.el(SmallFqElement(None))
187    }
188
189    fn one(self) -> FieldElement<Self> {
190        self.el(SmallFqElement(Some(0)))
191    }
192
193    #[cfg(feature = "proptest")]
194    fn arb_element(self) -> impl proptest::strategy::Strategy<Value = FieldElement<Self>> {
195        use proptest::prelude::*;
196
197        (0..self.q()).prop_map(move |i| {
198            // Note that, mathematically, the elements of SmallFq<P> are 0 and a^i for i in [0,q-1).
199            // Since a^(q-1) = 1 = a^0, we can use i == 0 to represent the zero element. Thanks to
200            // `FieldInternal::el`, the exponent of q - 1 will be reduced to 0.
201            self.el(SmallFqElement(if i == 0 { None } else { Some(i) }))
202        })
203    }
204}
205
206impl<P: Prime> FieldInternal for SmallFq<P> {
207    type ElementContainer = SmallFqElement;
208
209    fn el(self, value: Self::ElementContainer) -> FieldElement<Self> {
210        let reduced_value = value.0.map(|e| e % (self.q() - 1));
211        FieldElement::new(self, SmallFqElement(reduced_value))
212    }
213
214    fn add_assign(self, a: &mut FieldElement<Self>, b: FieldElement<Self>) {
215        // It's simpler to just define `add` directly, using the inherent symmetry. It doesn't
216        // really matter since `SmallFqElement` is `Copy`.
217        *a = self.add(*a, b);
218    }
219
220    fn add(self, a: FieldElement<Self>, b: FieldElement<Self>) -> FieldElement<Self> {
221        self.el(match (a.value, b.value) {
222            (SmallFqElement(None), b) => b,
223            (a, SmallFqElement(None)) => a,
224            (SmallFqElement(Some(a)), SmallFqElement(Some(b))) => {
225                // a^m + a^n = a^m (1 + a^(n - m)) = a^(m + Zech(n - m))
226                let (a, b) = if a >= b { (a, b) } else { (b, a) };
227                let zech = self.table[(a - b) as usize];
228                if let Some(zech) = zech.0 {
229                    SmallFqElement(Some(b + zech))
230                } else {
231                    SmallFqElement(None)
232                }
233            }
234        })
235    }
236
237    fn mul_assign(self, a: &mut FieldElement<Self>, b: FieldElement<Self>) {
238        if let (SmallFqElement(Some(a)), SmallFqElement(Some(b))) = (&mut a.value, b.value) {
239            *a += b;
240            *a %= self.q() - 1;
241        } else {
242            a.value = SmallFqElement(None);
243        }
244    }
245
246    fn neg(self, a: FieldElement<Self>) -> FieldElement<Self> {
247        self.mul(a, self.negative_one())
248    }
249
250    fn inv(self, a: FieldElement<Self>) -> Option<FieldElement<Self>> {
251        let complement = match a.0? {
252            0 => 0,
253            x => self.q() - 1 - x,
254        };
255        Some(self.el(SmallFqElement(Some(complement))))
256    }
257
258    fn frobenius(self, a: FieldElement<Self>) -> FieldElement<Self> {
259        self.el(SmallFqElement(
260            a.0.map(|x| x * self.characteristic().as_u32()),
261        ))
262    }
263
264    /// This is 2n + 1 if `element` is a^n, and 0 otherwise.
265    fn encode(self, element: FieldElement<Self>) -> Limb {
266        element.value.0.map(|x| ((x as Limb) << 1) | 1).unwrap_or(0)
267    }
268
269    fn decode(self, element: Limb) -> FieldElement<Self> {
270        self.el(if element & 1 == 0 {
271            // This only checks that the element is even, but by the definition of `encode`, this
272            // only happens if the element is zero.
273            SmallFqElement(None)
274        } else {
275            SmallFqElement(Some((element >> 1) as u32))
276        })
277    }
278
279    fn bit_length(self) -> usize {
280        // A field has q - 1 units, so SmallFqElement is either Some(a) where a is in [0, q - 2], or
281        // None. We add 1 bit to account for encoding the None case.
282        BITS_PER_LIMB - (self.q() - 1).leading_zeros() as usize + 1
283    }
284
285    fn fma_limb(self, limb_a: Limb, limb_b: Limb, coeff: FieldElement<Self>) -> Limb {
286        let bit_length = self.bit_length();
287        let mut result: Limb = 0;
288        let mut shift = 0;
289        for (a, b) in self.unpack(limb_a).zip(self.unpack(limb_b)) {
290            result += self.encode(a + coeff * b) << shift;
291            shift += bit_length;
292        }
293        result
294    }
295
296    fn reduce(self, limb: Limb) -> Limb {
297        limb
298    }
299}
300
301#[cfg(feature = "proptest")]
302impl<P: Prime> proptest::arbitrary::Arbitrary for SmallFq<P> {
303    type Parameters = ();
304    type Strategy = proptest::strategy::BoxedStrategy<Self>;
305
306    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
307        use proptest::prelude::*;
308
309        // Can't be a closure because the return value references this function, and a closure
310        // would not live long enough
311        fn largest_degree(p: impl Prime) -> u32 {
312            let mut d = 2;
313            while p.pow(d) < 1 << 16 {
314                d += 1;
315            }
316            d - 1
317        }
318
319        any_with::<P>(std::num::NonZeroU32::new(256)) // prime is at most 256
320            .prop_flat_map(|p| (2..=largest_degree(p)).prop_map(move |d| Self::new(p, d)))
321            .boxed()
322    }
323}
324
325impl<P: Prime> crate::MaybeArbitrary<()> for SmallFq<P> {}
326
327/// A field element, stored as the exponent of a distinguished generator of the group of units.
328/// `None` if the element is zero.
329#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
330pub struct SmallFqElement(pub(super) Option<u32>);
331
332impl FieldElementContainer for SmallFqElement {}
333
334impl std::fmt::Display for SmallFqElement {
335    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
336        match self.0 {
337            None => write!(f, "0"),
338            Some(0) => write!(f, "1"),
339            Some(1) => write!(f, "a"),
340            Some(x) => write!(f, "a^{x}"),
341        }
342    }
343}
344
345#[cfg(test)]
346pub(crate) mod tests {
347    use super::SmallFq;
348    use crate::field::tests::field_tests;
349
350    mod validprime {
351        use super::*;
352        use crate::prime::ValidPrime;
353
354        field_tests!(SmallFq<ValidPrime>);
355    }
356
357    macro_rules! static_smallfq_tests {
358        ($p:tt) => {
359            paste::paste! {
360                static_smallfq_tests!(@ [<$p:lower>], $p, $p);
361            }
362        };
363        (@ $mod_name:ident, $p_ident:ident, $p_ty:ty) => {
364            mod $mod_name {
365                use super::*;
366                use crate::prime::$p_ident;
367
368                field_tests!(SmallFq<$p_ty>);
369            }
370        };
371    }
372
373    static_smallfq_tests!(P2);
374    cfg_if::cfg_if! { if #[cfg(feature = "odd-primes")] {
375        static_smallfq_tests!(P3);
376        static_smallfq_tests!(P5);
377        static_smallfq_tests!(P7);
378    }}
379}
fp/field/smallfq.rs

fp/field/
smallfq.rs
