miniextendr_api/altrep_data/iter/
sparse.rs

1//! Sparse iterator-backed ALTREP data adaptors with skipping support.
2//!
3//! Provides `SparseIterState<I, T>` which uses `Iterator::nth()` to skip elements
4//! efficiently, and data-adaptor types for each ALTREP family.
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::collections::BTreeMap;
13
14use crate::altrep_data::{
15    AltComplexData, AltIntegerData, AltLogicalData, AltRawData, AltRealData, AltrepLen, Logical,
16    fill_region,
17};
18
19/// Core state for sparse iterator-backed ALTREP vectors.
20///
21/// Unlike [`super::IterState`], this variant uses `Iterator::nth()` to skip elements
22/// efficiently, only caching the elements that are actually accessed.
23///
24/// # Type Parameters
25///
26/// - `I`: The iterator type
27/// - `T`: The element type produced by the iterator
28///
29/// # Design
30///
31/// - **Sparse:** Only accessed elements are cached (uses `BTreeMap`)
32/// - **Skipping:** Uses `nth()` to skip directly to requested indices
33/// - **Trade-off:** Skipped elements are gone forever (iterator is consumed)
34/// - **Best for:** Large iterators where only a small subset of elements are accessed
35///
36/// # Comparison with `IterState`
37///
38/// | Feature | `IterState` | `SparseIterState` |
39/// |---------|-------------|-------------------|
40/// | Cache storage | Contiguous `Vec<T>` | Sparse `BTreeMap<usize, T>` |
41/// | Access pattern | Prefix (0..=i) cached | Only accessed indices cached |
42/// | Skipped elements | All cached | Gone forever (return NA) |
43/// | Memory for sparse access | O(max_index) | O(num_accessed) |
44/// | `as_slice()` support | Yes (after full materialization) | No (sparse) |
45///
46/// # Example
47///
48/// ```ignore
49/// use miniextendr_api::altrep_data::SparseIterIntData;
50///
51/// // Create from an infinite-ish iterator
52/// let data = SparseIterIntData::from_iter((0..).map(|x| x * 2), 1_000_000);
53///
54/// // Access only element 999_999 - skips directly there
55/// let last = data.elt(999_999);  // Only this element is cached
56///
57/// // Element 0 was skipped and is now inaccessible
58/// let first = data.elt(0);  // Returns NA_INTEGER
59/// ```
60pub struct SparseIterState<I, T> {
61    /// Vector length
62    len: usize,
63    /// Iterator state: (iterator, next index the iterator will produce)
64    iter: RefCell<Option<(I, usize)>>,
65    /// Sparse cache of accessed elements
66    cache: RefCell<BTreeMap<usize, T>>,
67}
68
69impl<I, T> SparseIterState<I, T>
70where
71    I: Iterator<Item = T>,
72{
73    /// Create a new sparse iterator state with an explicit length.
74    ///
75    /// # Arguments
76    ///
77    /// - `iter`: The iterator to wrap
78    /// - `len`: The expected number of elements
79    ///
80    /// The first `None` ends the source permanently, even for a non-fused
81    /// iterator. The source is also dropped once its declared prefix is read;
82    /// an empty prefix never polls it.
83    pub fn new(iter: I, len: usize) -> Self {
84        Self {
85            len,
86            iter: RefCell::new((len != 0).then_some((iter, 0))),
87            cache: RefCell::new(BTreeMap::new()),
88        }
89    }
90
91    /// Get an element, skipping intermediate elements if needed.
92    ///
93    /// Uses `Iterator::nth()` to skip efficiently. Skipped elements are
94    /// consumed from the iterator and cannot be retrieved later.
95    ///
96    /// # Returns
97    ///
98    /// - `Some(T)` if element exists and is accessible
99    /// - `None` if:
100    ///   - Index is out of bounds
101    ///   - Element was already skipped (iterator advanced past it)
102    ///   - Iterator exhausted before reaching the index
103    pub fn get_element(&self, i: usize) -> Option<T>
104    where
105        T: Copy,
106    {
107        // Check bounds
108        if i >= self.len {
109            return None;
110        }
111
112        // Check cache first
113        {
114            let cache = self.cache.borrow();
115            if let Some(&val) = cache.get(&i) {
116                return Some(val);
117            }
118        }
119
120        // Need to get from iterator
121        let mut iter_opt = self.iter.borrow_mut();
122        let (iter, pos) = iter_opt.as_mut()?;
123
124        // Element already passed? It was skipped.
125        if i < *pos {
126            return None;
127        }
128
129        // Skip to element i using nth()
130        let skip_count = i - *pos;
131        let elem = match iter.nth(skip_count) {
132            Some(elem) => elem,
133            None => {
134                // Like the prefix and windowed states, never resume a source
135                // after None: it need not implement FusedIterator.
136                iter_opt.take();
137                return None;
138            }
139        };
140        *pos = i + 1;
141        if *pos == self.len {
142            iter_opt.take();
143        }
144
145        // Cache the element
146        drop(iter_opt);
147        self.cache.borrow_mut().insert(i, elem);
148
149        Some(elem)
150    }
151
152    /// Get the current iterator position (next index to be produced).
153    ///
154    /// Returns `None` if the source ended or the declared prefix is complete.
155    pub fn iterator_position(&self) -> Option<usize> {
156        self.iter.borrow().as_ref().map(|(_, pos)| *pos)
157    }
158
159    /// Check if an element has been cached.
160    pub fn is_cached(&self, i: usize) -> bool {
161        self.cache.borrow().contains_key(&i)
162    }
163
164    /// Get the number of cached elements.
165    pub fn cached_count(&self) -> usize {
166        self.cache.borrow().len()
167    }
168
169    /// Get the current length.
170    pub fn len(&self) -> usize {
171        self.len
172    }
173
174    /// Check if the vector is empty.
175    pub fn is_empty(&self) -> bool {
176        self.len == 0
177    }
178}
179
180impl<I, T> SparseIterState<I, T>
181where
182    I: ExactSizeIterator<Item = T>,
183{
184    /// Create a new sparse iterator state from an `ExactSizeIterator`.
185    pub fn from_exact_size(iter: I) -> Self {
186        let len = iter.len();
187        Self::new(iter, len)
188    }
189}
190
191/// Sparse iterator-backed integer vector data adaptor.
192///
193/// Uses `Iterator::nth()` to skip directly to requested indices.
194/// Only accessed elements are cached; skipped elements return `NA_INTEGER`.
195///
196/// # Example
197///
198/// ```ignore
199/// use miniextendr_api::altrep_data::SparseIterIntData;
200///
201/// // Access only specific elements from a large range
202/// let data = SparseIterIntData::from_iter(0..1_000_000, 1_000_000);
203/// let elem = data.elt(500_000);  // Skips 0..499_999
204/// ```
205pub struct SparseIterIntData<I: Iterator<Item = i32>> {
206    state: SparseIterState<I, i32>,
207}
208
209impl<I: Iterator<Item = i32>> SparseIterIntData<I> {
210    /// Create from an iterator with explicit length.
211    pub fn from_iter(iter: I, len: usize) -> Self {
212        Self {
213            state: SparseIterState::new(iter, len),
214        }
215    }
216}
217
218impl<I: ExactSizeIterator<Item = i32>> SparseIterIntData<I> {
219    /// Create from an ExactSizeIterator (length auto-detected).
220    pub fn from_exact_iter(iter: I) -> Self {
221        Self {
222            state: SparseIterState::from_exact_size(iter),
223        }
224    }
225}
226
227impl<I: Iterator<Item = i32>> AltrepLen for SparseIterIntData<I> {
228    fn altrep_len(&self) -> usize {
229        self.state.len()
230    }
231}
232
233impl<I: Iterator<Item = i32>> AltIntegerData for SparseIterIntData<I> {
234    fn elt(&self, i: usize) -> i32 {
235        self.state
236            .get_element(i)
237            .unwrap_or(crate::altrep_traits::NA_INTEGER)
238    }
239
240    fn as_slice(&self) -> Option<&[i32]> {
241        // Sparse storage cannot provide contiguous slice
242        None
243    }
244
245    fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
246        fill_region(start, len, self.altrep_len(), buf, |idx| self.elt(idx))
247    }
248}
249
250/// Sparse iterator-backed real (f64) vector data adaptor.
251///
252/// Uses `Iterator::nth()` to skip directly to requested indices.
253/// Only accessed elements are cached; skipped elements return `NA_real_`.
254pub struct SparseIterRealData<I: Iterator<Item = f64>> {
255    state: SparseIterState<I, f64>,
256}
257
258impl<I: Iterator<Item = f64>> SparseIterRealData<I> {
259    /// Create from an iterator with explicit length.
260    pub fn from_iter(iter: I, len: usize) -> Self {
261        Self {
262            state: SparseIterState::new(iter, len),
263        }
264    }
265}
266
267impl<I: ExactSizeIterator<Item = f64>> SparseIterRealData<I> {
268    /// Create from an ExactSizeIterator (length auto-detected).
269    pub fn from_exact_iter(iter: I) -> Self {
270        Self {
271            state: SparseIterState::from_exact_size(iter),
272        }
273    }
274}
275
276impl<I: Iterator<Item = f64>> AltrepLen for SparseIterRealData<I> {
277    fn altrep_len(&self) -> usize {
278        self.state.len()
279    }
280}
281
282impl<I: Iterator<Item = f64>> AltRealData for SparseIterRealData<I> {
283    fn elt(&self, i: usize) -> f64 {
284        self.state
285            .get_element(i)
286            .unwrap_or(crate::altrep_traits::NA_REAL)
287    }
288
289    fn as_slice(&self) -> Option<&[f64]> {
290        None
291    }
292
293    fn get_region(&self, start: usize, len: usize, buf: &mut [f64]) -> usize {
294        fill_region(start, len, self.altrep_len(), buf, |idx| self.elt(idx))
295    }
296}
297
298/// Sparse iterator-backed logical vector data adaptor.
299pub struct SparseIterLogicalData<I: Iterator<Item = bool>> {
300    state: SparseIterState<I, bool>,
301}
302
303impl<I: Iterator<Item = bool>> SparseIterLogicalData<I> {
304    /// Create from an iterator with explicit length.
305    pub fn from_iter(iter: I, len: usize) -> Self {
306        Self {
307            state: SparseIterState::new(iter, len),
308        }
309    }
310}
311
312impl<I: ExactSizeIterator<Item = bool>> SparseIterLogicalData<I> {
313    /// Create from an ExactSizeIterator (length auto-detected).
314    pub fn from_exact_iter(iter: I) -> Self {
315        Self {
316            state: SparseIterState::from_exact_size(iter),
317        }
318    }
319}
320
321impl<I: Iterator<Item = bool>> AltrepLen for SparseIterLogicalData<I> {
322    fn altrep_len(&self) -> usize {
323        self.state.len()
324    }
325}
326
327impl<I: Iterator<Item = bool>> AltLogicalData for SparseIterLogicalData<I> {
328    fn elt(&self, i: usize) -> Logical {
329        self.state
330            .get_element(i)
331            .map(Logical::from_bool)
332            .unwrap_or(Logical::Na)
333    }
334
335    fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
336        fill_region(start, len, self.altrep_len(), buf, |idx| {
337            self.elt(idx).to_r_int()
338        })
339    }
340}
341
342/// Sparse iterator-backed raw (u8) vector data adaptor.
343pub struct SparseIterRawData<I: Iterator<Item = u8>> {
344    state: SparseIterState<I, u8>,
345}
346
347impl<I: Iterator<Item = u8>> SparseIterRawData<I> {
348    /// Create from an iterator with explicit length.
349    pub fn from_iter(iter: I, len: usize) -> Self {
350        Self {
351            state: SparseIterState::new(iter, len),
352        }
353    }
354}
355
356impl<I: ExactSizeIterator<Item = u8>> SparseIterRawData<I> {
357    /// Create from an ExactSizeIterator (length auto-detected).
358    pub fn from_exact_iter(iter: I) -> Self {
359        Self {
360            state: SparseIterState::from_exact_size(iter),
361        }
362    }
363}
364
365impl<I: Iterator<Item = u8>> AltrepLen for SparseIterRawData<I> {
366    fn altrep_len(&self) -> usize {
367        self.state.len()
368    }
369}
370
371impl<I: Iterator<Item = u8>> AltRawData for SparseIterRawData<I> {
372    fn elt(&self, i: usize) -> u8 {
373        self.state.get_element(i).unwrap_or(0)
374    }
375
376    fn as_slice(&self) -> Option<&[u8]> {
377        None
378    }
379
380    fn get_region(&self, start: usize, len: usize, buf: &mut [u8]) -> usize {
381        fill_region(start, len, self.altrep_len(), buf, |idx| self.elt(idx))
382    }
383}
384
385/// Sparse iterator-backed complex number vector data adaptor.
386pub struct SparseIterComplexData<I>
387where
388    I: Iterator<Item = crate::Rcomplex>,
389{
390    state: SparseIterState<I, crate::Rcomplex>,
391}
392
393impl<I> SparseIterComplexData<I>
394where
395    I: Iterator<Item = crate::Rcomplex>,
396{
397    /// Create from an iterator with explicit length.
398    pub fn from_iter(iter: I, len: usize) -> Self {
399        Self {
400            state: SparseIterState::new(iter, len),
401        }
402    }
403}
404
405impl<I> SparseIterComplexData<I>
406where
407    I: ExactSizeIterator<Item = crate::Rcomplex>,
408{
409    /// Create from an ExactSizeIterator (length auto-detected).
410    pub fn from_exact_iter(iter: I) -> Self {
411        Self {
412            state: SparseIterState::from_exact_size(iter),
413        }
414    }
415}
416
417impl<I> AltrepLen for SparseIterComplexData<I>
418where
419    I: Iterator<Item = crate::Rcomplex>,
420{
421    fn altrep_len(&self) -> usize {
422        self.state.len()
423    }
424}
425
426impl<I> AltComplexData for SparseIterComplexData<I>
427where
428    I: Iterator<Item = crate::Rcomplex>,
429{
430    fn elt(&self, i: usize) -> crate::Rcomplex {
431        self.state
432            .get_element(i)
433            .unwrap_or(<crate::Rcomplex as crate::RNativeType>::R_NA)
434    }
435
436    fn as_slice(&self) -> Option<&[crate::Rcomplex]> {
437        None
438    }
439
440    fn get_region(&self, start: usize, len: usize, buf: &mut [crate::Rcomplex]) -> usize {
441        fill_region(start, len, self.altrep_len(), buf, |idx| self.elt(idx))
442    }
443}
444
445#[cfg(test)]
446mod tests {
447    use super::SparseIterState;
448    use std::cell::Cell;
449    use std::rc::Rc;
450
451    #[test]
452    fn exhausted_sparse_sources_never_resume() {
453        // The third poll would resurrect the source if we kept it after None.
454        let polls = Cell::new(0);
455        let source = std::iter::from_fn(|| {
456            let poll = polls.get() + 1;
457            polls.set(poll);
458            match poll {
459                1 => Some(10),
460                3 => Some(20),
461                _ => None,
462            }
463        });
464        let state = SparseIterState::new(source, 4);
465        assert_eq!(state.get_element(0), Some(10));
466        assert_eq!(state.get_element(2), None); // Ends while nth skips index 1.
467        assert_eq!(state.iterator_position(), None);
468        assert_eq!(state.get_element(2), None);
469        assert_eq!(state.get_element(1), None);
470        assert_eq!(state.get_element(0), Some(10)); // Cache survives EOF.
471        assert_eq!(polls.get(), 2);
472    }
473
474    #[test]
475    fn declared_prefix_drops_the_source_without_an_extra_poll() {
476        for len in [0, 3] {
477            // The closure owns the second handle until the state drops it.
478            let owner = Rc::new(());
479            let captured = Rc::clone(&owner);
480            let polls = Cell::new(0);
481            let poll_count = &polls;
482            let source = std::iter::from_fn(move || {
483                let _keep_alive = &captured;
484                let poll = poll_count.get();
485                assert!(poll < len, "read beyond the declared prefix");
486                poll_count.set(poll + 1);
487                Some(poll)
488            });
489            let state = SparseIterState::new(source, len);
490            if len != 0 {
491                assert_eq!(state.iterator_position(), Some(0));
492                assert_eq!(state.get_element(len - 1), Some(len - 1));
493                assert_eq!(state.get_element(0), None); // Skipped, not cached.
494                assert_eq!(state.get_element(len - 1), Some(len - 1));
495            }
496            assert_eq!(state.iterator_position(), None);
497            assert_eq!(state.get_element(len), None);
498            assert_eq!(polls.get(), len);
499            assert_eq!(Rc::strong_count(&owner), 1);
500        }
501    }
502}
miniextendr_api/altrep_data/iter/sparse.rs

miniextendr_api/altrep_data/iter/
sparse.rs
