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