Skip to main content

miniextendr_api/altrep_data/iter/
windowed.rs

1//! Windowed iterator-backed ALTREP data adaptors.
2//!
3//! Provides `WindowedIterState<I, T>` which keeps a sliding window of elements
4//! in memory, and data-adaptor types for the integer/real ALTREP families.
5//!
6//! See the iterator-adaptor section in the [`altrep_data`](crate::altrep_data)
7//! module docs for how to expose
8//! these adaptors to R (wrap in a `#[derive(Altrep*)]` + `#[altrep(manual)]`
9//! struct).
10
11use std::cell::RefCell;
12use std::sync::OnceLock;
13
14use crate::altrep_data::{AltIntegerData, AltRealData, AltrepLen, fill_region};
15
16// region: WindowedIterState
17
18/// Core state for windowed iterator-backed ALTREP vectors.
19///
20/// Like [`super::IterState`], but only keeps a sliding window of elements in memory.
21/// Access within the window is O(1), and reading past it slides the window
22/// forward. The iterator cannot be rewound, so reading an element that has left
23/// the window is an error. Use [`super::IterState`] when R may read elements
24/// again or out of order.
25///
26/// This is useful for large iterators where only a small region is accessed
27/// at a time (e.g., streaming data processed in order).
28///
29/// # Type Parameters
30///
31/// - `I`: The iterator type
32/// - `T`: The element type produced by the iterator
33pub struct WindowedIterState<I, T> {
34    len: usize,
35    iter: RefCell<Option<I>>,
36    consumed: RefCell<usize>,
37    window: RefCell<Vec<T>>,
38    window_start: RefCell<usize>,
39    window_size: usize,
40    materialized: OnceLock<Vec<T>>,
41}
42
43impl<I, T> WindowedIterState<I, T>
44where
45    I: Iterator<Item = T>,
46    T: Copy,
47{
48    /// Create a new windowed iterator state.
49    pub fn new(iter: I, len: usize, window_size: usize) -> Self {
50        let window_size = window_size.max(1);
51        Self {
52            len,
53            iter: RefCell::new(Some(iter)),
54            consumed: RefCell::new(0),
55            window: RefCell::new(Vec::with_capacity(window_size)),
56            window_start: RefCell::new(0),
57            window_size,
58            materialized: OnceLock::new(),
59        }
60    }
61
62    /// Get element at index `i`.
63    pub fn get_element(&self, i: usize) -> Option<T> {
64        if i >= self.len {
65            return None;
66        }
67
68        // Check materialized first
69        if let Some(vec) = self.materialized.get() {
70            return vec.get(i).copied();
71        }
72
73        let window_start = *self.window_start.borrow();
74        let window = self.window.borrow();
75
76        // Check if in current window
77        if i >= window_start && i < window_start + window.len() {
78            return Some(window[i - window_start]);
79        }
80        drop(window);
81
82        // Check if we can advance to reach this index
83        let consumed = *self.consumed.borrow();
84        if i >= consumed {
85            // Forward access — advance iterator to fill window containing i
86            self.advance_to(i);
87            let window = self.window.borrow();
88            let window_start = *self.window_start.borrow();
89            if i >= window_start && i < window_start + window.len() {
90                return Some(window[i - window_start]);
91            }
92            return None; // iterator exhausted
93        }
94
95        // Backward access: the element was discarded when the window moved on.
96        {
    ::core::panicking::panic_fmt(format_args!("windowed iterator ALTREP: element {0} left the window, which now starts at element {1}; a windowed vector can only be read forward",
            i + 1, *self.window_start.borrow() + 1));
};panic!(
97            "windowed iterator ALTREP: element {} left the window, which now starts at element {}; \
98             a windowed vector can only be read forward",
99            i + 1,
100            *self.window_start.borrow() + 1
101        );
102    }
103
104    /// Advance the iterator to fill a window containing index `i`.
105    fn advance_to(&self, i: usize) {
106        let mut iter_opt = self.iter.borrow_mut();
107        let iter = match iter_opt.as_mut() {
108            Some(it) => it,
109            None => return,
110        };
111
112        let mut consumed = self.consumed.borrow_mut();
113        let mut window = self.window.borrow_mut();
114        let mut window_start = self.window_start.borrow_mut();
115
116        // Skip elements before the target window
117        let target_window_start = if i >= self.window_size {
118            i - self.window_size + 1
119        } else {
120            0
121        };
122
123        // Skip elements we need to discard
124        while *consumed < target_window_start {
125            if iter.next().is_some() {
126                *consumed += 1;
127            } else {
128                // The skipped elements are gone too; keep the window adjacent to `consumed`.
129                iter_opt.take();
130                window.clear();
131                *window_start = *consumed;
132                return;
133            }
134        }
135
136        // Fill the window
137        window.clear();
138        *window_start = *consumed;
139
140        while window.len() < self.window_size && *consumed < self.len {
141            if let Some(elem) = iter.next() {
142                window.push(elem);
143                *consumed += 1;
144            } else {
145                // Never poll an exhausted iterator again: it need not be fused.
146                iter_opt.take();
147                break;
148            }
149            // Stop once we've passed index i
150            if *consumed > i + 1 && window.len() >= self.window_size {
151                break;
152            }
153        }
154    }
155
156    /// Materialize all elements.
157    ///
158    /// # Panics
159    ///
160    /// If elements have already left the window, or if the iterator ends before
161    /// the declared length.
162    pub fn materialize_all(&self) -> &[T] {
163        if let Some(vec) = self.materialized.get() {
164            return vec;
165        }
166
167        let mut iter_opt = self.iter.borrow_mut();
168        let mut window = self.window.borrow_mut();
169        let mut consumed = self.consumed.borrow_mut();
170        if !(*consumed == window.len()) {
    {
        ::core::panicking::panic_fmt(format_args!("windowed iterator ALTREP: the first {0} elements left the window and cannot be read again",
                *consumed - window.len()));
    }
};assert!(
171            *consumed == window.len(),
172            "windowed iterator ALTREP: the first {} elements left the window and cannot be read again",
173            *consumed - window.len()
174        );
175
176        // Nothing was discarded, so the window holds every element read so far.
177        // Grow it in place: on early exhaustion those elements stay readable.
178        if let Some(iter) = iter_opt.take() {
179            window.extend(iter.take(self.len - *consumed));
180            *consumed = window.len();
181        }
182        if !(window.len() == self.len) {
    {
        ::core::panicking::panic_fmt(format_args!("ALTREP iterator ended before element {0} of its declared length {1}",
                window.len() + 1, self.len));
    }
};assert!(
183            window.len() == self.len,
184            "ALTREP iterator ended before element {} of its declared length {}",
185            window.len() + 1,
186            self.len
187        );
188
189        let result = std::mem::take(&mut *window);
190        drop(consumed);
191        drop(window);
192        drop(iter_opt);
193        self.materialized.get_or_init(|| result)
194    }
195
196    /// Get materialized slice if available.
197    pub fn as_materialized(&self) -> Option<&[T]> {
198        self.materialized.get().map(|v| v.as_slice())
199    }
200
201    /// Get the length.
202    pub fn len(&self) -> usize {
203        self.len
204    }
205
206    /// Check if empty.
207    pub fn is_empty(&self) -> bool {
208        self.len == 0
209    }
210}
211
212impl<I, T> WindowedIterState<I, T>
213where
214    I: ExactSizeIterator<Item = T>,
215    T: Copy,
216{
217    /// Create from an ExactSizeIterator.
218    pub fn from_exact_size(iter: I, window_size: usize) -> Self {
219        let len = iter.len();
220        Self::new(iter, len, window_size)
221    }
222}
223// endregion
224
225// region: Windowed Iterator data adaptors
226
227/// Windowed iterator-backed integer vector data adaptor.
228///
229/// Like [`super::IterIntData`], but only keeps a sliding window of elements in memory.
230/// Forward reads are O(1); reading an element that has left the window is an
231/// error (see [`WindowedIterState`]).
232pub struct WindowedIterIntData<I: Iterator<Item = i32>> {
233    state: WindowedIterState<I, i32>,
234}
235
236impl<I: Iterator<Item = i32>> WindowedIterIntData<I> {
237    /// Create from an iterator with explicit length and window size.
238    pub fn from_iter(iter: I, len: usize, window_size: usize) -> Self {
239        Self {
240            state: WindowedIterState::new(iter, len, window_size),
241        }
242    }
243}
244
245impl<I: ExactSizeIterator<Item = i32>> WindowedIterIntData<I> {
246    /// Create from an ExactSizeIterator with window size (length auto-detected).
247    pub fn from_exact_iter(iter: I, window_size: usize) -> Self {
248        Self {
249            state: WindowedIterState::from_exact_size(iter, window_size),
250        }
251    }
252}
253
254impl<I: Iterator<Item = i32>> AltrepLen for WindowedIterIntData<I> {
255    fn altrep_len(&self) -> usize {
256        self.state.len()
257    }
258}
259
260impl<I: Iterator<Item = i32>> AltIntegerData for WindowedIterIntData<I> {
261    fn elt(&self, i: usize) -> i32 {
262        self.state
263            .get_element(i)
264            .unwrap_or(crate::altrep_traits::NA_INTEGER)
265    }
266
267    fn as_slice(&self) -> Option<&[i32]> {
268        self.state.as_materialized()
269    }
270
271    fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
272        fill_region(start, len, self.altrep_len(), buf, |idx| self.elt(idx))
273    }
274}
275
276/// Windowed iterator-backed real (f64) vector data adaptor.
277///
278/// Like [`super::IterRealData`], but only keeps a sliding window of elements in memory.
279/// Forward reads are O(1); reading an element that has left the window is an
280/// error (see [`WindowedIterState`]).
281pub struct WindowedIterRealData<I: Iterator<Item = f64>> {
282    state: WindowedIterState<I, f64>,
283}
284
285impl<I: Iterator<Item = f64>> WindowedIterRealData<I> {
286    /// Create from an iterator with explicit length and window size.
287    pub fn from_iter(iter: I, len: usize, window_size: usize) -> Self {
288        Self {
289            state: WindowedIterState::new(iter, len, window_size),
290        }
291    }
292}
293
294impl<I: ExactSizeIterator<Item = f64>> WindowedIterRealData<I> {
295    /// Create from an ExactSizeIterator with window size (length auto-detected).
296    pub fn from_exact_iter(iter: I, window_size: usize) -> Self {
297        Self {
298            state: WindowedIterState::from_exact_size(iter, window_size),
299        }
300    }
301}
302
303impl<I: Iterator<Item = f64>> AltrepLen for WindowedIterRealData<I> {
304    fn altrep_len(&self) -> usize {
305        self.state.len()
306    }
307}
308
309impl<I: Iterator<Item = f64>> AltRealData for WindowedIterRealData<I> {
310    fn elt(&self, i: usize) -> f64 {
311        self.state
312            .get_element(i)
313            .unwrap_or(crate::altrep_traits::NA_REAL)
314    }
315
316    fn as_slice(&self) -> Option<&[f64]> {
317        self.state.as_materialized()
318    }
319
320    fn get_region(&self, start: usize, len: usize, buf: &mut [f64]) -> usize {
321        fill_region(start, len, self.altrep_len(), buf, |idx| self.elt(idx))
322    }
323}
324
325// endregion
326
327#[cfg(test)]
328mod tests {
329    use super::WindowedIterState;
330    use std::panic::{AssertUnwindSafe, catch_unwind};
331
332    #[test]
333    fn a_read_behind_the_window_is_an_error() {
334        let state = WindowedIterState::new(0..10, 10, 3);
335        assert_eq!(state.get_element(0), Some(0));
336        assert_eq!(state.get_element(5), Some(5));
337        assert!(catch_unwind(AssertUnwindSafe(|| state.get_element(1))).is_err());
338        assert_eq!(state.get_element(4), Some(4));
339        assert_eq!(state.get_element(6), Some(6));
340    }
341
342    #[test]
343    fn materialization_before_the_window_moves_reads_exactly_len() {
344        let source = || (0..4).chain(std::iter::once_with(|| panic!("read beyond the length")));
345        let state = WindowedIterState::new(source(), 4, 2);
346        assert_eq!(state.materialize_all(), [0, 1, 2, 3]);
347
348        let state = WindowedIterState::new(source(), 4, 2);
349        assert_eq!(state.get_element(1), Some(1));
350        assert_eq!(state.materialize_all(), [0, 1, 2, 3]);
351        assert_eq!(state.get_element(3), Some(3));
352    }
353
354    #[test]
355    fn materialization_after_the_window_moves_is_an_error() {
356        let state = WindowedIterState::new(0..6, 6, 2);
357        assert_eq!(state.get_element(4), Some(4));
358        assert!(catch_unwind(AssertUnwindSafe(|| state.materialize_all())).is_err());
359        assert!(state.as_materialized().is_none());
360        assert_eq!(state.get_element(4), Some(4));
361        assert_eq!(state.get_element(5), Some(5));
362    }
363
364    #[test]
365    fn early_exhaustion_never_publishes_a_partial_buffer() {
366        let state = WindowedIterState::new(0..2, 3, 4);
367        for _ in 0..2 {
368            assert!(catch_unwind(AssertUnwindSafe(|| state.materialize_all())).is_err());
369            assert!(state.as_materialized().is_none());
370            assert_eq!(state.get_element(0), Some(0));
371            assert_eq!(state.get_element(1), Some(1));
372            assert_eq!(state.get_element(2), None);
373        }
374    }
375
376    #[test]
377    fn an_exhausted_source_is_never_polled_again() {
378        // Yields 0, then None, then 1 if polled again (not fused).
379        struct Flaky(u8);
380        impl Iterator for Flaky {
381            type Item = i32;
382            fn next(&mut self) -> Option<i32> {
383                self.0 += 1;
384                match self.0 {
385                    1 => Some(0),
386                    3 => Some(1),
387                    _ => None,
388                }
389            }
390        }
391
392        // Exhausted while filling the window.
393        let state = WindowedIterState::new(Flaky(0), 3, 2);
394        assert_eq!(state.get_element(0), Some(0));
395        assert_eq!(state.get_element(1), None);
396        assert_eq!(state.get_element(1), None);
397
398        // Exhausted while skipping to the window.
399        let state = WindowedIterState::new(Flaky(0), 5, 1);
400        assert_eq!(state.get_element(2), None);
401        assert_eq!(state.get_element(2), None);
402        assert!(catch_unwind(AssertUnwindSafe(|| state.get_element(0))).is_err());
403    }
404}