1use p3_dft::{NaiveDft, TwoAdicSubgroupDft};
2use p3_field::{ExtensionField, TwoAdicField};
3use p3_matrix::Matrix;
4use p3_matrix::dense::RowMajorMatrix;
5use rand::SeedableRng;
6use rand::distr::{Distribution, StandardUniform};
7use rand::rngs::SmallRng;
8
9pub fn test_dft_matches_naive<F, Dft>()
10where
11 F: TwoAdicField,
12 StandardUniform: Distribution<F>,
13 Dft: TwoAdicSubgroupDft<F>,
14{
15 let dft = Dft::default();
16 let mut rng = SmallRng::seed_from_u64(1);
17 for log_h in 0..5 {
18 let h = 1 << log_h;
19 let mat = RowMajorMatrix::<F>::rand(&mut rng, h, 3);
20 let dft_naive = NaiveDft.dft_batch(mat.clone());
21 let dft_result = dft.dft_batch(mat);
22 assert_eq!(dft_naive, dft_result.to_row_major_matrix());
23 }
24}
25
26pub fn test_coset_dft_matches_naive<F, Dft>()
27where
28 F: TwoAdicField,
29 StandardUniform: Distribution<F>,
30 Dft: TwoAdicSubgroupDft<F>,
31{
32 let dft = Dft::default();
33 let mut rng = SmallRng::seed_from_u64(1);
34 for log_h in 0..5 {
35 let h = 1 << log_h;
36 let mat = RowMajorMatrix::<F>::rand(&mut rng, h, 3);
37 let shift = F::GENERATOR;
38 let coset_dft_naive = NaiveDft.coset_dft_batch(mat.clone(), shift);
39 let coset_dft_result = dft.coset_dft_batch(mat, shift);
40 assert_eq!(coset_dft_naive, coset_dft_result.to_row_major_matrix());
41 }
42}
43
44pub fn test_idft_matches_naive<F, Dft>()
45where
46 F: TwoAdicField,
47 StandardUniform: Distribution<F>,
48 Dft: TwoAdicSubgroupDft<F>,
49{
50 let dft = Dft::default();
51 let mut rng = SmallRng::seed_from_u64(1);
52 for log_h in 0..5 {
53 let h = 1 << log_h;
54 let mat = RowMajorMatrix::<F>::rand(&mut rng, h, 3);
55 let idft_naive = NaiveDft.idft_batch(mat.clone());
56 let idft_result = dft.idft_batch(mat.clone());
57 assert_eq!(idft_naive, idft_result.to_row_major_matrix());
58 }
59}
60
61pub fn test_coset_idft_matches_naive<F, Dft>()
62where
63 F: TwoAdicField,
64 StandardUniform: Distribution<F>,
65 Dft: TwoAdicSubgroupDft<F>,
66{
67 let dft = Dft::default();
68 let mut rng = SmallRng::seed_from_u64(1);
69 for log_h in 0..5 {
70 let h = 1 << log_h;
71 let mat = RowMajorMatrix::<F>::rand(&mut rng, h, 3);
72 let shift = F::GENERATOR;
73 let idft_naive = NaiveDft.coset_idft_batch(mat.clone(), shift);
74 let idft_result = dft.coset_idft_batch(mat, shift);
75 assert_eq!(idft_naive, idft_result.to_row_major_matrix());
76 }
77}
78
79pub fn test_lde_matches_naive<F, Dft>()
80where
81 F: TwoAdicField,
82 StandardUniform: Distribution<F>,
83 Dft: TwoAdicSubgroupDft<F>,
84{
85 let dft = Dft::default();
86 let mut rng = SmallRng::seed_from_u64(1);
87 for log_h in 0..5 {
88 let h = 1 << log_h;
89 let mat = RowMajorMatrix::<F>::rand(&mut rng, h, 3);
90 let lde_naive = NaiveDft.lde_batch(mat.clone(), 1);
91 let lde_result = dft.lde_batch(mat, 1);
92 assert_eq!(lde_naive, lde_result.to_row_major_matrix());
93 }
94}
95
96pub fn test_coset_lde_matches_naive<F, Dft>()
97where
98 F: TwoAdicField,
99 StandardUniform: Distribution<F>,
100 Dft: TwoAdicSubgroupDft<F>,
101{
102 let dft = Dft::default();
103 let mut rng = SmallRng::seed_from_u64(1);
104 for log_h in 0..5 {
105 let h = 1 << log_h;
106 let mat = RowMajorMatrix::<F>::rand(&mut rng, h, 3);
107 let shift = F::GENERATOR;
108 let coset_lde_naive = NaiveDft.coset_lde_batch(mat.clone(), 1, shift);
109 let coset_lde_result = dft.coset_lde_batch(mat, 1, shift);
110 assert_eq!(coset_lde_naive, coset_lde_result.to_row_major_matrix());
111 }
112}
113
114pub fn test_dft_idft_consistency<F, Dft>()
115where
116 F: TwoAdicField,
117 StandardUniform: Distribution<F>,
118 Dft: TwoAdicSubgroupDft<F>,
119{
120 let dft = Dft::default();
121 let mut rng = SmallRng::seed_from_u64(1);
122 for log_h in 0..5 {
123 let h = 1 << log_h;
124 let original = RowMajorMatrix::<F>::rand(&mut rng, h, 3);
125 let dft_output = dft.dft_batch(original.clone());
126 let idft_output = dft.idft_batch(dft_output.to_row_major_matrix());
127 assert_eq!(original, idft_output.to_row_major_matrix());
128 }
129}
130
131pub fn test_dft_algebra_matches_naive<F, EF, Dft>()
132where
133 F: TwoAdicField,
134 EF: ExtensionField<F> + TwoAdicField,
135 StandardUniform: Distribution<EF>,
136 Dft: TwoAdicSubgroupDft<F>,
137{
138 let dft = Dft::default();
139 let mut rng = SmallRng::seed_from_u64(1);
140 for log_h in 0..5 {
141 let h = 1 << log_h;
142 let mat = RowMajorMatrix::<EF>::rand(&mut rng, h, 3);
143 let dft_naive = NaiveDft.dft_batch(mat.clone());
144 let dft_result = dft.dft_algebra_batch(mat);
145 assert_eq!(dft_naive, dft_result.to_row_major_matrix());
146 }
147}
148
149pub fn test_coset_dft_algebra_matches_naive<F, EF, Dft>()
150where
151 F: TwoAdicField,
152 EF: ExtensionField<F> + TwoAdicField,
153 StandardUniform: Distribution<EF>,
154 Dft: TwoAdicSubgroupDft<F>,
155{
156 let dft = Dft::default();
157 let mut rng = SmallRng::seed_from_u64(1);
158 for log_h in 0..5 {
159 let h = 1 << log_h;
160 let mat = RowMajorMatrix::<EF>::rand(&mut rng, h, 3);
161 let shift = F::GENERATOR;
162 let coset_dft_naive = NaiveDft.coset_dft_batch(mat.clone(), EF::from(shift));
163 let coset_dft_result = dft.coset_dft_algebra_batch(mat, shift);
164 assert_eq!(coset_dft_naive, coset_dft_result.to_row_major_matrix());
165 }
166}
167
168pub fn test_idft_algebra_matches_naive<F, EF, Dft>()
169where
170 F: TwoAdicField,
171 EF: ExtensionField<F> + TwoAdicField,
172 StandardUniform: Distribution<EF>,
173 Dft: TwoAdicSubgroupDft<F>,
174{
175 let dft = Dft::default();
176 let mut rng = SmallRng::seed_from_u64(1);
177 for log_h in 0..5 {
178 let h = 1 << log_h;
179 let mat = RowMajorMatrix::<EF>::rand(&mut rng, h, 3);
180 let idft_naive = NaiveDft.idft_batch(mat.clone());
181 let idft_result = dft.idft_algebra_batch(mat.clone());
182 assert_eq!(idft_naive, idft_result.to_row_major_matrix());
183 }
184}
185
186pub fn test_coset_idft_algebra_matches_naive<F, EF, Dft>()
187where
188 F: TwoAdicField,
189 EF: ExtensionField<F> + TwoAdicField,
190 StandardUniform: Distribution<EF>,
191 Dft: TwoAdicSubgroupDft<F>,
192{
193 let dft = Dft::default();
194 let mut rng = SmallRng::seed_from_u64(1);
195 for log_h in 0..5 {
196 let h = 1 << log_h;
197 let mat = RowMajorMatrix::<EF>::rand(&mut rng, h, 3);
198 let shift = F::GENERATOR;
199 let idft_naive = NaiveDft.coset_idft_batch(mat.clone(), EF::from(shift));
200 let idft_result = dft.coset_idft_algebra_batch(mat, shift);
201 assert_eq!(idft_naive, idft_result.to_row_major_matrix());
202 }
203}
204
205pub fn test_lde_algebra_matches_naive<F, EF, Dft>()
206where
207 F: TwoAdicField,
208 EF: ExtensionField<F> + TwoAdicField,
209 StandardUniform: Distribution<EF>,
210 Dft: TwoAdicSubgroupDft<F>,
211{
212 let dft = Dft::default();
213 let mut rng = SmallRng::seed_from_u64(1);
214 for log_h in 0..5 {
215 let h = 1 << log_h;
216 let mat = RowMajorMatrix::<EF>::rand(&mut rng, h, 3);
217 let lde_naive = NaiveDft.lde_batch(mat.clone(), 1);
218 let lde_result = dft.lde_algebra_batch(mat, 1);
219 assert_eq!(lde_naive, lde_result.to_row_major_matrix());
220 }
221}
222
223pub fn test_coset_lde_algebra_matches_naive<F, EF, Dft>()
224where
225 F: TwoAdicField,
226 EF: ExtensionField<F> + TwoAdicField,
227 StandardUniform: Distribution<EF>,
228 Dft: TwoAdicSubgroupDft<F>,
229{
230 let dft = Dft::default();
231 let mut rng = SmallRng::seed_from_u64(1);
232 for log_h in 0..5 {
233 let h = 1 << log_h;
234 let mat = RowMajorMatrix::<EF>::rand(&mut rng, h, 3);
235 let shift = F::GENERATOR;
236 let coset_lde_naive = NaiveDft.coset_lde_batch(mat.clone(), 1, EF::from(shift));
237 let coset_lde_result = dft.coset_lde_algebra_batch(mat, 1, shift);
238 assert_eq!(coset_lde_naive, coset_lde_result.to_row_major_matrix());
239 }
240}
241
242pub fn test_dft_idft_algebra_consistency<F, EF, Dft>()
243where
244 F: TwoAdicField,
245 EF: ExtensionField<F> + TwoAdicField,
246 StandardUniform: Distribution<EF>,
247 Dft: TwoAdicSubgroupDft<F>,
248{
249 let dft = Dft::default();
250 let mut rng = SmallRng::seed_from_u64(1);
251 for log_h in 0..5 {
252 let h = 1 << log_h;
253 let original = RowMajorMatrix::<EF>::rand(&mut rng, h, 3);
254 let dft_output = dft.dft_algebra_batch(original.clone());
255 let idft_output = dft.idft_algebra_batch(dft_output.to_row_major_matrix());
256 assert_eq!(original, idft_output.to_row_major_matrix());
257 }
258}
259
260pub fn test_dft_idft_algebra_consistency_large<F, EF, Dft>()
261where
262 F: TwoAdicField,
263 EF: ExtensionField<F> + TwoAdicField,
264 StandardUniform: Distribution<EF>,
265 Dft: TwoAdicSubgroupDft<F>,
266{
267 let dft = Dft::default();
268 let mut rng = SmallRng::seed_from_u64(1);
269 for log_h in &[14, 15, 16, 17] {
270 let h = 1 << log_h;
271 let original = RowMajorMatrix::<EF>::rand(&mut rng, h, 3);
272 let dft_output = dft.dft_algebra_batch(original.clone());
273 let idft_output = dft.idft_algebra_batch(dft_output.to_row_major_matrix());
274 assert_eq!(
275 original,
276 idft_output.to_row_major_matrix(),
277 "Error Found in size: {}",
278 h
279 );
280 }
281}
282
283pub fn test_large_coset_ldes_agree<F, Dft1, Dft2>()
284where
285 F: TwoAdicField,
286 StandardUniform: Distribution<F>,
287 Dft1: TwoAdicSubgroupDft<F>,
288 Dft2: TwoAdicSubgroupDft<F>,
289{
290 let dft1 = Dft1::default();
291 let dft2 = Dft2::default();
292 let mut rng = SmallRng::seed_from_u64(1);
293 for log_h in &[14, 15, 16, 17] {
294 let h = 1 << log_h;
295 let mat = RowMajorMatrix::<F>::rand(&mut rng, h, 16);
296 let shift = F::GENERATOR;
297 let coset_lde_result1 = dft1.coset_lde_batch(mat.clone(), 1, shift);
298 let coset_lde_result2 = dft2.coset_lde_batch(mat, 1, shift);
299
300 assert_eq!(
301 coset_lde_result1.to_row_major_matrix(),
302 coset_lde_result2.to_row_major_matrix()
303 );
304 }
305}
306
307#[macro_export]
308macro_rules! test_field_dft {
309 ($mod:ident, $field:ty, $extfield:ty, $dft:ty) => {
310 mod $mod {
311 #[test]
312 fn dft_matches_naive() {
313 $crate::test_dft_matches_naive::<$field, $dft>();
314 }
315
316 #[test]
317 fn coset_dft_matches_naive() {
318 $crate::test_coset_dft_matches_naive::<$field, $dft>();
319 }
320
321 #[test]
322 fn idft_matches_naive() {
323 $crate::test_idft_matches_naive::<$field, $dft>();
324 }
325
326 #[test]
327 fn coset_idft_matches_naive() {
328 $crate::test_coset_idft_matches_naive::<$field, $dft>();
329 }
330
331 #[test]
332 fn lde_matches_naive() {
333 $crate::test_lde_matches_naive::<$field, $dft>();
334 }
335
336 #[test]
337 fn coset_lde_matches_naive() {
338 $crate::test_coset_lde_matches_naive::<$field, $dft>();
339 }
340
341 #[test]
342 fn dft_idft_consistency() {
343 $crate::test_dft_idft_consistency::<$field, $dft>();
344 }
345
346 #[test]
347 fn dft_algebra_matches_naive() {
348 $crate::test_dft_algebra_matches_naive::<$field, $extfield, $dft>();
349 }
350
351 #[test]
352 fn coset_dft_algebra_matches_naive() {
353 $crate::test_coset_dft_algebra_matches_naive::<$field, $extfield, $dft>();
354 }
355
356 #[test]
357 fn idft_algebra_matches_naive() {
358 $crate::test_idft_algebra_matches_naive::<$field, $extfield, $dft>();
359 }
360
361 #[test]
362 fn coset_idft_algebra_matches_naive() {
363 $crate::test_coset_idft_algebra_matches_naive::<$field, $extfield, $dft>();
364 }
365
366 #[test]
367 fn lde_algebra_matches_naive() {
368 $crate::test_lde_algebra_matches_naive::<$field, $extfield, $dft>();
369 }
370
371 #[test]
372 fn coset_lde_algebra_matches_naive() {
373 $crate::test_coset_lde_algebra_matches_naive::<$field, $extfield, $dft>();
374 }
375
376 #[test]
377 fn dft_idft_algebra_consistency() {
378 $crate::test_dft_idft_algebra_consistency::<$field, $extfield, $dft>();
379 }
380 }
381 };
382}
383
384#[macro_export]
385macro_rules! test_field_dft_large {
386 ($mod:ident, $field:ty, $extfield:ty, $dft1:ty, $dft2:ty) => {
387 mod $mod {
388 #[test]
389 fn test_dft1_idft_algebra_consistency_large() {
390 $crate::test_dft_idft_algebra_consistency_large::<$field, $extfield, $dft1>();
391 }
392
393 #[test]
394 fn test_dft2_idft_algebra_consistency_large() {
395 $crate::test_dft_idft_algebra_consistency_large::<$field, $extfield, $dft2>();
396 }
397
398 #[test]
399 fn test_large_coset_ldes_agree() {
400 $crate::test_large_coset_ldes_agree::<$field, $dft1, $dft2>();
401 }
402 }
403 };
404}