Skip to main content

ark_ff/fields/
fft_friendly.rs

1use num_bigint::BigUint;
2/// The interface for fields that are able to be used in FFTs.
3pub trait FftField: crate::Field {
4    /// The generator of the multiplicative group of the field
5    const GENERATOR: Self;
6
7    /// Let `N` be the size of the multiplicative group defined by the field.
8    /// Then `TWO_ADICITY` is the two-adicity of `N`, i.e. the integer `s`
9    /// such that `N = 2^s * t` for some odd integer `t`.
10    const TWO_ADICITY: u32;
11
12    /// 2^s root of unity computed by GENERATOR^t
13    const TWO_ADIC_ROOT_OF_UNITY: Self;
14
15    /// An integer `b` such that there exists a multiplicative subgroup
16    /// of size `b^k` for some integer `k`.
17    const SMALL_SUBGROUP_BASE: Option<u32> = None;
18
19    /// The integer `k` such that there exists a multiplicative subgroup
20    /// of size `Self::SMALL_SUBGROUP_BASE^k`.
21    const SMALL_SUBGROUP_BASE_ADICITY: Option<u32> = None;
22
23    /// GENERATOR^((MODULUS-1) / (2^s *
24    /// SMALL_SUBGROUP_BASE^SMALL_SUBGROUP_BASE_ADICITY)) Used for mixed-radix
25    /// FFT.
26    const LARGE_SUBGROUP_ROOT_OF_UNITY: Option<Self> = None;
27
28    /// Returns the root of unity of order n, if one exists.
29    /// If no small multiplicative subgroup is defined, this is the 2-adic root
30    /// of unity of order n (for n a power of 2).
31    /// If a small multiplicative subgroup is defined, this is the root of unity
32    /// of order n for the larger subgroup generated by
33    /// `FftConfig::LARGE_SUBGROUP_ROOT_OF_UNITY`
34    /// (for n = 2^i * FftConfig::SMALL_SUBGROUP_BASE^j for some i, j).
35    fn get_root_of_unity(n: u64) -> Option<Self> {
36        let mut omega: Self;
37        if let Some(large_subgroup_root_of_unity) = Self::LARGE_SUBGROUP_ROOT_OF_UNITY {
38            let q = Self::SMALL_SUBGROUP_BASE.expect(
39                "LARGE_SUBGROUP_ROOT_OF_UNITY should only be set in conjunction with SMALL_SUBGROUP_BASE",
40            ) as u64;
41            let small_subgroup_base_adicity = Self::SMALL_SUBGROUP_BASE_ADICITY.expect(
42                "LARGE_SUBGROUP_ROOT_OF_UNITY should only be set in conjunction with SMALL_SUBGROUP_BASE_ADICITY",
43            );
44
45            let q_adicity = crate::utils::k_adicity(q, n);
46            let q_part = q.checked_pow(q_adicity)?;
47
48            let two_adicity = crate::utils::k_adicity(2, n);
49            let two_part = 2u64.checked_pow(two_adicity)?;
50
51            if n != two_part * q_part
52                || (two_adicity > Self::TWO_ADICITY)
53                || (q_adicity > small_subgroup_base_adicity)
54            {
55                return None;
56            }
57
58            omega = large_subgroup_root_of_unity;
59            for _ in q_adicity..small_subgroup_base_adicity {
60                omega = omega.pow([q]);
61            }
62
63            for _ in two_adicity..Self::TWO_ADICITY {
64                omega.square_in_place();
65            }
66        } else {
67            // Compute the next power of 2.
68            let size = n.next_power_of_two();
69            let log_size_of_group = ark_std::log2(usize::try_from(size).expect("too large"));
70
71            if n != size || log_size_of_group > Self::TWO_ADICITY {
72                return None;
73            }
74
75            // Compute the generator for the multiplicative subgroup.
76            // It should be 2^(log_size_of_group) root of unity.
77            omega = Self::TWO_ADIC_ROOT_OF_UNITY;
78            for _ in log_size_of_group..Self::TWO_ADICITY {
79                omega.square_in_place();
80            }
81        }
82        Some(omega)
83    }
84
85    /// Returns the root of unity of order n, if one exists.
86    /// If no small multiplicative subgroup is defined, this is the 2-adic root
87    /// of unity of order n (for n a power of 2).
88    /// If a small multiplicative subgroup is defined, this is the root of unity
89    /// of order n for the larger subgroup generated by
90    /// `FftConfig::LARGE_SUBGROUP_ROOT_OF_UNITY`
91    /// (for n = 2^i * FftConfig::SMALL_SUBGROUP_BASE^j for some i, j).
92    fn get_root_of_unity_big_int(n: BigUint) -> Option<Self> {
93        let mut omega: Self;
94        if let Some(large_subgroup_root_of_unity) = Self::LARGE_SUBGROUP_ROOT_OF_UNITY {
95            let q = Self::SMALL_SUBGROUP_BASE.expect(
96                "LARGE_SUBGROUP_ROOT_OF_UNITY should only be set in conjunction with SMALL_SUBGROUP_BASE",
97            ) as u64;
98            let small_subgroup_base_adicity = Self::SMALL_SUBGROUP_BASE_ADICITY.expect(
99                "LARGE_SUBGROUP_ROOT_OF_UNITY should only be set in conjunction with SMALL_SUBGROUP_BASE_ADICITY",
100            );
101
102            let q_adicity = crate::utils::k_adicity_big_int(q.into(), n.clone());
103            let q_part = q.checked_pow(q_adicity)?;
104
105            let two_adicity = crate::utils::k_adicity_big_int(2_u8.into(), n.clone());
106            let two_part = 2u64.checked_pow(two_adicity)?;
107
108            if n != (two_part * q_part).into()
109                || (two_adicity > Self::TWO_ADICITY)
110                || (q_adicity > small_subgroup_base_adicity)
111            {
112                return None;
113            }
114
115            omega = large_subgroup_root_of_unity;
116            for _ in q_adicity..small_subgroup_base_adicity {
117                omega = omega.pow([q]);
118            }
119
120            for _ in two_adicity..Self::TWO_ADICITY {
121                omega.square_in_place();
122            }
123        } else {
124            // Compute the next power of 2.
125            let (size, log_size_of_group) = if n == BigUint::from(1_u8) {
126                (BigUint::from(1_u8), 0)
127            } else {
128                let log_size_of_group = (n).bits() - 1;
129                (
130                    BigUint::from(1_u8) << (log_size_of_group),
131                    log_size_of_group,
132                )
133            };
134
135            if n != size || log_size_of_group > Self::TWO_ADICITY.into() {
136                return None;
137            }
138
139            // Compute the generator for the multiplicative subgroup.
140            // It should be 2^(log_size_of_group) root of unity.
141            omega = Self::TWO_ADIC_ROOT_OF_UNITY;
142            for _ in log_size_of_group..Self::TWO_ADICITY.into() {
143                omega.square_in_place();
144            }
145        }
146        Some(omega)
147    }
148}