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