Skip to main content

dspk_core/window/
schedule.rs

1//! Halo window scheduling: which samples each window owns and which it reads.
2//!
3//! When streaming a large recording through causal/non-causal processing, each non-overlapping
4//! valid window `[s0 .. s1)` is padded with a `left_halo` (e.g. filter settling, lookback) and a
5//! `right_halo` (e.g. lookahead), clamped to `[0 .. total_samples)`.
6
7use std::ops::Range;
8
9/// A single halo-padded window in a chunked streaming schedule.
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct HaloWindow {
12    /// 0-based window index in the schedule.
13    pub index: usize,
14    /// Non-overlapping global sample range `[s0 .. s1)` owned by this window.
15    pub valid_global: Range<u64>,
16    /// Halo-padded global sample range `[read_start .. read_end)` to read from `RecordingSource`.
17    pub read_global: Range<u64>,
18    /// Interior slice range `[left_pad .. left_pad + valid_len)` within the padded buffer.
19    pub valid_local: Range<usize>,
20}
21
22impl HaloWindow {
23    /// Total number of samples in the padded read buffer (`read_global`).
24    #[inline]
25    pub fn read_len(&self) -> usize {
26        (self.read_global.end - self.read_global.start) as usize
27    }
28
29    /// Number of valid interior samples (`valid_global`).
30    #[inline]
31    pub fn valid_len(&self) -> usize {
32        self.valid_local.len()
33    }
34
35    /// Returns `true` if a local sample index (relative to the padded buffer) lies inside
36    /// the valid interior owned by this window.
37    #[inline]
38    pub fn is_interior_local(&self, local_sample: usize) -> bool {
39        self.valid_local.contains(&local_sample)
40    }
41
42    /// Converts a local sample index (relative to the padded buffer) to a global recording sample index.
43    #[inline]
44    pub fn to_global_sample(&self, local_sample: usize) -> u64 {
45        self.read_global.start + local_sample as u64
46    }
47
48    /// Global sample of a local event (e.g. a spike found in the padded buffer) if this window
49    /// owns it (it lies in the valid interior), `None` if it belongs to a neighbouring window.
50    #[inline]
51    pub fn remap_event(&self, local_sample: usize) -> Option<u64> {
52        self.is_interior_local(local_sample).then(|| self.to_global_sample(local_sample))
53    }
54
55    /// The window owning `valid` (clamped to `[0 .. total_samples)`), read with `left_halo` and
56    /// `right_halo` samples of context, also clamped to the recording.
57    pub fn around(index: usize, valid: Range<u64>, left_halo: u64, right_halo: u64, total_samples: u64) -> Self {
58        let end = valid.end.min(total_samples);
59        let start = valid.start.min(end);
60        let read_start = start.saturating_sub(left_halo);
61        let read_end = (end + right_halo).min(total_samples);
62        let left_pad = (start - read_start) as usize;
63        Self {
64            index,
65            valid_global: start..end,
66            read_global: read_start..read_end,
67            valid_local: left_pad..left_pad + (end - start) as usize,
68        }
69    }
70}
71
72/// Schedule of contiguous `HaloWindow`s covering a sample range `[start .. end)` of a recording.
73#[derive(Debug, Clone)]
74pub struct ChunkSchedule {
75    windows: Vec<HaloWindow>,
76    max_read_samples: usize,
77}
78
79impl ChunkSchedule {
80    /// Builds a schedule over `target_range` within a recording of length `total_samples`.
81    pub fn new(
82        target_range: Range<u64>,
83        batch_samples: u64,
84        left_halo: u64,
85        right_halo: u64,
86        total_samples: u64,
87    ) -> Self {
88        let start = target_range.start.min(total_samples);
89        let end = target_range.end.min(total_samples);
90        let step = batch_samples.max(1);
91
92        let mut windows = Vec::new();
93        let mut max_read_samples = 0usize;
94        let mut s0 = start;
95        let mut index = 0usize;
96
97        while s0 < end {
98            let s1 = (s0 + step).min(end);
99            let win = HaloWindow::around(index, s0..s1, left_halo, right_halo, total_samples);
100            max_read_samples = max_read_samples.max(win.read_len());
101            windows.push(win);
102
103            s0 = s1;
104            index += 1;
105        }
106
107        Self {
108            windows,
109            max_read_samples,
110        }
111    }
112
113    /// Builds a schedule covering the entire recording `0 .. total_samples`.
114    pub fn full_recording(
115        total_samples: u64,
116        batch_samples: u64,
117        left_halo: u64,
118        right_halo: u64,
119    ) -> Self {
120        Self::new(
121            0..total_samples,
122            batch_samples,
123            left_halo,
124            right_halo,
125            total_samples,
126        )
127    }
128
129    /// Slice of scheduled `HaloWindow`s.
130    pub fn windows(&self) -> &[HaloWindow] {
131        &self.windows
132    }
133
134    /// Consumes the schedule and returns the underlying vector of windows.
135    pub fn into_windows(self) -> Vec<HaloWindow> {
136        self.windows
137    }
138
139    /// Number of scheduled windows.
140    pub fn len(&self) -> usize {
141        self.windows.len()
142    }
143
144    /// Returns `true` if no windows are scheduled.
145    pub fn is_empty(&self) -> bool {
146        self.windows.is_empty()
147    }
148
149    /// Maximum `read_len()` across all windows (used to pre-allocate persistent VRAM buffers).
150    pub fn max_read_samples(&self) -> usize {
151        self.max_read_samples
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn test_chunk_schedule_partitions_without_gaps_or_overlap() {
161        let total = 10_000u64;
162        let sched = ChunkSchedule::full_recording(total, 3_000, 500, 200);
163        assert_eq!(sched.len(), 4);
164        assert_eq!(sched.max_read_samples(), 3_700);
165
166        let w = sched.windows();
167        // Window 0: valid 0..3000, read 0..3200, local 0..3000
168        assert_eq!(w[0].valid_global, 0..3_000);
169        assert_eq!(w[0].read_global, 0..3_200);
170        assert_eq!(w[0].valid_local, 0..3_000);
171        assert_eq!(w[0].to_global_sample(w[0].valid_local.start), 0);
172
173        // Window 1: valid 3000..6000, read 2500..6200, local 500..3500
174        assert_eq!(w[1].valid_global, 3_000..6_000);
175        assert_eq!(w[1].read_global, 2_500..6_200);
176        assert_eq!(w[1].valid_local, 500..3_500);
177        assert!(!w[1].is_interior_local(499));
178        assert!(w[1].is_interior_local(500));
179        assert_eq!(w[1].to_global_sample(500), 3_000);
180        assert_eq!(w[1].to_global_sample(3_499), 5_999);
181
182        // Window 3 (last): valid 9000..10000, read 8500..10000, local 500..1500
183        assert_eq!(w[3].valid_global, 9_000..10_000);
184        assert_eq!(w[3].read_global, 8_500..10_000);
185        assert_eq!(w[3].valid_local, 500..1_500);
186        assert_eq!(w[3].to_global_sample(1_499), 9_999);
187    }
188
189    #[test]
190    fn around_clamps_and_remaps_events() {
191        let w = HaloWindow::around(7, 30..60, 5, 5, 62);
192        assert_eq!((w.index, w.read_global.clone(), w.valid_local.clone()), (7, 25..62, 5..35));
193        assert_eq!(w.remap_event(4), None);
194        assert_eq!(w.remap_event(5), Some(30));
195        assert_eq!(w.remap_event(34), Some(59));
196        assert_eq!(w.remap_event(35), None);
197        assert_eq!(HaloWindow::around(0, 90..120, 5, 5, 100).valid_global, 90..100);
198    }
199}