1#![deny(missing_debug_implementations)]
2#![doc = include_str!("../README.md")]
3
4mod fft;
5mod two_stage;
6mod utilities;
7use crate::fft::Fft;
8use crate::utilities::{
9 complex_multiply_accumulate, complex_size, copy_and_pad, next_power_of_2, sum,
10};
11use realfft::num_complex::Complex;
12use realfft::num_traits::Zero;
13use realfft::{FftError, FftNum};
14use rtsan_standalone::nonblocking;
15use thiserror::Error;
16pub use two_stage::TwoStageFFTConvolver;
17pub use utilities::compute_tail_block_size;
18
19#[derive(Error, Debug)]
20pub enum FFTConvolverError {
21 #[error("block size is not allowed to be zero")]
22 BlockSizeZero,
23 #[error("impulse response exceeds configured capacity")]
24 ImpulseResponseExceedsCapacity,
25 #[error("input and output buffers must have the same length")]
26 InputOutputLengthMismatch,
27 #[error("error in fft: {0}")]
28 Fft(#[from] FftError),
29}
30
31#[derive(Clone, Debug)]
47pub struct FFTConvolver<F: FftNum> {
48 ir_len: usize,
49 block_size: usize,
50 seg_size: usize,
51 seg_count: usize,
52 active_seg_count: usize,
53 fft_complex_size: usize,
54 segments: Vec<Vec<Complex<F>>>,
55 segments_ir: Vec<Vec<Complex<F>>>,
56 fft_buffer: Vec<F>,
57 fft: Fft<F>,
58 pre_multiplied: Vec<Complex<F>>,
59 conv: Vec<Complex<F>>,
60 overlap: Vec<F>,
61 current: usize,
62 input_buffer: Vec<F>,
63 input_buffer_fill: usize,
64}
65
66impl<F: FftNum> Default for FFTConvolver<F> {
67 fn default() -> Self {
68 Self {
69 ir_len: Default::default(),
70 block_size: Default::default(),
71 seg_size: Default::default(),
72 seg_count: Default::default(),
73 active_seg_count: Default::default(),
74 fft_complex_size: Default::default(),
75 segments: Default::default(),
76 segments_ir: Default::default(),
77 fft_buffer: Default::default(),
78 fft: Default::default(),
79 pre_multiplied: Default::default(),
80 conv: Default::default(),
81 overlap: Default::default(),
82 current: Default::default(),
83 input_buffer: Default::default(),
84 input_buffer_fill: Default::default(),
85 }
86 }
87}
88
89impl<F: FftNum> FFTConvolver<F> {
90 pub fn init(
119 &mut self,
120 block_size: usize,
121 impulse_response: &[F],
122 ) -> Result<(), FFTConvolverError> {
123 if block_size == 0 {
124 return Err(FFTConvolverError::BlockSizeZero);
125 }
126
127 *self = Self::default();
128
129 self.ir_len = impulse_response.len();
130
131 if self.ir_len == 0 {
132 return Ok(());
133 }
134
135 self.block_size = next_power_of_2(block_size);
136 self.seg_size = 2 * self.block_size;
137 self.seg_count = (self.ir_len as f64 / self.block_size as f64).ceil() as usize;
138 self.active_seg_count = self.seg_count;
139 self.fft_complex_size = complex_size(self.seg_size);
140
141 self.fft.init(self.seg_size);
143 self.fft_buffer = vec![F::zero(); self.seg_size];
144
145 self.segments = vec![vec![Complex::zero(); self.fft_complex_size]; self.seg_count];
147
148 self.segments_ir = vec![vec![Complex::zero(); self.fft_complex_size]; self.seg_count];
150 for (i, segment) in self.segments_ir.iter_mut().enumerate() {
151 let remaining = self.ir_len - (i * self.block_size);
152 let size_copy = if remaining >= self.block_size {
153 self.block_size
154 } else {
155 remaining
156 };
157 copy_and_pad(
158 &mut self.fft_buffer,
159 &impulse_response[i * self.block_size..],
160 size_copy,
161 );
162 self.fft.forward(&mut self.fft_buffer, segment)?;
163 }
164
165 self.pre_multiplied = vec![Complex::zero(); self.fft_complex_size];
167 self.conv = vec![Complex::zero(); self.fft_complex_size];
168 self.overlap.resize(self.block_size, F::zero());
169
170 self.input_buffer = vec![F::zero(); self.block_size];
172 self.input_buffer_fill = 0;
173
174 self.current = 0;
176
177 Ok(())
178 }
179
180 #[nonblocking]
209 pub fn set_response(&mut self, impulse_response: &[F]) -> Result<(), FFTConvolverError> {
210 if impulse_response.len() > self.ir_len {
211 return Err(FFTConvolverError::ImpulseResponseExceedsCapacity);
212 }
213
214 self.fft_buffer.fill(F::zero());
215 self.conv.fill(Complex::zero());
216 self.pre_multiplied.fill(Complex::zero());
217 self.overlap.fill(F::zero());
218
219 self.active_seg_count =
220 (impulse_response.len() as f64 / self.block_size as f64).ceil() as usize;
221
222 for (i, segment) in self
224 .segments_ir
225 .iter_mut()
226 .enumerate()
227 .take(self.active_seg_count)
228 {
229 let remaining = impulse_response.len() - (i * self.block_size);
230 let size_copy = if remaining >= self.block_size {
231 self.block_size
232 } else {
233 remaining
234 };
235 copy_and_pad(
236 &mut self.fft_buffer,
237 &impulse_response[i * self.block_size..],
238 size_copy,
239 );
240 self.fft.forward(&mut self.fft_buffer, segment)?;
241 }
242
243 for segment in self.segments_ir.iter_mut().skip(self.active_seg_count) {
245 segment.fill(Complex::zero());
246 }
247
248 self.input_buffer.fill(F::zero());
249 self.input_buffer_fill = 0;
250 self.current = 0;
251 for segment in &mut self.segments {
252 segment.fill(Complex::zero());
253 }
254
255 Ok(())
256 }
257
258 #[nonblocking]
290 pub fn process(&mut self, input: &[F], output: &mut [F]) -> Result<(), FFTConvolverError> {
291 if input.len() != output.len() {
292 return Err(FFTConvolverError::InputOutputLengthMismatch);
293 }
294
295 if self.active_seg_count == 0 {
296 output.fill(F::zero());
297 return Ok(());
298 }
299
300 let mut processed = 0;
301 while processed < output.len() {
302 let input_buffer_was_empty = self.input_buffer_fill == 0;
303 let processing = std::cmp::min(
304 output.len() - processed,
305 self.block_size - self.input_buffer_fill,
306 );
307
308 let input_buffer_pos = self.input_buffer_fill;
309 self.input_buffer[input_buffer_pos..input_buffer_pos + processing]
310 .copy_from_slice(&input[processed..processed + processing]);
311
312 copy_and_pad(&mut self.fft_buffer, &self.input_buffer, self.block_size);
314 if let Err(err) = self
315 .fft
316 .forward(&mut self.fft_buffer, &mut self.segments[self.current])
317 {
318 output.fill(F::zero());
319 return Err(err.into());
320 }
321
322 if input_buffer_was_empty {
324 self.pre_multiplied.fill(Complex::zero());
325 for i in 1..self.active_seg_count {
326 let index_ir = i;
327 let index_audio = (self.current + i) % self.active_seg_count;
328 complex_multiply_accumulate(
329 &mut self.pre_multiplied,
330 &self.segments_ir[index_ir],
331 &self.segments[index_audio],
332 );
333 }
334 }
335 self.conv.copy_from_slice(&self.pre_multiplied);
336 complex_multiply_accumulate(
337 &mut self.conv,
338 &self.segments[self.current],
339 &self.segments_ir[0],
340 );
341
342 if let Err(err) = self.fft.inverse(&mut self.conv, &mut self.fft_buffer) {
344 output.fill(F::zero());
345 return Err(err.into());
346 }
347
348 sum(
350 &mut output[processed..processed + processing],
351 &self.fft_buffer[input_buffer_pos..input_buffer_pos + processing],
352 &self.overlap[input_buffer_pos..input_buffer_pos + processing],
353 );
354
355 self.input_buffer_fill += processing;
357 if self.input_buffer_fill == self.block_size {
358 self.input_buffer.fill(F::zero());
360 self.input_buffer_fill = 0;
361 self.overlap
363 .copy_from_slice(&self.fft_buffer[self.block_size..self.block_size * 2]);
364
365 self.current = if self.current > 0 {
367 self.current - 1
368 } else {
369 self.active_seg_count - 1
370 };
371 }
372 processed += processing;
373 }
374 Ok(())
375 }
376
377 #[nonblocking]
412 pub fn reset(&mut self) {
413 self.input_buffer.fill(F::zero());
414 self.input_buffer_fill = 0;
415
416 self.fft_buffer.fill(F::zero());
417 for segment in &mut self.segments {
418 segment.fill(Complex::zero());
419 }
420
421 self.conv.fill(Complex::zero());
422 self.pre_multiplied.fill(Complex::zero());
423
424 self.overlap.fill(F::zero());
425 self.current = 0;
426 }
427}
428
429#[cfg(test)]
431mod tests {
432 use crate::{FFTConvolver, FFTConvolverError};
433
434 #[test]
435 fn init_test() {
436 let mut convolver = FFTConvolver::default();
437 let ir = vec![1., 0., 0., 0.];
438 convolver.init(10, &ir).unwrap();
439
440 assert_eq!(convolver.ir_len, 4);
441 assert_eq!(convolver.block_size, 16);
442 assert_eq!(convolver.seg_size, 32);
443 assert_eq!(convolver.seg_count, 1);
444 assert_eq!(convolver.active_seg_count, 1);
445 assert_eq!(convolver.fft_complex_size, 17);
446
447 assert_eq!(convolver.segments.len(), 1);
448 assert_eq!(convolver.segments.first().unwrap().len(), 17);
449 for seg in &convolver.segments {
450 for num in seg {
451 assert_eq!(num.re, 0.);
452 assert_eq!(num.im, 0.);
453 }
454 }
455
456 assert_eq!(convolver.segments_ir.len(), 1);
457 assert_eq!(convolver.segments_ir.first().unwrap().len(), 17);
458 for seg in &convolver.segments_ir {
459 for num in seg {
460 assert_eq!(num.re, 1.);
461 assert_eq!(num.im, 0.);
462 }
463 }
464
465 assert_eq!(convolver.fft_buffer.len(), 32);
466 assert_eq!(*convolver.fft_buffer.first().unwrap(), 1.);
467 for i in 1..convolver.fft_buffer.len() {
468 assert_eq!(convolver.fft_buffer[i], 0.);
469 }
470
471 assert_eq!(convolver.pre_multiplied.len(), 17);
472 for num in &convolver.pre_multiplied {
473 assert_eq!(num.re, 0.);
474 assert_eq!(num.im, 0.);
475 }
476
477 assert_eq!(convolver.conv.len(), 17);
478 for num in &convolver.conv {
479 assert_eq!(num.re, 0.);
480 assert_eq!(num.im, 0.);
481 }
482
483 assert_eq!(convolver.overlap.len(), 16);
484 for num in &convolver.overlap {
485 assert_eq!(*num, 0.);
486 }
487
488 assert_eq!(convolver.input_buffer.len(), 16);
489 for num in &convolver.input_buffer {
490 assert_eq!(*num, 0.);
491 }
492
493 assert_eq!(convolver.input_buffer_fill, 0);
494 }
495
496 #[test]
497 fn process_test() {
498 let mut convolver = FFTConvolver::<f32>::default();
499 let ir = vec![1., 0., 0., 0.];
500 convolver.init(2, &ir).unwrap();
501
502 let input = vec![0., 1., 2., 3.];
503 let mut output = vec![0.; 4];
504 convolver.process(&input, &mut output).unwrap();
505
506 for i in 0..output.len() {
507 assert_eq!(input[i], output[i]);
508 }
509 }
510
511 #[test]
512 fn reset_test() {
513 let ir = vec![0.5, 0.3, 0.2, 0.1];
515 let block_size = 4;
516
517 let mut convolver1 = FFTConvolver::<f32>::default();
519 convolver1.init(block_size, &ir).unwrap();
520
521 let history_input = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
523 let mut history_output = vec![0.0; 8];
524 convolver1
525 .process(&history_input, &mut history_output)
526 .unwrap();
527
528 convolver1.reset();
530
531 let test_input = vec![1.0, 1.0, 1.0, 1.0];
533 let mut output1 = vec![0.0; 4];
534 convolver1.process(&test_input, &mut output1).unwrap();
535
536 let mut convolver2 = FFTConvolver::<f32>::default();
538 convolver2.init(block_size, &ir).unwrap();
539 let mut output2 = vec![0.0; 4];
540 convolver2.process(&test_input, &mut output2).unwrap();
541
542 for i in 0..output1.len() {
544 assert!(
545 (output1[i] - output2[i]).abs() < 1e-5,
546 "Mismatch at index {}: cleared convolver produced {}, fresh convolver produced {}",
547 i,
548 output1[i],
549 output2[i]
550 );
551 }
552 }
553
554 #[test]
555 fn reset_preserves_configuration() {
556 let ir = vec![0.5, 0.3, 0.2, 0.1];
558 let block_size = 4;
559
560 let mut convolver = FFTConvolver::<f32>::default();
561 convolver.init(block_size, &ir).unwrap();
562
563 let ir_len = convolver.ir_len;
564 let block_size_actual = convolver.block_size;
565 let seg_size = convolver.seg_size;
566 let seg_count = convolver.seg_count;
567
568 let input = vec![1.0, 2.0, 3.0, 4.0];
570 let mut output = vec![0.0; 4];
571 convolver.process(&input, &mut output).unwrap();
572
573 convolver.reset();
575
576 assert_eq!(convolver.ir_len, ir_len);
578 assert_eq!(convolver.block_size, block_size_actual);
579 assert_eq!(convolver.seg_size, seg_size);
580 assert_eq!(convolver.seg_count, seg_count);
581 }
582
583 #[test]
584 fn set_response_equals_init() {
585 let ir1 = vec![0.5, 0.3, 0.2, 0.1];
587 let ir2 = vec![0.8, 0.6, 0.4, 0.2];
588 let block_size = 4;
589
590 let mut convolver1 = FFTConvolver::<f32>::default();
592 convolver1.init(block_size, &ir1).unwrap();
593 convolver1.set_response(&ir2).unwrap();
594
595 let mut convolver2 = FFTConvolver::<f32>::default();
597 convolver2.init(block_size, &ir2).unwrap();
598
599 let input = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
601 let mut output1 = vec![0.0; 8];
602 let mut output2 = vec![0.0; 8];
603
604 convolver1.process(&input, &mut output1).unwrap();
605 convolver2.process(&input, &mut output2).unwrap();
606
607 for i in 0..output1.len() {
609 assert!(
610 (output1[i] - output2[i]).abs() < 1e-5,
611 "Mismatch at index {}: set_response produced {}, init produced {}",
612 i,
613 output1[i],
614 output2[i]
615 );
616 }
617 }
618
619 #[test]
620 fn set_response_with_shorter_ir() {
621 let ir1 = vec![0.5, 0.3, 0.2, 0.1, 0.05, 0.02];
623 let ir2 = vec![0.8, 0.6, 0.4];
624 let block_size = 4;
625
626 let mut convolver1 = FFTConvolver::<f32>::default();
628 convolver1.init(block_size, &ir1).unwrap();
629 convolver1.set_response(&ir2).unwrap();
630
631 let mut convolver2 = FFTConvolver::<f32>::default();
633 convolver2.init(block_size, &ir2).unwrap();
634
635 let input = vec![1.0, 1.0, 1.0, 1.0];
637 let mut output1 = vec![0.0; 4];
638 let mut output2 = vec![0.0; 4];
639
640 convolver1.process(&input, &mut output1).unwrap();
641 convolver2.process(&input, &mut output2).unwrap();
642
643 for i in 0..output1.len() {
645 assert!(
646 (output1[i] - output2[i]).abs() < 1e-5,
647 "Mismatch at index {}: set_response produced {}, init produced {}",
648 i,
649 output1[i],
650 output2[i]
651 );
652 }
653 }
654
655 #[test]
656 fn set_response_too_long_returns_error() {
657 let ir1 = vec![0.5, 0.3, 0.2, 0.1];
659 let ir2 = vec![0.8, 0.6, 0.4, 0.2, 0.1, 0.05];
660 let block_size = 4;
661
662 let mut convolver = FFTConvolver::<f32>::default();
663 convolver.init(block_size, &ir1).unwrap();
664
665 let result = convolver.set_response(&ir2);
667 assert!(result.is_err());
668 assert!(matches!(
669 result.unwrap_err(),
670 FFTConvolverError::ImpulseResponseExceedsCapacity
671 ));
672 }
673
674 #[test]
675 fn test_zero_latency() {
676 let mut convolver = FFTConvolver::<f32>::default();
679
680 let ir = vec![0.5, 0.3, 0.2, 0.1];
682 convolver.init(4, &ir).unwrap();
683
684 let mut input = vec![0.0; 16];
686 input[0] = 1.0; let mut output = vec![0.0; 16];
689 convolver.process(&input, &mut output).unwrap();
690
691 assert!(
694 output[0].abs() > 0.0,
695 "Output[0] should be non-zero, indicating zero latency. Got: {}",
696 output[0]
697 );
698
699 assert!(
701 (output[0] - 0.5).abs() < 1e-5,
702 "output[0] should be 0.5, got {}",
703 output[0]
704 );
705 assert!(
706 (output[1] - 0.3).abs() < 1e-5,
707 "output[1] should be 0.3, got {}",
708 output[1]
709 );
710 assert!(
711 (output[2] - 0.2).abs() < 1e-5,
712 "output[2] should be 0.2, got {}",
713 output[2]
714 );
715 assert!(
716 (output[3] - 0.1).abs() < 1e-5,
717 "output[3] should be 0.1, got {}",
718 output[3]
719 );
720 }
721
722 #[test]
723 fn reinit_with_shorter_ir_matches_fresh_init() {
724 let long_ir = vec![0.5_f32; 1000];
725 let short_ir = vec![0.8_f32, 0.6, 0.4, 0.2];
726 let block_size = 4;
727
728 let mut reinit = FFTConvolver::<f32>::default();
730 reinit.init(block_size, &long_ir).unwrap();
731 reinit.init(block_size, &short_ir).unwrap();
732
733 let mut fresh = FFTConvolver::<f32>::default();
735 fresh.init(block_size, &short_ir).unwrap();
736
737 let input = vec![1.0_f32; 64];
738 let mut output_reinit = vec![0.0_f32; 64];
739 let mut output_fresh = vec![0.0_f32; 64];
740
741 reinit.process(&input, &mut output_reinit).unwrap();
742 fresh.process(&input, &mut output_fresh).unwrap();
743
744 for i in 0..output_reinit.len() {
745 assert!(
746 (output_reinit[i] - output_fresh[i]).abs() < 1e-5,
747 "Mismatch at index {}: reinit produced {}, fresh produced {}",
748 i,
749 output_reinit[i],
750 output_fresh[i]
751 );
752 }
753 }
754
755 #[test]
756 fn process_mismatched_lengths_returns_error() {
757 let mut convolver = FFTConvolver::<f32>::default();
758 let ir = vec![1., 0., 0., 0.];
759 convolver.init(4, &ir).unwrap();
760
761 let input = vec![1.0; 4];
762 let mut output = vec![0.0; 8];
763 let result = convolver.process(&input, &mut output);
764 assert!(matches!(
765 result.unwrap_err(),
766 FFTConvolverError::InputOutputLengthMismatch
767 ));
768 }
769
770 #[test]
771 fn reinit_with_empty_ir_produces_silence() {
772 let long_ir = vec![0.5_f32; 1000];
773 let block_size = 4;
774
775 let mut convolver = FFTConvolver::<f32>::default();
776 convolver.init(block_size, &long_ir).unwrap();
777 convolver.init(block_size, &[]).unwrap();
778
779 let input = vec![1.0_f32; 64];
780 let mut output = vec![0.0_f32; 64];
781 convolver.process(&input, &mut output).unwrap();
782
783 for (i, &sample) in output.iter().enumerate() {
784 assert!(
785 sample.abs() < 1e-10,
786 "Expected silence at index {}, got {}",
787 i,
788 sample
789 );
790 }
791 }
792
793 #[test]
794 fn test_block_size_one() {
795 let mut convolver = FFTConvolver::<f32>::default();
796 let ir = vec![0.5_f32, 0.3, 0.2, 0.1];
797 convolver.init(1, &ir).unwrap();
798 assert_eq!(convolver.block_size, 1);
799
800 let mut input = vec![0.0_f32; 16];
801 input[0] = 1.0;
802 let mut output = vec![0.0_f32; 16];
803 convolver.process(&input, &mut output).unwrap();
804
805 assert!((output[0] - 0.5).abs() < 1e-5);
806 assert!((output[1] - 0.3).abs() < 1e-5);
807 assert!((output[2] - 0.2).abs() < 1e-5);
808 assert!((output[3] - 0.1).abs() < 1e-5);
809 }
810
811 #[test]
812 fn test_large_ir() {
813 let ir_len = 8192_usize;
815 let block_size = 512;
816
817 let mut ir = vec![0.0_f32; ir_len];
818 ir[0] = 1.0; let mut convolver = FFTConvolver::<f32>::default();
821 convolver.init(block_size, &ir).unwrap();
822
823 let input: Vec<f32> = (0..ir_len).map(|i| i as f32 * 0.001).collect();
824 let mut output = vec![0.0_f32; ir_len];
825 convolver.process(&input, &mut output).unwrap();
826
827 for i in 0..ir_len {
828 assert!(
829 (output[i] - input[i]).abs() < 1e-4,
830 "Mismatch at {}: input={}, output={}",
831 i,
832 input[i],
833 output[i]
834 );
835 }
836 }
837}