lambdaworks_math/fft/cpu/
roots_of_unity.rs

1use crate::field::{
2    element::FieldElement,
3    traits::{IsFFTField, RootsConfig},
4};
5use alloc::vec::Vec;
6
7use crate::fft::errors::FFTError;
8
9use super::bit_reversing::in_place_bit_reverse_permute;
10
11/// Returns a `Vec` of the powers of a `2^n`th primitive root of unity in some configuration
12/// `config`. For example, in a `Natural` config this would yield: w^0, w^1, w^2...
13pub fn get_powers_of_primitive_root<F: IsFFTField>(
14    n: u64,
15    count: usize,
16    config: RootsConfig,
17) -> Result<Vec<FieldElement<F>>, FFTError> {
18    if count == 0 {
19        return Ok(Vec::new());
20    }
21
22    let root = match config {
23        RootsConfig::Natural | RootsConfig::BitReverse => F::get_primitive_root_of_unity(n)?,
24        _ => F::get_primitive_root_of_unity(n)?.inv().unwrap(),
25    };
26    let up_to = match config {
27        RootsConfig::Natural | RootsConfig::NaturalInversed => count,
28        // In bit reverse form we could need as many as `(1 << count.bits()) - 1` roots
29        _ => count.next_power_of_two(),
30    };
31
32    let mut results = Vec::with_capacity(up_to);
33    // NOTE: a nice version would be using `core::iter::successors`. However, this is 10% faster.
34    results.extend((0..up_to).scan(FieldElement::one(), |state, _| {
35        let res = state.clone();
36        *state = &(*state) * &root;
37        Some(res)
38    }));
39
40    if matches!(
41        config,
42        RootsConfig::BitReverse | RootsConfig::BitReverseInversed
43    ) {
44        in_place_bit_reverse_permute(&mut results);
45    }
46
47    Ok(results)
48}
49
50/// Returns a `Vec` of the powers of a `2^n`th primitive root of unity, scaled `offset` times,
51/// in a Natural configuration.
52pub fn get_powers_of_primitive_root_coset<F: IsFFTField>(
53    n: u64,
54    count: usize,
55    offset: &FieldElement<F>,
56) -> Result<Vec<FieldElement<F>>, FFTError> {
57    let root = F::get_primitive_root_of_unity(n)?;
58    let results = (0..count).map(|i| offset * root.pow(i));
59
60    Ok(results.collect())
61}
62
63/// Returns 2^`order` / 2 twiddle factors for FFT in some configuration `config`.
64/// Twiddle factors are powers of a primitive root of unity of the field, used for FFT
65/// computations. FFT only requires the first half of all the powers
66pub fn get_twiddles<F: IsFFTField>(
67    order: u64,
68    config: RootsConfig,
69) -> Result<Vec<FieldElement<F>>, FFTError> {
70    if order > 63 {
71        return Err(FFTError::OrderError(order));
72    }
73
74    get_powers_of_primitive_root(order, (1 << order) / 2, config)
75}
76
77#[cfg(test)]
78mod tests {
79    use crate::{
80        fft::{
81            cpu::{bit_reversing::in_place_bit_reverse_permute, roots_of_unity::get_twiddles},
82            errors::FFTError,
83        },
84        field::{test_fields::u64_test_field::U64TestField, traits::RootsConfig},
85    };
86    use proptest::prelude::*;
87
88    type F = U64TestField;
89
90    proptest! {
91        #[test]
92        fn test_gen_twiddles_bit_reversed_validity(n in 1..8_u64) {
93            let twiddles = get_twiddles::<F>(n, RootsConfig::Natural).unwrap();
94            let mut twiddles_to_reorder = get_twiddles(n, RootsConfig::BitReverse).unwrap();
95            in_place_bit_reverse_permute(&mut twiddles_to_reorder); // so now should be naturally ordered
96
97            prop_assert_eq!(twiddles, twiddles_to_reorder);
98        }
99    }
100
101    #[test]
102    fn gen_twiddles_with_order_greater_than_63_should_fail() {
103        let twiddles = get_twiddles::<F>(64, RootsConfig::Natural);
104
105        assert!(matches!(twiddles, Err(FFTError::OrderError(_))));
106    }
107}
lambdaworks_math/fft/cpu/roots_of_unity.rs

lambdaworks_math/fft/cpu/
roots_of_unity.rs
