lambdaworks_math/circle/
cfft.rs

1extern crate alloc;
2use crate::field::{element::FieldElement, fields::mersenne31::field::Mersenne31Field};
3use alloc::vec::Vec;
4
5#[cfg(feature = "alloc")]
6/// fft in place algorithm used to evaluate a polynomial of degree 2^n - 1 in 2^n points.
7/// Input must be of size 2^n for some n.
8pub fn cfft(
9    input: &mut [FieldElement<Mersenne31Field>],
10    twiddles: Vec<Vec<FieldElement<Mersenne31Field>>>,
11) {
12    // If the input size is 2^n, then log_2_size is n.
13    let log_2_size = input.len().trailing_zeros();
14
15    // The cfft has n layers.
16    (0..log_2_size).for_each(|i| {
17        // In each layer i we split the current input in chunks of size 2^{i+1}.
18        let chunk_size = 1 << (i + 1);
19        let half_chunk_size = 1 << i;
20        input.chunks_mut(chunk_size).for_each(|chunk| {
21            // We split each chunk in half, calling the first half hi_part and the second hal low_part.
22            let (hi_part, low_part) = chunk.split_at_mut(half_chunk_size);
23
24            // We apply the corresponding butterfly for every element j of the high and low part.
25            hi_part
26                .iter_mut()
27                .zip(low_part)
28                .enumerate()
29                .for_each(|(j, (hi, low))| {
30                    let temp = *low * twiddles[i as usize][j];
31                    *low = *hi - temp;
32                    *hi += temp
33                });
34        });
35    });
36}
37
38#[cfg(feature = "alloc")]
39/// The inverse fft algorithm used to interpolate 2^n points.
40/// Input must be of size 2^n for some n.
41pub fn icfft(
42    input: &mut [FieldElement<Mersenne31Field>],
43    twiddles: Vec<Vec<FieldElement<Mersenne31Field>>>,
44) {
45    // If the input size is 2^n, then log_2_size is n.
46    let log_2_size = input.len().trailing_zeros();
47
48    // The icfft has n layers.
49    (0..log_2_size).for_each(|i| {
50        // In each layer i we split the current input in chunks of size 2^{n - i}.
51        let chunk_size = 1 << (log_2_size - i);
52        let half_chunk_size = chunk_size >> 1;
53        input.chunks_mut(chunk_size).for_each(|chunk| {
54            // We split each chunk in half, calling the first half hi_part and the second hal low_part.
55            let (hi_part, low_part) = chunk.split_at_mut(half_chunk_size);
56
57            // We apply the corresponding butterfly for every element j of the high and low part.
58            hi_part
59                .iter_mut()
60                .zip(low_part)
61                .enumerate()
62                .for_each(|(j, (hi, low))| {
63                    let temp = *hi + *low;
64                    *low = (*hi - *low) * twiddles[i as usize][j];
65                    *hi = temp;
66                });
67        });
68    });
69}
70
71/// This function permutes a slice of field elements to order the result of the cfft in the natural way.
72/// We call the natural order to [P(x0, y0), P(x1, y1), P(x2, y2), ...],
73/// where (x0, y0) is the first point of the corresponding coset.
74/// The cfft doesn't return the evaluations in the natural order.
75/// For example, if we apply the cfft to 8 coefficients of a polynomial of degree 7 we'll get the evaluations in this order:
76/// [P(x0, y0), P(x2, y2), P(x4, y4), P(x6, y6), P(x7, y7), P(x5, y5), P(x3, y3), P(x1, y1)],
77/// where the even indices are found first in ascending order and then the odd indices in descending order.
78/// This function permutes the slice [0, 2, 4, 6, 7, 5, 3, 1] into [0, 1, 2, 3, 4, 5, 6, 7].
79/// TODO: This can be optimized by performing in-place value swapping (WIP).  
80pub fn order_cfft_result_naive(
81    input: &[FieldElement<Mersenne31Field>],
82) -> Vec<FieldElement<Mersenne31Field>> {
83    let mut result = Vec::new();
84    let length = input.len();
85    for i in 0..length / 2 {
86        result.push(input[i]); // We push the left index.
87        result.push(input[length - i - 1]); // We push the right index.
88    }
89    result
90}
91
92/// This function permutes a slice of field elements to order the input of the icfft in a specific way.
93/// For example, if we want to interpolate 8 points we should input them in the icfft in this order:
94/// [(x0, y0), (x2, y2), (x4, y4), (x6, y6), (x7, y7), (x5, y5), (x3, y3), (x1, y1)],
95/// where the even indices are found first in ascending order and then the odd indices in descending order.
96/// This function permutes the slice [0, 1, 2, 3, 4, 5, 6, 7] into [0, 2, 4, 6, 7, 5, 3, 1].
97/// TODO: This can be optimized by performing in-place value swapping (WIP).  
98pub fn order_icfft_input_naive(
99    input: &mut [FieldElement<Mersenne31Field>],
100) -> Vec<FieldElement<Mersenne31Field>> {
101    let mut result = Vec::new();
102
103    // We push the even indices.
104    (0..input.len()).step_by(2).for_each(|i| {
105        result.push(input[i]);
106    });
107
108    // We push the odd indices.
109    (1..input.len()).step_by(2).rev().for_each(|i| {
110        result.push(input[i]);
111    });
112    result
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    type FE = FieldElement<Mersenne31Field>;
119
120    #[test]
121    fn ordering_cfft_result_works_for_4_points() {
122        let expected_slice = [FE::from(0), FE::from(1), FE::from(2), FE::from(3)];
123
124        let slice = [FE::from(0), FE::from(2), FE::from(3), FE::from(1)];
125
126        let res = order_cfft_result_naive(&slice);
127
128        assert_eq!(res, expected_slice)
129    }
130
131    #[test]
132    fn ordering_cfft_result_works_for_16_points() {
133        let expected_slice = [
134            FE::from(0),
135            FE::from(1),
136            FE::from(2),
137            FE::from(3),
138            FE::from(4),
139            FE::from(5),
140            FE::from(6),
141            FE::from(7),
142            FE::from(8),
143            FE::from(9),
144            FE::from(10),
145            FE::from(11),
146            FE::from(12),
147            FE::from(13),
148            FE::from(14),
149            FE::from(15),
150        ];
151
152        let slice = [
153            FE::from(0),
154            FE::from(2),
155            FE::from(4),
156            FE::from(6),
157            FE::from(8),
158            FE::from(10),
159            FE::from(12),
160            FE::from(14),
161            FE::from(15),
162            FE::from(13),
163            FE::from(11),
164            FE::from(9),
165            FE::from(7),
166            FE::from(5),
167            FE::from(3),
168            FE::from(1),
169        ];
170
171        let res = order_cfft_result_naive(&slice);
172
173        assert_eq!(res, expected_slice)
174    }
175
176    #[test]
177    fn from_natural_to_icfft_input_order_works() {
178        let mut slice = [
179            FE::from(0),
180            FE::from(1),
181            FE::from(2),
182            FE::from(3),
183            FE::from(4),
184            FE::from(5),
185            FE::from(6),
186            FE::from(7),
187            FE::from(8),
188            FE::from(9),
189            FE::from(10),
190            FE::from(11),
191            FE::from(12),
192            FE::from(13),
193            FE::from(14),
194            FE::from(15),
195        ];
196
197        let expected_slice = [
198            FE::from(0),
199            FE::from(2),
200            FE::from(4),
201            FE::from(6),
202            FE::from(8),
203            FE::from(10),
204            FE::from(12),
205            FE::from(14),
206            FE::from(15),
207            FE::from(13),
208            FE::from(11),
209            FE::from(9),
210            FE::from(7),
211            FE::from(5),
212            FE::from(3),
213            FE::from(1),
214        ];
215
216        let res = order_icfft_input_naive(&mut slice);
217
218        assert_eq!(res, expected_slice)
219    }
220}
lambdaworks_math/circle/cfft.rs

lambdaworks_math/circle/
cfft.rs
