p3_uni_stark/
zerofier_coset.rs

1use alloc::vec::Vec;
2
3use itertools::Itertools;
4use p3_field::{
5    batch_multiplicative_inverse, cyclic_subgroup_coset_known_order, Field, PackedField,
6    TwoAdicField,
7};
8
9/// Precomputations of the evaluation of `Z_H(X) = X^n - 1` on a coset `s K` with `H <= K`.
10#[derive(Debug)]
11pub struct ZerofierOnCoset<F: Field> {
12    /// `n = |H|`.
13    log_n: usize,
14    /// `rate = |K|/|H|`.
15    rate_bits: usize,
16    coset_shift: F,
17    /// Holds `g^n * (w^n)^i - 1 = g^n * v^i - 1` for `i in 0..rate`, with `w` a generator of `K` and `v` a
18    /// `rate`-primitive root of unity.
19    evals: Vec<F>,
20    /// Holds the multiplicative inverses of `evals`.
21    inverses: Vec<F>,
22}
23
24impl<F: TwoAdicField> ZerofierOnCoset<F> {
25    pub fn new(log_n: usize, rate_bits: usize, coset_shift: F) -> Self {
26        let s_pow_n = coset_shift.exp_power_of_2(log_n);
27        let evals = F::two_adic_generator(rate_bits)
28            .powers()
29            .take(1 << rate_bits)
30            .map(|x| s_pow_n * x - F::ONE)
31            .collect::<Vec<_>>();
32        let inverses = batch_multiplicative_inverse(&evals);
33        Self {
34            log_n,
35            rate_bits,
36            coset_shift,
37            evals,
38            inverses,
39        }
40    }
41
42    /// Returns `Z_H(g * w^i)`.
43    pub fn eval(&self, i: usize) -> F {
44        self.evals[i & ((1 << self.rate_bits) - 1)]
45    }
46
47    /// Returns `1 / Z_H(g * w^i)`.
48    pub fn eval_inverse(&self, i: usize) -> F {
49        self.inverses[i & ((1 << self.rate_bits) - 1)]
50    }
51
52    /// Like `eval_inverse`, but for a range of indices starting with `i_start`.
53    pub fn eval_inverse_packed<P: PackedField<Scalar = F>>(&self, i_start: usize) -> P {
54        let mut packed = P::ZERO;
55        packed
56            .as_slice_mut()
57            .iter_mut()
58            .enumerate()
59            .for_each(|(j, packed_j)| *packed_j = self.eval_inverse(i_start + j));
60        packed
61    }
62
63    /// Evaluate the Langrange basis polynomial, `L_i(x) = Z_H(x) / (x - g_H^i)`, on our coset `s K`.
64    /// Here `L_i(x)` is unnormalized in the sense that it evaluates to some nonzero value at `g_H^i`,
65    /// not necessarily 1.
66    pub fn lagrange_basis_unnormalized(&self, i: usize) -> Vec<F> {
67        let log_coset_size = self.log_n + self.rate_bits;
68        let coset_size = 1 << log_coset_size;
69        let g_h = F::two_adic_generator(self.log_n);
70        let g_k = F::two_adic_generator(log_coset_size);
71
72        let target_point = g_h.exp_u64(i as u64);
73        let denominators = cyclic_subgroup_coset_known_order(g_k, self.coset_shift, coset_size)
74            .map(|x| x - target_point)
75            .collect_vec();
76        let inverses = batch_multiplicative_inverse(&denominators);
77
78        self.evals
79            .iter()
80            .cycle()
81            .zip(inverses)
82            .map(|(&z_h, inv)| z_h * inv)
83            .collect()
84    }
85}