Skip to main content

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 generated
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    pub fn new(iter: I, len: usize) -> Self {
80        Self {
81            len,
82            iter: RefCell::new(Some((iter, 0))),
83            cache: RefCell::new(BTreeMap::new()),
84        }
85    }
86
87    /// Get an element, skipping intermediate elements if needed.
88    ///
89    /// Uses `Iterator::nth()` to skip efficiently. Skipped elements are
90    /// consumed from the iterator and cannot be retrieved later.
91    ///
92    /// # Returns
93    ///
94    /// - `Some(T)` if element exists and is accessible
95    /// - `None` if:
96    ///   - Index is out of bounds
97    ///   - Element was already skipped (iterator advanced past it)
98    ///   - Iterator exhausted before reaching the index
99    pub fn get_element(&self, i: usize) -> Option<T>
100    where
101        T: Copy,
102    {
103        // Check bounds
104        if i >= self.len {
105            return None;
106        }
107
108        // Check cache first
109        {
110            let cache = self.cache.borrow();
111            if let Some(&val) = cache.get(&i) {
112                return Some(val);
113            }
114        }
115
116        // Need to get from iterator
117        let mut iter_opt = self.iter.borrow_mut();
118        let (iter, pos) = iter_opt.as_mut()?;
119
120        // Element already passed? It was skipped.
121        if i < *pos {
122            return None;
123        }
124
125        // Skip to element i using nth()
126        let skip_count = i - *pos;
127        let elem = iter.nth(skip_count)?;
128        *pos = i + 1;
129
130        // Cache the element
131        drop(iter_opt);
132        self.cache.borrow_mut().insert(i, elem);
133
134        Some(elem)
135    }
136
137    /// Get the current iterator position (next index to be produced).
138    ///
139    /// Returns `None` if the iterator has been exhausted.
140    pub fn iterator_position(&self) -> Option<usize> {
141        self.iter.borrow().as_ref().map(|(_, pos)| *pos)
142    }
143
144    /// Check if an element has been cached.
145    pub fn is_cached(&self, i: usize) -> bool {
146        self.cache.borrow().contains_key(&i)
147    }
148
149    /// Get the number of cached elements.
150    pub fn cached_count(&self) -> usize {
151        self.cache.borrow().len()
152    }
153
154    /// Get the current length.
155    pub fn len(&self) -> usize {
156        self.len
157    }
158
159    /// Check if the vector is empty.
160    pub fn is_empty(&self) -> bool {
161        self.len == 0
162    }
163}
164
165impl<I, T> SparseIterState<I, T>
166where
167    I: ExactSizeIterator<Item = T>,
168{
169    /// Create a new sparse iterator state from an `ExactSizeIterator`.
170    pub fn from_exact_size(iter: I) -> Self {
171        let len = iter.len();
172        Self::new(iter, len)
173    }
174}
175
176/// Sparse iterator-backed integer vector data adaptor.
177///
178/// Uses `Iterator::nth()` to skip directly to requested indices.
179/// Only accessed elements are cached; skipped elements return `NA_INTEGER`.
180///
181/// # Example
182///
183/// ```ignore
184/// use miniextendr_api::altrep_data::SparseIterIntData;
185///
186/// // Access only specific elements from a large range
187/// let data = SparseIterIntData::from_iter(0..1_000_000, 1_000_000);
188/// let elem = data.elt(500_000);  // Skips 0..499_999
189/// ```
190pub struct SparseIterIntData<I: Iterator<Item = i32>> {
191    state: SparseIterState<I, i32>,
192}
193
194impl<I: Iterator<Item = i32>> SparseIterIntData<I> {
195    /// Create from an iterator with explicit length.
196    pub fn from_iter(iter: I, len: usize) -> Self {
197        Self {
198            state: SparseIterState::new(iter, len),
199        }
200    }
201}
202
203impl<I: ExactSizeIterator<Item = i32>> SparseIterIntData<I> {
204    /// Create from an ExactSizeIterator (length auto-detected).
205    pub fn from_exact_iter(iter: I) -> Self {
206        Self {
207            state: SparseIterState::from_exact_size(iter),
208        }
209    }
210}
211
212impl<I: Iterator<Item = i32>> AltrepLen for SparseIterIntData<I> {
213    fn len(&self) -> usize {
214        self.state.len()
215    }
216}
217
218impl<I: Iterator<Item = i32>> AltIntegerData for SparseIterIntData<I> {
219    fn elt(&self, i: usize) -> i32 {
220        self.state
221            .get_element(i)
222            .unwrap_or(crate::altrep_traits::NA_INTEGER)
223    }
224
225    fn as_slice(&self) -> Option<&[i32]> {
226        // Sparse storage cannot provide contiguous slice
227        None
228    }
229
230    fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
231        fill_region(start, len, self.len(), buf, |idx| self.elt(idx))
232    }
233}
234
235/// Sparse iterator-backed real (f64) vector data adaptor.
236///
237/// Uses `Iterator::nth()` to skip directly to requested indices.
238/// Only accessed elements are cached; skipped elements return `NaN`.
239pub struct SparseIterRealData<I: Iterator<Item = f64>> {
240    state: SparseIterState<I, f64>,
241}
242
243impl<I: Iterator<Item = f64>> SparseIterRealData<I> {
244    /// Create from an iterator with explicit length.
245    pub fn from_iter(iter: I, len: usize) -> Self {
246        Self {
247            state: SparseIterState::new(iter, len),
248        }
249    }
250}
251
252impl<I: ExactSizeIterator<Item = f64>> SparseIterRealData<I> {
253    /// Create from an ExactSizeIterator (length auto-detected).
254    pub fn from_exact_iter(iter: I) -> Self {
255        Self {
256            state: SparseIterState::from_exact_size(iter),
257        }
258    }
259}
260
261impl<I: Iterator<Item = f64>> AltrepLen for SparseIterRealData<I> {
262    fn len(&self) -> usize {
263        self.state.len()
264    }
265}
266
267impl<I: Iterator<Item = f64>> AltRealData for SparseIterRealData<I> {
268    fn elt(&self, i: usize) -> f64 {
269        self.state.get_element(i).unwrap_or(f64::NAN)
270    }
271
272    fn as_slice(&self) -> Option<&[f64]> {
273        None
274    }
275
276    fn get_region(&self, start: usize, len: usize, buf: &mut [f64]) -> usize {
277        fill_region(start, len, self.len(), buf, |idx| self.elt(idx))
278    }
279}
280
281/// Sparse iterator-backed logical vector data adaptor.
282pub struct SparseIterLogicalData<I: Iterator<Item = bool>> {
283    state: SparseIterState<I, bool>,
284}
285
286impl<I: Iterator<Item = bool>> SparseIterLogicalData<I> {
287    /// Create from an iterator with explicit length.
288    pub fn from_iter(iter: I, len: usize) -> Self {
289        Self {
290            state: SparseIterState::new(iter, len),
291        }
292    }
293}
294
295impl<I: ExactSizeIterator<Item = bool>> SparseIterLogicalData<I> {
296    /// Create from an ExactSizeIterator (length auto-detected).
297    pub fn from_exact_iter(iter: I) -> Self {
298        Self {
299            state: SparseIterState::from_exact_size(iter),
300        }
301    }
302}
303
304impl<I: Iterator<Item = bool>> AltrepLen for SparseIterLogicalData<I> {
305    fn len(&self) -> usize {
306        self.state.len()
307    }
308}
309
310impl<I: Iterator<Item = bool>> AltLogicalData for SparseIterLogicalData<I> {
311    fn elt(&self, i: usize) -> Logical {
312        self.state
313            .get_element(i)
314            .map(Logical::from_bool)
315            .unwrap_or(Logical::Na)
316    }
317
318    fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
319        fill_region(start, len, self.len(), buf, |idx| self.elt(idx).to_r_int())
320    }
321}
322
323/// Sparse iterator-backed raw (u8) vector data adaptor.
324pub struct SparseIterRawData<I: Iterator<Item = u8>> {
325    state: SparseIterState<I, u8>,
326}
327
328impl<I: Iterator<Item = u8>> SparseIterRawData<I> {
329    /// Create from an iterator with explicit length.
330    pub fn from_iter(iter: I, len: usize) -> Self {
331        Self {
332            state: SparseIterState::new(iter, len),
333        }
334    }
335}
336
337impl<I: ExactSizeIterator<Item = u8>> SparseIterRawData<I> {
338    /// Create from an ExactSizeIterator (length auto-detected).
339    pub fn from_exact_iter(iter: I) -> Self {
340        Self {
341            state: SparseIterState::from_exact_size(iter),
342        }
343    }
344}
345
346impl<I: Iterator<Item = u8>> AltrepLen for SparseIterRawData<I> {
347    fn len(&self) -> usize {
348        self.state.len()
349    }
350}
351
352impl<I: Iterator<Item = u8>> AltRawData for SparseIterRawData<I> {
353    fn elt(&self, i: usize) -> u8 {
354        self.state.get_element(i).unwrap_or(0)
355    }
356
357    fn as_slice(&self) -> Option<&[u8]> {
358        None
359    }
360
361    fn get_region(&self, start: usize, len: usize, buf: &mut [u8]) -> usize {
362        fill_region(start, len, self.len(), buf, |idx| self.elt(idx))
363    }
364}
365
366/// Sparse iterator-backed complex number vector data adaptor.
367pub struct SparseIterComplexData<I>
368where
369    I: Iterator<Item = crate::Rcomplex>,
370{
371    state: SparseIterState<I, crate::Rcomplex>,
372}
373
374impl<I> SparseIterComplexData<I>
375where
376    I: Iterator<Item = crate::Rcomplex>,
377{
378    /// Create from an iterator with explicit length.
379    pub fn from_iter(iter: I, len: usize) -> Self {
380        Self {
381            state: SparseIterState::new(iter, len),
382        }
383    }
384}
385
386impl<I> SparseIterComplexData<I>
387where
388    I: ExactSizeIterator<Item = crate::Rcomplex>,
389{
390    /// Create from an ExactSizeIterator (length auto-detected).
391    pub fn from_exact_iter(iter: I) -> Self {
392        Self {
393            state: SparseIterState::from_exact_size(iter),
394        }
395    }
396}
397
398impl<I> AltrepLen for SparseIterComplexData<I>
399where
400    I: Iterator<Item = crate::Rcomplex>,
401{
402    fn len(&self) -> usize {
403        self.state.len()
404    }
405}
406
407impl<I> AltComplexData for SparseIterComplexData<I>
408where
409    I: Iterator<Item = crate::Rcomplex>,
410{
411    fn elt(&self, i: usize) -> crate::Rcomplex {
412        self.state.get_element(i).unwrap_or(crate::Rcomplex {
413            r: f64::NAN,
414            i: f64::NAN,
415        })
416    }
417
418    fn as_slice(&self) -> Option<&[crate::Rcomplex]> {
419        None
420    }
421
422    fn get_region(&self, start: usize, len: usize, buf: &mut [crate::Rcomplex]) -> usize {
423        fill_region(start, len, self.len(), buf, |idx| self.elt(idx))
424    }
425}