dspk_core/window/
loader.rs

1//! Double-buffered window reading over a [`RecordingSource`].
2//!
3//! A scoped background thread reads window `i + 1` (disk, decompression) while the caller
4//! processes window `i`. Two buffers are recycled between the threads, so host memory stays at two
5//! windows whatever the recording length, with no allocation per window.
6
7use std::ops::Range;
8use std::sync::mpsc::sync_channel;
9
10use super::HaloWindow;
11use crate::{DspResult, RecordingSource};
12
13/// Buffers in flight between the reading thread and the caller.
14const BUFFERS_IN_FLIGHT: usize = 2;
15
16/// Streams [`HaloWindow`]s of a [`RecordingSource`], reading ahead on a background thread.
17pub struct WindowLoader<'a> {
18    source: &'a dyn RecordingSource,
19    channels: Vec<usize>,
20}
21
22impl<'a> WindowLoader<'a> {
23    /// Loader over all channels of `source`.
24    pub fn new(source: &'a dyn RecordingSource) -> Self {
25        let channels = (0..source.info().channel_count()).collect();
26        Self { source, channels }
27    }
28
29    /// Loader over `channels` of `source`, in that order.
30    pub fn with_channels(source: &'a dyn RecordingSource, channels: Vec<usize>) -> Self {
31        Self { source, channels }
32    }
33
34    /// The recording read.
35    #[inline]
36    pub fn source(&self) -> &'a dyn RecordingSource {
37        self.source
38    }
39
40    /// Channels read, in buffer order.
41    #[inline]
42    pub fn channels(&self) -> &[usize] {
43        &self.channels
44    }
45
46    /// Calls `f(&window, samples)` for every window of `windows`, in order: `samples` is the
47    /// channel-major `[channels, window.read_len()]` block of scaled values
48    /// ([`RecordingSource::read`]). Any subset works (a full [`super::ChunkSchedule`], every
49    /// n-th window, isolated windows); the next window is read while `f` runs.
50    pub fn stream<F>(&self, windows: &[HaloWindow], mut f: F) -> DspResult<()>
51    where
52        F: FnMut(&HaloWindow, &[f32]) -> DspResult<()>,
53    {
54        self.stream_while(windows, |win, buf| f(win, buf).map(|()| true))
55    }
56
57    /// Like [`Self::stream`] until `f` returns `false` (e.g. once enough data is gathered): no
58    /// window after that one is passed to `f`, and reading stops.
59    pub fn stream_while<F>(&self, windows: &[HaloWindow], f: F) -> DspResult<()>
60    where
61        F: FnMut(&HaloWindow, &[f32]) -> DspResult<bool>,
62    {
63        let source = self.source;
64        self.stream_with(windows, 1, |ch, range, buf: &mut [f32]| source.read(ch, range, buf), f)
65    }
66
67    /// Like [`Self::stream`] with the stored values ([`RecordingSource::read_stored`]):
68    /// `info().format` little-endian bytes, channel-major, before gain and offset.
69    pub fn stream_stored<F>(&self, windows: &[HaloWindow], mut f: F) -> DspResult<()>
70    where
71        F: FnMut(&HaloWindow, &[u8]) -> DspResult<()>,
72    {
73        self.stream_stored_while(windows, |win, buf| f(win, buf).map(|()| true))
74    }
75
76    /// Like [`Self::stream_stored`] until `f` returns `false`.
77    pub fn stream_stored_while<F>(&self, windows: &[HaloWindow], f: F) -> DspResult<()>
78    where
79        F: FnMut(&HaloWindow, &[u8]) -> DspResult<bool>,
80    {
81        let source = self.source;
82        let bytes = source.info().format.bytes();
83        let read = |ch: &[usize], range, buf: &mut [u8]| source.read_stored(ch, range, buf);
84        self.stream_with(windows, bytes, read, f)
85    }
86
87    /// Like [`Self::stream_while`], handing `f` each window as an **owned** buffer: the reading
88    /// thread allocates and fills it, so a consumer that moves it on (e.g. into a device upload)
89    /// copies nothing on the calling thread. At most one window is read ahead.
90    pub fn stream_owned_while<F>(&self, windows: &[HaloWindow], f: F) -> DspResult<()>
91    where
92        F: FnMut(&HaloWindow, Vec<f32>) -> DspResult<bool>,
93    {
94        let source = self.source;
95        self.stream_owned(windows, 1, |ch: &[usize], range, buf: &mut [f32]| source.read(ch, range, buf), f)
96    }
97
98    /// [`Self::stream_owned_while`] with the stored values ([`RecordingSource::read_stored`]).
99    pub fn stream_stored_owned_while<F>(&self, windows: &[HaloWindow], f: F) -> DspResult<()>
100    where
101        F: FnMut(&HaloWindow, Vec<u8>) -> DspResult<bool>,
102    {
103        let source = self.source;
104        let bytes = source.info().format.bytes();
105        self.stream_owned(windows, bytes, |ch: &[usize], range, buf: &mut [u8]| source.read_stored(ch, range, buf), f)
106    }
107
108    /// Read-ahead loop handing over a fresh buffer per window (`per_sample` elements per channel
109    /// sample), until `f` returns `false`.
110    fn stream_owned<T, R, F>(&self, windows: &[HaloWindow], per_sample: usize, read: R, mut f: F) -> DspResult<()>
111    where
112        T: Copy + Default + Send,
113        R: Fn(&[usize], Range<u64>, &mut [T]) -> DspResult<()> + Sync,
114        F: FnMut(&HaloWindow, Vec<T>) -> DspResult<bool>,
115    {
116        if windows.is_empty() || self.channels.is_empty() {
117            return Ok(());
118        }
119        let n_ch = self.channels.len();
120        let channels = &self.channels;
121        let read = &read;
122        std::thread::scope(|s| {
123            // Capacity 1: the reader fills the next window while the caller holds the current one
124            let (ready_tx, ready_rx) = sync_channel::<DspResult<(usize, Vec<T>)>>(1);
125            s.spawn(move || {
126                for (i, win) in windows.iter().enumerate() {
127                    let mut buf = vec![T::default(); n_ch * win.read_len() * per_sample];
128                    let item = read(channels, win.read_global.clone(), &mut buf).map(|()| (i, buf));
129                    let failed = item.is_err();
130                    if ready_tx.send(item).is_err() || failed {
131                        break;
132                    }
133                }
134            });
135            // Returning early drops the receiver, which stops the reading thread
136            for item in ready_rx {
137                let (i, buf) = item?;
138                if !f(&windows[i], buf)? {
139                    break;
140                }
141            }
142            Ok(())
143        })
144    }
145
146    /// Double-buffered loop over `per_sample` elements per channel sample, until `f` returns
147    /// `false`.
148    fn stream_with<T, R, F>(&self, windows: &[HaloWindow], per_sample: usize, read: R, mut f: F) -> DspResult<()>
149    where
150        T: Copy + Default + Send,
151        R: Fn(&[usize], Range<u64>, &mut [T]) -> DspResult<()> + Sync,
152        F: FnMut(&HaloWindow, &[T]) -> DspResult<bool>,
153    {
154        if windows.is_empty() || self.channels.is_empty() {
155            return Ok(());
156        }
157        let n_ch = self.channels.len();
158        let max_len = n_ch * windows.iter().map(HaloWindow::read_len).max().unwrap_or(0) * per_sample;
159        let channels = &self.channels;
160        let read = &read;
161
162        std::thread::scope(|s| {
163            let (ready_tx, ready_rx) = sync_channel::<DspResult<(usize, Vec<T>)>>(1);
164            let (recycle_tx, recycle_rx) = sync_channel::<Vec<T>>(BUFFERS_IN_FLIGHT);
165            for _ in 0..BUFFERS_IN_FLIGHT {
166                let _ = recycle_tx.send(vec![T::default(); max_len]);
167            }
168
169            s.spawn(move || {
170                for (i, win) in windows.iter().enumerate() {
171                    let Ok(mut buf) = recycle_rx.recv() else { break };
172                    let len = n_ch * win.read_len() * per_sample;
173                    buf.resize(len, T::default());
174                    let item = read(channels, win.read_global.clone(), &mut buf).map(|()| (i, buf));
175                    let failed = item.is_err();
176                    if ready_tx.send(item).is_err() || failed {
177                        break;
178                    }
179                }
180            });
181
182            // Returning early drops both channel ends, which stops the reading thread
183            for item in ready_rx {
184                let (i, buf) = item?;
185                if !f(&windows[i], &buf)? {
186                    break;
187                }
188                let _ = recycle_tx.send(buf);
189            }
190            Ok(())
191        })
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use crate::window::ChunkSchedule;
199    use crate::MemoryRecording;
200
201    fn recording() -> MemoryRecording {
202        let data = (0..2).flat_map(|c| (0..100).map(move |t| (c * 1000 + t) as f32)).collect();
203        MemoryRecording::new("test", data, 2, 1000.0).unwrap()
204    }
205
206    #[test]
207    fn streams_a_full_schedule() {
208        let rec = recording();
209        let sched = ChunkSchedule::full_recording(100, 30, 5, 5);
210        let mut visited = Vec::new();
211        WindowLoader::new(&rec)
212            .stream(sched.windows(), |win, buf| {
213                let n = win.read_len();
214                visited.push((win.index, win.valid_global.clone(), n, buf[win.valid_local.start], buf[n]));
215                Ok(())
216            })
217            .unwrap();
218        assert_eq!(visited.len(), 4);
219        assert_eq!(visited[0], (0, 0..30, 35, 0.0, 1000.0));
220        assert_eq!(visited[1], (1, 30..60, 40, 30.0, 1025.0));
221        assert_eq!(visited[3], (3, 90..100, 15, 90.0, 1085.0));
222    }
223
224    /// The owned stream hands over the same windows as the borrowed one, and stops when asked.
225    #[test]
226    fn owned_stream_matches_borrowed_stream() {
227        let rec = recording();
228        let sched = ChunkSchedule::full_recording(100, 30, 5, 5);
229        let mut borrowed = Vec::new();
230        WindowLoader::new(&rec).stream(sched.windows(), |_, buf| {
231            borrowed.push(buf.to_vec());
232            Ok(())
233        })
234        .unwrap();
235        let mut owned = Vec::new();
236        WindowLoader::new(&rec)
237            .stream_owned_while(sched.windows(), |win, buf| {
238                owned.push(buf);
239                Ok(win.index < 2)
240            })
241            .unwrap();
242        assert_eq!(owned.len(), 3, "stopped after the window that returned false");
243        assert_eq!(owned[..], borrowed[..3]);
244    }
245
246    #[test]
247    fn streams_a_subset_and_stops_on_error() {
248        let rec = recording();
249        let sched = ChunkSchedule::full_recording(100, 30, 5, 5);
250        let subset: Vec<HaloWindow> = sched.windows().iter().step_by(2).cloned().collect();
251        let mut seen = Vec::new();
252        WindowLoader::new(&rec)
253            .stream(&subset, |win, _| {
254                seen.push(win.index);
255                Ok(())
256            })
257            .unwrap();
258        assert_eq!(seen, vec![0, 2]);
259
260        let mut calls = 0;
261        let err = WindowLoader::new(&rec).stream(sched.windows(), |_, _| {
262            calls += 1;
263            Err(crate::DspError::InvalidConfig("stop".into()))
264        });
265        assert!(err.is_err());
266        assert_eq!(calls, 1);
267
268        let mut seen = Vec::new();
269        WindowLoader::new(&rec)
270            .stream_while(sched.windows(), |win, _| {
271                seen.push(win.index);
272                Ok(win.index < 1)
273            })
274            .unwrap();
275        assert_eq!(seen, vec![0, 1]);
276    }
277}
dspk_core/window/loader.rs

dspk_core/window/
loader.rs
