fft_convolver/
two_stage.rs

1use crate::utilities::{compute_tail_block_size, next_power_of_2};
2use crate::{FFTConvolver, FFTConvolverError};
3use realfft::FftNum;
4use rtsan_standalone::nonblocking;
5
6/// TwoStageFFTConvolver
7/// Implementation of a partitioned FFT convolution algorithm with non-uniform block sizes.
8///
9/// Uses two different block sizes: a small head block size for low latency and a
10/// larger tail block size for efficient processing of long impulse responses.
11/// Internally uses three [`FFTConvolver`] instances:
12///
13/// - **Head**: Processes the first portion of the IR at the small block size (zero latency)
14/// - **Tail0**: Transition stage that bridges the latency gap between head and tail
15/// - **Tail**: Processes the remaining IR at the larger block size (efficient)
16///
17/// Some notes on how to use it:
18/// - After initialization with an impulse response, subsequent data portions of
19///   arbitrary length can be convolved. The convolver internally can handle
20///   this by using appropriate buffering.
21/// - The convolver works without "latency" (except for the required
22///   processing time, of course), i.e. the output always is the convolved
23///   input for each processing call.
24/// - The convolver is suitable for real-time processing which means that no
25///   "unpredictable" operations like allocations, locking, API calls, etc. are
26///   performed during processing (all necessary allocations and preparations take
27///   place during initialization).
28#[derive(Clone, Debug)]
29pub struct TwoStageFFTConvolver<F: FftNum> {
30    ir_len: usize,
31    head_block_size: usize,
32    tail_block_size: usize,
33
34    head_convolver: FFTConvolver<F>,
35    tail_convolver0: FFTConvolver<F>,
36    tail_convolver: FFTConvolver<F>,
37
38    tail_input: Vec<F>,
39    tail_input_fill: usize,
40
41    tail_output0: Vec<F>,
42    tail_precalculated0: Vec<F>,
43
44    tail_output: Vec<F>,
45    tail_precalculated: Vec<F>,
46
47    // Pre-allocated buffer to avoid allocations during process().
48    // Used to copy tail_input before passing to tail convolvers,
49    // since both are fields of self and cannot be borrowed simultaneously.
50    processing_buffer: Vec<F>,
51
52    precalculated_pos: usize,
53}
54
55impl<F: FftNum> Default for TwoStageFFTConvolver<F> {
56    fn default() -> Self {
57        Self {
58            ir_len: Default::default(),
59            head_block_size: Default::default(),
60            tail_block_size: Default::default(),
61            head_convolver: Default::default(),
62            tail_convolver0: Default::default(),
63            tail_convolver: Default::default(),
64            tail_input: Default::default(),
65            tail_input_fill: Default::default(),
66            tail_output0: Default::default(),
67            tail_precalculated0: Default::default(),
68            tail_output: Default::default(),
69            tail_precalculated: Default::default(),
70            processing_buffer: Default::default(),
71            precalculated_pos: Default::default(),
72        }
73    }
74}
75
76impl<F: FftNum> TwoStageFFTConvolver<F> {
77    /// Initializes the two-stage convolver with an impulse response
78    ///
79    /// This method sets up all internal buffers and prepares the convolver for processing.
80    /// The head block size determines latency and will be rounded up to the next power of 2.
81    /// The tail block size determines efficiency for the tail portion and will also be
82    /// rounded up to the next power of 2.
83    ///
84    /// Use [`compute_tail_block_size`] to calculate an
85    /// optimal tail block size based on García's formula.
86    ///
87    /// All memory allocations happen during initialization, making subsequent processing
88    /// operations real-time safe.
89    ///
90    /// # Arguments
91    ///
92    /// * `head_block_size` - Block size for the head convolver (determines latency).
93    ///   Will be rounded up to the next power of 2. Must be > 0.
94    /// * `tail_block_size` - Block size for the tail convolver (determines efficiency).
95    ///   Will be rounded up to the next power of 2. Must be > 0.
96    ///   If smaller than `head_block_size`, the two values are automatically swapped.
97    /// * `impulse_response` - The impulse response to convolve with. Can be empty.
98    ///
99    /// # Returns
100    ///
101    /// Returns `BlockSizeZero` if either block size is 0.
102    ///
103    /// # Example
104    ///
105    /// ```
106    /// use fft_convolver::{TwoStageFFTConvolver, compute_tail_block_size};
107    ///
108    /// let ir = vec![0.5_f32; 10000];
109    /// let head_block_size = 64;
110    /// let tail_block_size = compute_tail_block_size(head_block_size, ir.len());
111    ///
112    /// let mut convolver = TwoStageFFTConvolver::default();
113    /// convolver.init(head_block_size, tail_block_size, &ir).unwrap();
114    /// ```
115    pub fn init(
116        &mut self,
117        head_block_size: usize,
118        tail_block_size: usize,
119        impulse_response: &[F],
120    ) -> Result<(), FFTConvolverError> {
121        if head_block_size == 0 || tail_block_size == 0 {
122            return Err(FFTConvolverError::BlockSizeZero);
123        }
124
125        *self = Self::default();
126
127        self.head_block_size = next_power_of_2(head_block_size);
128        self.tail_block_size = next_power_of_2(tail_block_size);
129
130        if self.head_block_size > self.tail_block_size {
131            std::mem::swap(&mut self.head_block_size, &mut self.tail_block_size);
132        }
133
134        self.ir_len = impulse_response.len();
135        let ir_len = self.ir_len;
136
137        if ir_len == 0 {
138            return Ok(());
139        }
140
141        let head_ir_len = ir_len.min(self.tail_block_size);
142        self.head_convolver
143            .init(self.head_block_size, &impulse_response[..head_ir_len])?;
144
145        if ir_len > self.tail_block_size {
146            let conv1_ir_len = (ir_len - self.tail_block_size).min(self.tail_block_size);
147            self.tail_convolver0.init(
148                self.head_block_size,
149                &impulse_response[self.tail_block_size..self.tail_block_size + conv1_ir_len],
150            )?;
151            self.tail_output0 = vec![F::zero(); self.tail_block_size];
152            self.tail_precalculated0 = vec![F::zero(); self.tail_block_size];
153        }
154
155        if ir_len > 2 * self.tail_block_size {
156            let tail_ir_len = ir_len - 2 * self.tail_block_size;
157            self.tail_convolver.init(
158                self.tail_block_size,
159                &impulse_response[2 * self.tail_block_size..2 * self.tail_block_size + tail_ir_len],
160            )?;
161            self.tail_output = vec![F::zero(); self.tail_block_size];
162            self.tail_precalculated = vec![F::zero(); self.tail_block_size];
163        }
164
165        if !self.tail_precalculated0.is_empty() || !self.tail_precalculated.is_empty() {
166            self.tail_input = vec![F::zero(); self.tail_block_size];
167            self.processing_buffer = vec![F::zero(); self.tail_block_size];
168        }
169
170        self.tail_input_fill = 0;
171        self.precalculated_pos = 0;
172
173        Ok(())
174    }
175
176    /// Initializes the two-stage convolver with an automatically computed tail block size
177    ///
178    /// This is a convenience method that computes the optimal tail block size using
179    /// García's formula and then calls [`init`](Self::init).
180    ///
181    /// # Arguments
182    ///
183    /// * `head_block_size` - Block size for the head convolver (determines latency).
184    ///   Will be rounded up to the next power of 2. Must be > 0.
185    /// * `impulse_response` - The impulse response to convolve with. Can be empty.
186    ///
187    /// # Returns
188    ///
189    /// Returns `BlockSizeZero` if head_block_size is 0.
190    ///
191    /// # Example
192    ///
193    /// ```
194    /// use fft_convolver::TwoStageFFTConvolver;
195    ///
196    /// let ir = vec![0.5_f32; 10000];
197    /// let mut convolver = TwoStageFFTConvolver::default();
198    /// convolver.init_default(64, &ir).unwrap();
199    /// ```
200    pub fn init_default(
201        &mut self,
202        head_block_size: usize,
203        impulse_response: &[F],
204    ) -> Result<(), FFTConvolverError> {
205        let tail_block_size = compute_tail_block_size(head_block_size, impulse_response.len());
206        self.init(head_block_size, tail_block_size, impulse_response)
207    }
208
209    /// Updates the impulse response without reallocating buffers
210    ///
211    /// This method allows changing the impulse response at runtime while maintaining
212    /// real-time safety by avoiding allocations. The new impulse response must not
213    /// exceed the length of the original impulse response used during initialization.
214    ///
215    /// # Arguments
216    ///
217    /// * `impulse_response` - The new impulse response (must be ≤ original length)
218    ///
219    /// # Returns
220    ///
221    /// Returns `ImpulseResponseExceedsCapacity` if the new impulse response is longer
222    /// than the original one.
223    ///
224    /// # Example
225    ///
226    /// ```
227    /// use fft_convolver::{TwoStageFFTConvolver, compute_tail_block_size};
228    ///
229    /// let ir1 = vec![0.5_f32; 10000];
230    /// let head_block_size = 64;
231    /// let tail_block_size = compute_tail_block_size(head_block_size, ir1.len());
232    ///
233    /// let mut convolver = TwoStageFFTConvolver::default();
234    /// convolver.init(head_block_size, tail_block_size, &ir1).unwrap();
235    ///
236    /// let ir2 = vec![0.8_f32; 5000];
237    /// convolver.set_response(&ir2).unwrap();
238    /// ```
239    #[nonblocking]
240    pub fn set_response(&mut self, impulse_response: &[F]) -> Result<(), FFTConvolverError> {
241        if impulse_response.len() > self.ir_len {
242            return Err(FFTConvolverError::ImpulseResponseExceedsCapacity);
243        }
244
245        let ir_len = impulse_response.len();
246
247        // Head
248        let head_ir_len = ir_len.min(self.tail_block_size);
249        self.head_convolver
250            .set_response(&impulse_response[..head_ir_len])?;
251
252        // Tail0
253        if !self.tail_precalculated0.is_empty() {
254            if ir_len > self.tail_block_size {
255                let conv1_ir_len = (ir_len - self.tail_block_size).min(self.tail_block_size);
256                self.tail_convolver0.set_response(
257                    &impulse_response[self.tail_block_size..self.tail_block_size + conv1_ir_len],
258                )?;
259            } else {
260                self.tail_convolver0.set_response(&[])?;
261            }
262            self.tail_output0.fill(F::zero());
263            self.tail_precalculated0.fill(F::zero());
264        }
265
266        // Tail
267        if !self.tail_precalculated.is_empty() {
268            if ir_len > 2 * self.tail_block_size {
269                let tail_ir_len = ir_len - 2 * self.tail_block_size;
270                self.tail_convolver.set_response(
271                    &impulse_response
272                        [2 * self.tail_block_size..2 * self.tail_block_size + tail_ir_len],
273                )?;
274            } else {
275                self.tail_convolver.set_response(&[])?;
276            }
277            self.tail_output.fill(F::zero());
278            self.tail_precalculated.fill(F::zero());
279        }
280
281        self.tail_input.fill(F::zero());
282        self.tail_input_fill = 0;
283        self.precalculated_pos = 0;
284
285        Ok(())
286    }
287
288    /// Convolves the input samples with the impulse response and outputs the result
289    ///
290    /// This is a real-time safe operation that performs no allocations. Internal buffering
291    /// handles arbitrary sizes and ensures the output is always properly aligned with the
292    /// input (zero latency except for processing time).
293    ///
294    /// # Arguments
295    ///
296    /// * `input` - The input samples to convolve
297    /// * `output` - Buffer to write the convolution result. Must have the same length as `input`.
298    ///
299    /// # Returns
300    ///
301    /// Returns `InputOutputLengthMismatch` if `input` and `output` have different lengths.
302    /// Returns `Fft` error if an FFT operation fails.
303    ///
304    /// # Example
305    ///
306    /// ```
307    /// use fft_convolver::{TwoStageFFTConvolver, compute_tail_block_size};
308    ///
309    /// let ir = vec![0.5_f32; 10000];
310    /// let head_block_size = 64;
311    /// let tail_block_size = compute_tail_block_size(head_block_size, ir.len());
312    ///
313    /// let mut convolver = TwoStageFFTConvolver::default();
314    /// convolver.init(head_block_size, tail_block_size, &ir).unwrap();
315    ///
316    /// let input = vec![1.0_f32; 256];
317    /// let mut output = vec![0.0_f32; 256];
318    /// convolver.process(&input, &mut output).unwrap();
319    /// ```
320    #[nonblocking]
321    pub fn process(&mut self, input: &[F], output: &mut [F]) -> Result<(), FFTConvolverError> {
322        if input.len() != output.len() {
323            return Err(FFTConvolverError::InputOutputLengthMismatch);
324        }
325        self.head_convolver.process(input, output)?;
326
327        if self.tail_input.is_empty() {
328            return Ok(());
329        }
330
331        let len = input.len();
332        let mut processed = 0;
333
334        while processed < len {
335            let remaining = len - processed;
336            let processing =
337                remaining.min(self.head_block_size - (self.tail_input_fill % self.head_block_size));
338
339            let sum_begin = processed;
340            let sum_end = processed + processing;
341
342            // Add precalculated tail0 output
343            if !self.tail_precalculated0.is_empty() {
344                let precalc = &self.tail_precalculated0;
345                let mut pos = self.precalculated_pos;
346                #[allow(clippy::explicit_counter_loop)]
347                for sample in &mut output[sum_begin..sum_end] {
348                    *sample = *sample + precalc[pos];
349                    pos += 1;
350                }
351            }
352
353            // Add precalculated tail output
354            if !self.tail_precalculated.is_empty() {
355                let precalc = &self.tail_precalculated;
356                let mut pos = self.precalculated_pos;
357                #[allow(clippy::explicit_counter_loop)]
358                for sample in &mut output[sum_begin..sum_end] {
359                    *sample = *sample + precalc[pos];
360                    pos += 1;
361                }
362            }
363
364            self.precalculated_pos += processing;
365
366            // Buffer input for tail processing
367            self.tail_input[self.tail_input_fill..self.tail_input_fill + processing]
368                .copy_from_slice(&input[processed..processed + processing]);
369            self.tail_input_fill += processing;
370
371            // Process tail0 incrementally (every head_block_size samples)
372            if !self.tail_precalculated0.is_empty()
373                && self.tail_input_fill.is_multiple_of(self.head_block_size)
374            {
375                let block_offset = self.tail_input_fill - self.head_block_size;
376                self.processing_buffer[..self.head_block_size].copy_from_slice(
377                    &self.tail_input[block_offset..block_offset + self.head_block_size],
378                );
379                self.tail_convolver0.process(
380                    &self.processing_buffer[..self.head_block_size],
381                    &mut self.tail_output0[block_offset..block_offset + self.head_block_size],
382                )?;
383                if self.tail_input_fill == self.tail_block_size {
384                    std::mem::swap(&mut self.tail_precalculated0, &mut self.tail_output0);
385                }
386            }
387
388            // Process tail (when full tail block accumulated)
389            if !self.tail_precalculated.is_empty() && self.tail_input_fill == self.tail_block_size {
390                std::mem::swap(&mut self.tail_precalculated, &mut self.tail_output);
391                self.processing_buffer.copy_from_slice(&self.tail_input);
392                self.tail_convolver
393                    .process(&self.processing_buffer, &mut self.tail_output)?;
394            }
395
396            if self.tail_input_fill == self.tail_block_size {
397                self.tail_input_fill = 0;
398                self.precalculated_pos = 0;
399            }
400
401            processed += processing;
402        }
403
404        Ok(())
405    }
406
407    /// Clears the internal processing state while preserving the impulse response
408    ///
409    /// This real-time safe operation resets all internal buffers that store the
410    /// convolution state, effectively removing any "history" or "tail" from previous
411    /// processing. The impulse response configuration remains intact, so processing
412    /// can continue immediately.
413    ///
414    /// This is useful when handling stream discontinuities such as:
415    /// - Seeking in audio playback
416    /// - Pause/resume operations with large time gaps
417    /// - Switching between different audio sources
418    ///
419    /// After calling `reset()`, the next `process()` call will produce output as if
420    /// the convolver had just been initialized.
421    ///
422    /// # Example
423    ///
424    /// ```
425    /// use fft_convolver::{TwoStageFFTConvolver, compute_tail_block_size};
426    ///
427    /// let ir = vec![0.5_f32; 10000];
428    /// let head_block_size = 64;
429    /// let tail_block_size = compute_tail_block_size(head_block_size, ir.len());
430    ///
431    /// let mut convolver = TwoStageFFTConvolver::default();
432    /// convolver.init(head_block_size, tail_block_size, &ir).unwrap();
433    ///
434    /// let input = vec![1.0_f32; 256];
435    /// let mut output = vec![0.0_f32; 256];
436    /// convolver.process(&input, &mut output).unwrap();
437    ///
438    /// convolver.reset();
439    /// convolver.process(&input, &mut output).unwrap();
440    /// ```
441    #[nonblocking]
442    pub fn reset(&mut self) {
443        self.head_convolver.reset();
444        self.tail_convolver0.reset();
445        self.tail_convolver.reset();
446
447        self.tail_input.fill(F::zero());
448        self.tail_input_fill = 0;
449
450        self.tail_output0.fill(F::zero());
451        self.tail_precalculated0.fill(F::zero());
452
453        self.tail_output.fill(F::zero());
454        self.tail_precalculated.fill(F::zero());
455
456        self.processing_buffer.fill(F::zero());
457        self.precalculated_pos = 0;
458    }
459}
460
461#[cfg(test)]
462mod tests {
463    use crate::{FFTConvolver, FFTConvolverError, TwoStageFFTConvolver, compute_tail_block_size};
464
465    #[test]
466    fn init_test() {
467        let mut convolver = TwoStageFFTConvolver::<f32>::default();
468        let ir = vec![0.5; 10000];
469        convolver.init(64, 4096, &ir).unwrap();
470
471        assert_eq!(convolver.head_block_size, 64);
472        assert_eq!(convolver.tail_block_size, 4096);
473        assert_eq!(convolver.ir_len, 10000);
474
475        // Head covers ir[0..4096], tail0 covers ir[4096..8192], tail covers ir[8192..]
476        assert_eq!(convolver.tail_input.len(), 4096);
477        assert_eq!(convolver.tail_output0.len(), 4096);
478        assert_eq!(convolver.tail_precalculated0.len(), 4096);
479        assert_eq!(convolver.tail_output.len(), 4096);
480        assert_eq!(convolver.tail_precalculated.len(), 4096);
481    }
482
483    #[test]
484    fn init_short_ir_test() {
485        // IR shorter than head block size — only head convolver active
486        let mut convolver = TwoStageFFTConvolver::<f32>::default();
487        let ir = vec![0.5, 0.3, 0.2, 0.1];
488        convolver.init(64, 4096, &ir).unwrap();
489
490        assert!(convolver.tail_input.is_empty());
491        assert!(convolver.tail_output0.is_empty());
492        assert!(convolver.tail_precalculated0.is_empty());
493        assert!(convolver.tail_output.is_empty());
494        assert!(convolver.tail_precalculated.is_empty());
495    }
496
497    #[test]
498    fn init_medium_ir_test() {
499        // IR between tail_block_size and 2*tail_block_size — head + tail0 only
500        let mut convolver = TwoStageFFTConvolver::<f32>::default();
501        let ir = vec![0.5; 5000];
502        convolver.init(64, 4096, &ir).unwrap();
503
504        assert_eq!(convolver.tail_input.len(), 4096);
505        assert_eq!(convolver.tail_output0.len(), 4096);
506        assert_eq!(convolver.tail_precalculated0.len(), 4096);
507        // No tail convolver since ir_len < 2 * tail_block_size
508        assert!(convolver.tail_output.is_empty());
509        assert!(convolver.tail_precalculated.is_empty());
510    }
511
512    #[test]
513    fn init_block_size_zero_returns_error() {
514        let mut convolver = TwoStageFFTConvolver::<f32>::default();
515        let ir = vec![0.5; 100];
516        assert!(matches!(
517            convolver.init(0, 4096, &ir),
518            Err(FFTConvolverError::BlockSizeZero)
519        ));
520        assert!(matches!(
521            convolver.init(64, 0, &ir),
522            Err(FFTConvolverError::BlockSizeZero)
523        ));
524    }
525
526    #[test]
527    fn init_swaps_block_sizes_if_inverted() {
528        let mut convolver = TwoStageFFTConvolver::<f32>::default();
529        let ir = vec![0.5; 10000];
530        convolver.init(4096, 64, &ir).unwrap();
531
532        assert_eq!(convolver.head_block_size, 64);
533        assert_eq!(convolver.tail_block_size, 4096);
534    }
535
536    #[test]
537    fn reinit_with_shorter_ir_does_not_leave_stale_buffers() {
538        let short_ir = vec![0.5_f32; 4];
539
540        // Re-initialized convolver
541        let mut reinit = TwoStageFFTConvolver::<f32>::default();
542        let long_ir = vec![0.5_f32; 10000];
543        reinit.init(64, 4096, &long_ir).unwrap();
544        reinit.init(64, 4096, &short_ir).unwrap();
545
546        // Tail buffers must be empty
547        assert!(reinit.tail_input.is_empty());
548        assert!(reinit.tail_output0.is_empty());
549        assert!(reinit.tail_precalculated0.is_empty());
550        assert!(reinit.tail_output.is_empty());
551        assert!(reinit.tail_precalculated.is_empty());
552
553        // Output must match a freshly initialized convolver
554        let mut fresh = TwoStageFFTConvolver::<f32>::default();
555        fresh.init(64, 4096, &short_ir).unwrap();
556
557        let input = vec![1.0_f32; 256];
558        let mut output_reinit = vec![0.0_f32; 256];
559        let mut output_fresh = vec![0.0_f32; 256];
560
561        reinit.process(&input, &mut output_reinit).unwrap();
562        fresh.process(&input, &mut output_fresh).unwrap();
563
564        for i in 0..output_reinit.len() {
565            assert!(
566                (output_reinit[i] - output_fresh[i]).abs() < 1e-5,
567                "Mismatch at index {}: reinit produced {}, fresh produced {}",
568                i,
569                output_reinit[i],
570                output_fresh[i]
571            );
572        }
573    }
574
575    #[test]
576    fn process_test() {
577        let mut convolver = TwoStageFFTConvolver::<f32>::default();
578        // IR longer than 2*tail_block_size to exercise all stages
579        let mut ir = vec![0.0; 300];
580        ir[0] = 1.0;
581        convolver.init(4, 16, &ir).unwrap();
582
583        let input = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
584        let mut output = vec![0.0; 8];
585        convolver.process(&input, &mut output).unwrap();
586
587        for i in 0..output.len() {
588            assert!(
589                (input[i] - output[i]).abs() < 1e-5,
590                "Mismatch at index {}: input={}, output={}",
591                i,
592                input[i],
593                output[i]
594            );
595        }
596    }
597
598    #[test]
599    fn equivalence_test() {
600        // Two-stage convolver should produce same output as uniform convolver
601        let ir: Vec<f32> = (0..500).map(|i| 1.0 / (i as f32 + 1.0)).collect();
602
603        let mut uniform = FFTConvolver::<f32>::default();
604        uniform.init(64, &ir).unwrap();
605
606        let tail_block_size = compute_tail_block_size(64, ir.len());
607        let mut two_stage = TwoStageFFTConvolver::<f32>::default();
608        two_stage.init(64, tail_block_size, &ir).unwrap();
609
610        let input: Vec<f32> = (0..1024).map(|i| (i as f32 * 0.1).sin()).collect();
611        let mut output_uniform = vec![0.0; 1024];
612        let mut output_two_stage = vec![0.0; 1024];
613
614        uniform.process(&input, &mut output_uniform).unwrap();
615        two_stage.process(&input, &mut output_two_stage).unwrap();
616
617        for i in 0..output_uniform.len() {
618            assert!(
619                (output_uniform[i] - output_two_stage[i]).abs() < 1e-3,
620                "Mismatch at index {}: uniform={}, two_stage={}",
621                i,
622                output_uniform[i],
623                output_two_stage[i]
624            );
625        }
626    }
627
628    #[test]
629    fn test_zero_latency() {
630        let mut convolver = TwoStageFFTConvolver::<f32>::default();
631        let ir = vec![0.5, 0.3, 0.2, 0.1];
632        convolver.init(4, 16, &ir).unwrap();
633
634        let mut input = vec![0.0; 16];
635        input[0] = 1.0;
636
637        let mut output = vec![0.0; 16];
638        convolver.process(&input, &mut output).unwrap();
639
640        assert!(
641            output[0].abs() > 0.0,
642            "Output[0] should be non-zero, indicating zero latency. Got: {}",
643            output[0]
644        );
645        assert!(
646            (output[0] - 0.5).abs() < 1e-5,
647            "output[0] should be 0.5, got {}",
648            output[0]
649        );
650    }
651
652    #[test]
653    fn reset_test() {
654        let ir: Vec<f32> = (0..500).map(|i| 1.0 / (i as f32 + 1.0)).collect();
655
656        let mut convolver1 = TwoStageFFTConvolver::<f32>::default();
657        convolver1.init(64, 256, &ir).unwrap();
658
659        // Process some data to build up history
660        let history_input: Vec<f32> = (0..512).map(|i| (i as f32 * 0.1).sin()).collect();
661        let mut history_output = vec![0.0; 512];
662        convolver1
663            .process(&history_input, &mut history_output)
664            .unwrap();
665
666        convolver1.reset();
667
668        let test_input = vec![1.0; 256];
669        let mut output1 = vec![0.0; 256];
670        convolver1.process(&test_input, &mut output1).unwrap();
671
672        let mut convolver2 = TwoStageFFTConvolver::<f32>::default();
673        convolver2.init(64, 256, &ir).unwrap();
674        let mut output2 = vec![0.0; 256];
675        convolver2.process(&test_input, &mut output2).unwrap();
676
677        for i in 0..output1.len() {
678            assert!(
679                (output1[i] - output2[i]).abs() < 1e-5,
680                "Mismatch at index {}: cleared={}, fresh={}",
681                i,
682                output1[i],
683                output2[i]
684            );
685        }
686    }
687
688    #[test]
689    fn reset_preserves_configuration() {
690        let ir = vec![0.5; 10000];
691        let mut convolver = TwoStageFFTConvolver::<f32>::default();
692        convolver.init(64, 4096, &ir).unwrap();
693
694        let ir_len = convolver.ir_len;
695        let head_block_size = convolver.head_block_size;
696        let tail_block_size = convolver.tail_block_size;
697
698        let input = vec![1.0; 256];
699        let mut output = vec![0.0; 256];
700        convolver.process(&input, &mut output).unwrap();
701
702        convolver.reset();
703
704        assert_eq!(convolver.ir_len, ir_len);
705        assert_eq!(convolver.head_block_size, head_block_size);
706        assert_eq!(convolver.tail_block_size, tail_block_size);
707    }
708
709    #[test]
710    fn set_response_equals_init() {
711        let ir1 = vec![0.5; 10000];
712        let ir2: Vec<f32> = (0..10000).map(|i| 1.0 / (i as f32 + 1.0)).collect();
713
714        let mut convolver1 = TwoStageFFTConvolver::<f32>::default();
715        convolver1.init(64, 4096, &ir1).unwrap();
716        convolver1.set_response(&ir2).unwrap();
717
718        let mut convolver2 = TwoStageFFTConvolver::<f32>::default();
719        convolver2.init(64, 4096, &ir2).unwrap();
720
721        let input: Vec<f32> = (0..1024).map(|i| (i as f32 * 0.1).sin()).collect();
722        let mut output1 = vec![0.0; 1024];
723        let mut output2 = vec![0.0; 1024];
724
725        convolver1.process(&input, &mut output1).unwrap();
726        convolver2.process(&input, &mut output2).unwrap();
727
728        for i in 0..output1.len() {
729            assert!(
730                (output1[i] - output2[i]).abs() < 1e-3,
731                "Mismatch at index {}: set_response={}, init={}",
732                i,
733                output1[i],
734                output2[i]
735            );
736        }
737    }
738
739    #[test]
740    fn set_response_with_shorter_ir() {
741        let ir1 = vec![0.5; 10000];
742        let ir2 = vec![0.8; 5000];
743
744        let mut convolver1 = TwoStageFFTConvolver::<f32>::default();
745        convolver1.init(64, 4096, &ir1).unwrap();
746        convolver1.set_response(&ir2).unwrap();
747
748        let mut convolver2 = TwoStageFFTConvolver::<f32>::default();
749        convolver2.init(64, 4096, &ir2).unwrap();
750
751        let input: Vec<f32> = (0..1024).map(|i| (i as f32 * 0.1).sin()).collect();
752        let mut output1 = vec![0.0; 1024];
753        let mut output2 = vec![0.0; 1024];
754
755        convolver1.process(&input, &mut output1).unwrap();
756        convolver2.process(&input, &mut output2).unwrap();
757
758        for i in 0..output1.len() {
759            assert!(
760                (output1[i] - output2[i]).abs() < 1e-3,
761                "Mismatch at index {}: set_response={}, init={}",
762                i,
763                output1[i],
764                output2[i]
765            );
766        }
767    }
768
769    #[test]
770    fn set_response_too_long_returns_error() {
771        let ir1 = vec![0.5; 5000];
772        let ir2 = vec![0.8; 10000];
773
774        let mut convolver = TwoStageFFTConvolver::<f32>::default();
775        convolver.init(64, 4096, &ir1).unwrap();
776
777        let result = convolver.set_response(&ir2);
778        assert!(result.is_err());
779        assert!(matches!(
780            result.unwrap_err(),
781            FFTConvolverError::ImpulseResponseExceedsCapacity
782        ));
783    }
784
785    #[test]
786    fn test_varying_chunk_sizes() {
787        let ir: Vec<f32> = (0..500).map(|i| 1.0 / (i as f32 + 1.0)).collect();
788        let input: Vec<f32> = (0..1024).map(|i| (i as f32 * 0.1).sin()).collect();
789
790        // Process all at once
791        let mut convolver_ref = TwoStageFFTConvolver::<f32>::default();
792        convolver_ref.init(64, 256, &ir).unwrap();
793        let mut output_ref = vec![0.0; 1024];
794        convolver_ref.process(&input, &mut output_ref).unwrap();
795
796        // Process in varying chunk sizes
797        let chunk_sizes = [1, 7, 64, 100, 3, 256, 13, 580];
798        let mut convolver = TwoStageFFTConvolver::<f32>::default();
799        convolver.init(64, 256, &ir).unwrap();
800        let mut output = vec![0.0; 1024];
801        let mut pos = 0;
802        for &chunk in &chunk_sizes {
803            let end = (pos + chunk).min(1024);
804            if pos >= end {
805                break;
806            }
807            convolver
808                .process(&input[pos..end], &mut output[pos..end])
809                .unwrap();
810            pos = end;
811        }
812
813        for i in 0..output_ref.len() {
814            assert!(
815                (output_ref[i] - output[i]).abs() < 1e-3,
816                "Mismatch at index {}: ref={}, chunked={}",
817                i,
818                output_ref[i],
819                output[i]
820            );
821        }
822    }
823
824    #[test]
825    fn process_mismatched_lengths_returns_error() {
826        let mut convolver = TwoStageFFTConvolver::<f32>::default();
827        let ir = vec![0.5; 10000];
828        convolver.init(64, 4096, &ir).unwrap();
829
830        let input = vec![1.0_f32; 64];
831        let mut output = vec![0.0_f32; 128];
832        let result = convolver.process(&input, &mut output);
833        assert!(matches!(
834            result.unwrap_err(),
835            FFTConvolverError::InputOutputLengthMismatch
836        ));
837    }
838
839    #[test]
840    fn test_short_ir_degenerate() {
841        // IR shorter than head block — should behave like plain FFTConvolver
842        let ir = vec![0.5, 0.3, 0.2, 0.1];
843
844        let mut uniform = FFTConvolver::<f32>::default();
845        uniform.init(64, &ir).unwrap();
846
847        let mut two_stage = TwoStageFFTConvolver::<f32>::default();
848        two_stage.init(64, 256, &ir).unwrap();
849
850        let input = vec![1.0; 256];
851        let mut output_uniform = vec![0.0; 256];
852        let mut output_two_stage = vec![0.0; 256];
853
854        uniform.process(&input, &mut output_uniform).unwrap();
855        two_stage.process(&input, &mut output_two_stage).unwrap();
856
857        for i in 0..output_uniform.len() {
858            assert!(
859                (output_uniform[i] - output_two_stage[i]).abs() < 1e-5,
860                "Mismatch at index {}: uniform={}, two_stage={}",
861                i,
862                output_uniform[i],
863                output_two_stage[i]
864            );
865        }
866    }
867
868    #[test]
869    fn test_small_head_block_size() {
870        // head_block_size=1 (rounds to 1); ensures no off-by-one in buffering
871        let ir = vec![0.5_f32, 0.3, 0.2, 0.1];
872        let mut convolver = TwoStageFFTConvolver::<f32>::default();
873        convolver.init(1, 16, &ir).unwrap();
874        assert_eq!(convolver.head_block_size, 1);
875
876        let mut input = vec![0.0_f32; 16];
877        input[0] = 1.0;
878        let mut output = vec![0.0_f32; 16];
879        convolver.process(&input, &mut output).unwrap();
880
881        assert!((output[0] - 0.5).abs() < 1e-5);
882        assert!((output[1] - 0.3).abs() < 1e-5);
883        assert!((output[2] - 0.2).abs() < 1e-5);
884        assert!((output[3] - 0.1).abs() < 1e-5);
885    }
886
887    #[test]
888    fn test_large_ir_equivalence() {
889        // Long IR — all three convolver stages active; compare against uniform convolver
890        let ir_len = 16384_usize;
891        let mut ir = vec![0.0_f32; ir_len];
892        ir[0] = 1.0; // identity
893
894        let mut uniform = FFTConvolver::<f32>::default();
895        uniform.init(512, &ir).unwrap();
896
897        let tail_block_size = compute_tail_block_size(512, ir_len);
898        let mut two_stage = TwoStageFFTConvolver::<f32>::default();
899        two_stage.init(512, tail_block_size, &ir).unwrap();
900
901        let input: Vec<f32> = (0..2048).map(|i| (i as f32 * 0.01).sin()).collect();
902        let mut out_uniform = vec![0.0_f32; 2048];
903        let mut out_two_stage = vec![0.0_f32; 2048];
904
905        uniform.process(&input, &mut out_uniform).unwrap();
906        two_stage.process(&input, &mut out_two_stage).unwrap();
907
908        for i in 0..input.len() {
909            assert!(
910                (out_uniform[i] - out_two_stage[i]).abs() < 1e-3,
911                "Mismatch at {}: uniform={}, two_stage={}",
912                i,
913                out_uniform[i],
914                out_two_stage[i]
915            );
916        }
917    }
918}
fft_convolver/two_stage.rs

fft_convolver/
two_stage.rs
