rustradio/
fft.rs

1use std::sync::Arc;
2
3use crate::block::{Block, BlockRet};
4use crate::stream::{NCReadStream, NCWriteStream};
5use crate::{Complex, Error, Float, Result};
6
7/// Run FFT on message.
8#[derive(rustradio_macros::Block)]
9#[rustradio(crate, new)]
10pub struct Fft {
11    #[rustradio(in)]
12    src: NCReadStream<Vec<Complex>>,
13    #[rustradio(out)]
14    dst: NCWriteStream<Vec<Complex>>,
15    fft: Arc<dyn rustfft::Fft<Float>>,
16}
17
18impl Fft {
19    fn process_one(&mut self, input: &[Complex]) -> Vec<Complex> {
20        let mut out = input.to_vec();
21        self.fft.process(&mut out);
22        out
23    }
24}
25
26impl Block for Fft {
27    fn work(&mut self) -> Result<BlockRet<'_>> {
28        loop {
29            if self.dst.remaining() == 0 {
30                return Ok(BlockRet::WaitForStream(&self.dst, 1));
31            }
32            let Some((msg, tags)) = self.src.pop() else {
33                return Ok(BlockRet::WaitForStream(&self.src, 1));
34            };
35            if msg.len() != self.fft.len() {
36                return Err(Error::msg(format!(
37                    "FFT expected {} samples, got {}",
38                    self.fft.len(),
39                    msg.len()
40                )));
41            }
42            let out = self.process_one(&msg);
43            self.dst.push(out, tags);
44        }
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51    use crate::stream::new_nocopy_stream;
52
53    #[test]
54    fn zeroes() -> Result<()> {
55        let mut planner = rustfft::FftPlanner::new();
56        let size = 1024;
57        let fft = planner.plan_fft_forward(size);
58        let (root, r) = new_nocopy_stream();
59        let (mut f, out) = Fft::new(r, fft);
60        assert!(out.pop().is_none());
61        assert!(matches![f.work()?, BlockRet::WaitForStream(_, 1)]);
62        assert!(out.pop().is_none());
63        root.push(vec![Complex::default(); size], &[]);
64        assert!(matches![f.work()?, BlockRet::WaitForStream(_, 1)]);
65        // Get the results.
66        let (omsg, tags) = out.pop().unwrap();
67        assert_eq!(omsg.len(), size);
68        assert_eq!(omsg, vec![Complex::default(); size]);
69        assert_eq!(tags, &[]);
70
71        // Should be empty now.
72        assert!(out.pop().is_none());
73        Ok(())
74    }
75
76    #[test]
77    fn rejects_wrong_size() -> Result<()> {
78        let mut planner = rustfft::FftPlanner::new();
79        let fft = planner.plan_fft_forward(4);
80        let (root, r) = new_nocopy_stream();
81        let (mut f, out) = Fft::new(r, fft);
82        root.push(vec![Complex::default(); 3], &[]);
83
84        let err = f.work().unwrap_err();
85        assert!(err.to_string().contains("expected 4 samples"));
86        assert!(out.pop().is_none());
87        Ok(())
88    }
89}
rustradio/fft.rs

rustradio/
fft.rs
