rusty_esp_audio_core/
pipeline.rs

1//! `Element` and `Pipeline`: ESP-ADF's `audio_element` / `audio_pipeline`
2//! remade as a fixed-block chain over two halves of one caller scratch buffer.
3//!
4//! Every element takes one block and writes its result into caller memory.
5//! The pipeline ping-pongs between the two halves of `scratch`, so a chain
6//! of any length needs exactly two blocks of scratch and no heap. An element
7//! may change the format (rate, channels, encoding) and may produce zero
8//! bytes for a block (a resampler priming, a gate closed); the pipeline
9//! reports `Ok(None)` for that block and the caller moves on.
10
11use rusty_esp_core::error::{Error, Result};
12use rusty_esp_core::pcm::{PcmBlock, PcmFormat};
13
14/// One processing stage.
15pub trait Element {
16    /// The format this element emits for blocks in `input`, or
17    /// `Err(Unsupported)` if it cannot take `input` at all.
18    fn output_format(&self, input: PcmFormat) -> Result<PcmFormat>;
19
20    /// Upper bound on the bytes [`process`](Self::process) writes for an input
21    /// block of `input_bytes` bytes in `input`. Defaults to `input_bytes`
22    /// (same-size elements); resamplers and channel converters override it.
23    fn max_output_bytes(&self, input: PcmFormat, input_bytes: usize) -> usize {
24        let _ = input;
25        input_bytes
26    }
27
28    /// Process one block into `out`, returning the bytes written — a whole
29    /// number of output frames, possibly zero.
30    fn process(&mut self, input: PcmBlock<'_>, out: &mut [u8]) -> Result<usize>;
31
32    /// Forget all state (filter memories, gains, phases).
33    fn reset(&mut self) {}
34}
35
36/// A fixed chain of `N` elements.
37pub struct Pipeline<'e, const N: usize> {
38    stages: [&'e mut dyn Element; N],
39    /// Blocks processed.
40    pub blocks: u64,
41    /// Blocks that produced no output.
42    pub empty_blocks: u64,
43}
44
45impl<'e, const N: usize> core::fmt::Debug for Pipeline<'e, N> {
46    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
47        f.debug_struct("Pipeline")
48            .field("stages", &N)
49            .field("blocks", &self.blocks)
50            .field("empty_blocks", &self.empty_blocks)
51            .finish()
52    }
53}
54
55impl<'e, const N: usize> Pipeline<'e, N> {
56    /// Chain `stages` in order.
57    pub fn new(stages: [&'e mut dyn Element; N]) -> Self {
58        Pipeline {
59            stages,
60            blocks: 0,
61            empty_blocks: 0,
62        }
63    }
64
65    /// The format that comes out for blocks in `input`.
66    pub fn output_format(&self, input: PcmFormat) -> Result<PcmFormat> {
67        self.stages
68            .iter()
69            .try_fold(input, |f, s| s.output_format(f))
70    }
71
72    /// Bytes of `scratch` needed for an input of `input_bytes` bytes in
73    /// `input`: twice the largest intermediate block.
74    pub fn scratch_bytes(&self, input: PcmFormat, input_bytes: usize) -> Result<usize> {
75        let mut fmt = input;
76        let mut bytes = input_bytes;
77        let mut largest = input_bytes;
78        for s in &self.stages {
79            bytes = s.max_output_bytes(fmt, bytes);
80            fmt = s.output_format(fmt)?;
81            largest = largest.max(bytes);
82        }
83        Ok(largest * 2)
84    }
85
86    /// Run one block through every stage. The result borrows `scratch`;
87    /// `Ok(None)` means a stage produced nothing for this block.
88    pub fn process<'b>(
89        &mut self,
90        input: PcmBlock<'_>,
91        scratch: &'b mut [u8],
92    ) -> Result<Option<PcmBlock<'b>>> {
93        self.blocks += 1;
94        let half = scratch.len() / 2;
95        let (a, b) = scratch.split_at_mut(half);
96        let ts = input.timestamp;
97        let mut fmt = input.format;
98
99        if N == 0 {
100            if a.len() < input.data.len() {
101                return Err(Error::BufferTooSmall {
102                    needed: input.data.len() * 2,
103                });
104            }
105            a[..input.data.len()].copy_from_slice(input.data);
106            return Ok(Some(PcmBlock::new(fmt, ts, &a[..input.data.len()])?));
107        }
108
109        // Where the current data lives: `true` = in `a`.
110        let mut in_a = false;
111        let mut len = 0usize;
112        for (i, stage) in self.stages.iter_mut().enumerate() {
113            let out_fmt = stage.output_format(fmt)?;
114            let in_len = if i == 0 { input.data.len() } else { len };
115            let need = stage.max_output_bytes(fmt, in_len);
116            if half < need {
117                return Err(Error::BufferTooSmall { needed: need * 2 });
118            }
119            let n = if i == 0 {
120                stage.process(input, a)?
121            } else if in_a {
122                let blk = PcmBlock::new(fmt, ts, &a[..len])?;
123                stage.process(blk, b)?
124            } else {
125                let blk = PcmBlock::new(fmt, ts, &b[..len])?;
126                stage.process(blk, a)?
127            };
128            if n % out_fmt.frame_bytes() != 0 || n > half {
129                return Err(Error::InvalidGeometry);
130            }
131            in_a = if i == 0 { true } else { !in_a };
132            fmt = out_fmt;
133            len = n;
134            if n == 0 {
135                self.empty_blocks += 1;
136                return Ok(None);
137            }
138        }
139        let data: &'b [u8] = if in_a { &a[..len] } else { &b[..len] };
140        Ok(Some(PcmBlock::new(fmt, ts, data)?))
141    }
142
143    /// Reset every stage.
144    pub fn reset(&mut self) {
145        for s in self.stages.iter_mut() {
146            s.reset();
147        }
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use crate::elements::{Gain, MonoToStereo};
155    use crate::{get_i16, put_i16};
156    use rusty_esp_core::time::Micros;
157
158    /// Doubles every sample; the simplest stateful-looking stage.
159    struct Double;
160    impl Element for Double {
161        fn output_format(&self, input: PcmFormat) -> Result<PcmFormat> {
162            Ok(input)
163        }
164        fn process(&mut self, input: PcmBlock<'_>, out: &mut [u8]) -> Result<usize> {
165            for (i, o) in input.data.chunks_exact(2).zip(out.chunks_exact_mut(2)) {
166                put_i16(o, get_i16(i).saturating_mul(2));
167            }
168            Ok(input.data.len())
169        }
170    }
171
172    /// Emits nothing on odd blocks.
173    struct Gate(u64);
174    impl Element for Gate {
175        fn output_format(&self, input: PcmFormat) -> Result<PcmFormat> {
176            Ok(input)
177        }
178        fn process(&mut self, input: PcmBlock<'_>, out: &mut [u8]) -> Result<usize> {
179            self.0 += 1;
180            if self.0 % 2 == 0 {
181                return Ok(0);
182            }
183            out[..input.data.len()].copy_from_slice(input.data);
184            Ok(input.data.len())
185        }
186    }
187
188    #[test]
189    fn three_stages_ping_pong_and_change_format() {
190        let f = PcmFormat::PCM16_16K_MONO;
191        let mut d1 = Double;
192        let mut d2 = Double;
193        let mut up = MonoToStereo;
194        let mut p = Pipeline::new([&mut d1 as &mut dyn Element, &mut d2, &mut up]);
195        let out_fmt = p.output_format(f).unwrap();
196        assert_eq!(out_fmt.channels, 2);
197        let mut input = [0u8; 8];
198        for (i, s) in input.chunks_exact_mut(2).enumerate() {
199            put_i16(s, i as i16 + 1);
200        }
201        let need = p.scratch_bytes(f, input.len()).unwrap();
202        assert_eq!(need, 32);
203        let mut scratch = [0u8; 32];
204        let blk = PcmBlock::new(f, Micros(5), &input).unwrap();
205        let out = p.process(blk, &mut scratch).unwrap().unwrap();
206        assert_eq!(out.format, out_fmt);
207        assert_eq!(out.timestamp, Micros(5));
208        let v: [i16; 8] = core::array::from_fn(|i| get_i16(&out.data[i * 2..]));
209        assert_eq!(v, [4, 4, 8, 8, 12, 12, 16, 16]);
210        assert_eq!(p.blocks, 1);
211    }
212
213    #[test]
214    fn empty_output_and_small_scratch_are_reported() {
215        let f = PcmFormat::PCM16_16K_MONO;
216        let mut g = Gate(0);
217        let mut unity = Gain::linear(1.0);
218        let mut p = Pipeline::new([&mut g as &mut dyn Element, &mut unity]);
219        let input = [1u8, 0, 2, 0];
220        let mut scratch = [0u8; 8];
221        let blk = PcmBlock::new(f, Micros::ZERO, &input).unwrap();
222        assert!(p.process(blk, &mut scratch).unwrap().is_some());
223        assert!(p.process(blk, &mut scratch).unwrap().is_none());
224        assert_eq!(p.empty_blocks, 1);
225        let mut tiny = [0u8; 6];
226        assert_eq!(
227            p.process(blk, &mut tiny).err(),
228            Some(Error::BufferTooSmall { needed: 8 })
229        );
230    }
231
232    #[test]
233    fn zero_stages_copies() {
234        let f = PcmFormat::PCM16_16K_MONO;
235        let mut p: Pipeline<'_, 0> = Pipeline::new([]);
236        let input = [1u8, 0, 2, 0];
237        let mut scratch = [0u8; 8];
238        let blk = PcmBlock::new(f, Micros::ZERO, &input).unwrap();
239        let out = p.process(blk, &mut scratch).unwrap().unwrap();
240        assert_eq!(out.data, &input);
241    }
242}
rusty_esp_audio_core/pipeline.rs

rusty_esp_audio_core/
pipeline.rs
