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/// Sequential access within the window is O(1). Access outside the window
22/// materializes the entire vector (falling back to full caching).
23///
24/// This is useful for large iterators where only a small region is accessed
25/// at a time (e.g., streaming data processed in order).
26///
27/// # Type Parameters
28///
29/// - `I`: The iterator type
30/// - `T`: The element type produced by the iterator
31pub struct WindowedIterState<I, T> {
32    len: usize,
33    iter: RefCell<Option<I>>,
34    consumed: RefCell<usize>,
35    window: RefCell<Vec<T>>,
36    window_start: RefCell<usize>,
37    window_size: usize,
38    materialized: OnceLock<Vec<T>>,
39}
40
41impl<I, T> WindowedIterState<I, T>
42where
43    I: Iterator<Item = T>,
44    T: Copy,
45{
46    /// Create a new windowed iterator state.
47    pub fn new(iter: I, len: usize, window_size: usize) -> Self {
48        let window_size = window_size.max(1);
49        Self {
50            len,
51            iter: RefCell::new(Some(iter)),
52            consumed: RefCell::new(0),
53            window: RefCell::new(Vec::with_capacity(window_size)),
54            window_start: RefCell::new(0),
55            window_size,
56            materialized: OnceLock::new(),
57        }
58    }
59
60    /// Get element at index `i`.
61    pub fn get_element(&self, i: usize) -> Option<T> {
62        if i >= self.len {
63            return None;
64        }
65
66        // Check materialized first
67        if let Some(vec) = self.materialized.get() {
68            return vec.get(i).copied();
69        }
70
71        let window_start = *self.window_start.borrow();
72        let window = self.window.borrow();
73
74        // Check if in current window
75        if i >= window_start && i < window_start + window.len() {
76            return Some(window[i - window_start]);
77        }
78        drop(window);
79
80        // Check if we can advance to reach this index
81        let consumed = *self.consumed.borrow();
82        if i >= consumed {
83            // Forward access — advance iterator to fill window containing i
84            self.advance_to(i);
85            let window = self.window.borrow();
86            let window_start = *self.window_start.borrow();
87            if i >= window_start && i < window_start + window.len() {
88                return Some(window[i - window_start]);
89            }
90            return None; // iterator exhausted
91        }
92
93        // Backward access — must materialize
94        self.materialize_all();
95        self.materialized.get().and_then(|v| v.get(i).copied())
96    }
97
98    /// Advance the iterator to fill a window containing index `i`.
99    fn advance_to(&self, i: usize) {
100        let mut iter_opt = self.iter.borrow_mut();
101        let iter = match iter_opt.as_mut() {
102            Some(it) => it,
103            None => return,
104        };
105
106        let mut consumed = self.consumed.borrow_mut();
107        let mut window = self.window.borrow_mut();
108        let mut window_start = self.window_start.borrow_mut();
109
110        // Skip elements before the target window
111        let target_window_start = if i >= self.window_size {
112            i - self.window_size + 1
113        } else {
114            0
115        };
116
117        // Skip elements we need to discard
118        while *consumed < target_window_start {
119            if iter.next().is_some() {
120                *consumed += 1;
121            } else {
122                return;
123            }
124        }
125
126        // Fill the window
127        window.clear();
128        *window_start = *consumed;
129
130        while window.len() < self.window_size && *consumed < self.len {
131            if let Some(elem) = iter.next() {
132                window.push(elem);
133                *consumed += 1;
134            } else {
135                break;
136            }
137            // Stop once we've passed index i
138            if *consumed > i + 1 && window.len() >= self.window_size {
139                break;
140            }
141        }
142    }
143
144    /// Materialize all elements.
145    pub fn materialize_all(&self) -> &[T] {
146        if let Some(vec) = self.materialized.get() {
147            return vec;
148        }
149
150        // We can only materialize elements from consumed onward
151        // For backward access, we'd need to restart the iterator
152        // Since iterators are consumed, materialize what we can
153        let mut iter_opt = self.iter.borrow_mut();
154        let window = self.window.borrow();
155        let window_start = *self.window_start.borrow();
156
157        let mut result = Vec::with_capacity(self.len);
158
159        // Copy window elements at their correct positions
160        // Start fresh approach: collect remaining from iterator
161        if let Some(iter) = iter_opt.take() {
162            // We have: window contents at window_start..window_start+window.len()
163            // And: unconsumed elements from consumed onward
164            // Elements before window_start are lost (consumed and discarded)
165
166            // Fill lost positions with default
167            for _ in 0..window_start {
168                // These elements were consumed — can't recover
169                result.push(window.first().copied().unwrap_or_else(|| {
170                    // This shouldn't happen with valid usage
171                    unsafe { std::mem::zeroed() }
172                }));
173            }
174
175            // Copy window
176            result.extend_from_slice(&window);
177
178            // Consume rest
179            for elem in iter {
180                if result.len() >= self.len {
181                    break;
182                }
183                result.push(elem);
184            }
185        }
186
187        drop(window);
188        drop(iter_opt);
189
190        if result.len() < self.len {
191            eprintln!(
192                "[miniextendr warning] windowed iterator ALTREP: could only recover {}/{} elements on materialization",
193                result.len(),
194                self.len
195            );
196        }
197
198        self.materialized.get_or_init(|| result)
199    }
200
201    /// Get materialized slice if available.
202    pub fn as_materialized(&self) -> Option<&[T]> {
203        self.materialized.get().map(|v| v.as_slice())
204    }
205
206    /// Get the length.
207    pub fn len(&self) -> usize {
208        self.len
209    }
210
211    /// Check if empty.
212    pub fn is_empty(&self) -> bool {
213        self.len == 0
214    }
215}
216
217impl<I, T> WindowedIterState<I, T>
218where
219    I: ExactSizeIterator<Item = T>,
220    T: Copy,
221{
222    /// Create from an ExactSizeIterator.
223    pub fn from_exact_size(iter: I, window_size: usize) -> Self {
224        let len = iter.len();
225        Self::new(iter, len, window_size)
226    }
227}
228// endregion
229
230// region: Windowed Iterator data adaptors
231
232/// Windowed iterator-backed integer vector data adaptor.
233///
234/// Like [`super::IterIntData`], but only keeps a sliding window of elements in memory.
235/// Sequential forward access within the window is O(1). Access outside the
236/// window triggers full materialization.
237pub struct WindowedIterIntData<I: Iterator<Item = i32>> {
238    state: WindowedIterState<I, i32>,
239}
240
241impl<I: Iterator<Item = i32>> WindowedIterIntData<I> {
242    /// Create from an iterator with explicit length and window size.
243    pub fn from_iter(iter: I, len: usize, window_size: usize) -> Self {
244        Self {
245            state: WindowedIterState::new(iter, len, window_size),
246        }
247    }
248}
249
250impl<I: ExactSizeIterator<Item = i32>> WindowedIterIntData<I> {
251    /// Create from an ExactSizeIterator with window size (length auto-detected).
252    pub fn from_exact_iter(iter: I, window_size: usize) -> Self {
253        Self {
254            state: WindowedIterState::from_exact_size(iter, window_size),
255        }
256    }
257}
258
259impl<I: Iterator<Item = i32>> AltrepLen for WindowedIterIntData<I> {
260    fn len(&self) -> usize {
261        self.state.len()
262    }
263}
264
265impl<I: Iterator<Item = i32>> AltIntegerData for WindowedIterIntData<I> {
266    fn elt(&self, i: usize) -> i32 {
267        self.state
268            .get_element(i)
269            .unwrap_or(crate::altrep_traits::NA_INTEGER)
270    }
271
272    fn as_slice(&self) -> Option<&[i32]> {
273        self.state.as_materialized()
274    }
275
276    fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
277        fill_region(start, len, self.len(), buf, |idx| self.elt(idx))
278    }
279}
280
281/// Windowed iterator-backed real (f64) vector data adaptor.
282///
283/// Like [`super::IterRealData`], but only keeps a sliding window of elements in memory.
284/// Sequential forward access within the window is O(1). Access outside the
285/// window triggers full materialization.
286pub struct WindowedIterRealData<I: Iterator<Item = f64>> {
287    state: WindowedIterState<I, f64>,
288}
289
290impl<I: Iterator<Item = f64>> WindowedIterRealData<I> {
291    /// Create from an iterator with explicit length and window size.
292    pub fn from_iter(iter: I, len: usize, window_size: usize) -> Self {
293        Self {
294            state: WindowedIterState::new(iter, len, window_size),
295        }
296    }
297}
298
299impl<I: ExactSizeIterator<Item = f64>> WindowedIterRealData<I> {
300    /// Create from an ExactSizeIterator with window size (length auto-detected).
301    pub fn from_exact_iter(iter: I, window_size: usize) -> Self {
302        Self {
303            state: WindowedIterState::from_exact_size(iter, window_size),
304        }
305    }
306}
307
308impl<I: Iterator<Item = f64>> AltrepLen for WindowedIterRealData<I> {
309    fn len(&self) -> usize {
310        self.state.len()
311    }
312}
313
314impl<I: Iterator<Item = f64>> AltRealData for WindowedIterRealData<I> {
315    fn elt(&self, i: usize) -> f64 {
316        self.state.get_element(i).unwrap_or(f64::NAN)
317    }
318
319    fn as_slice(&self) -> Option<&[f64]> {
320        self.state.as_materialized()
321    }
322
323    fn get_region(&self, start: usize, len: usize, buf: &mut [f64]) -> usize {
324        fill_region(start, len, self.len(), buf, |idx| self.elt(idx))
325    }
326}
327
328// endregion