Skip to main content

miniextendr_api/altrep_data/
stream.rs

1//! Streaming ALTREP data adaptors backed by chunk-cached reader closures.
2//!
3//! These types provide backing data for ALTREP vectors where elements are
4//! loaded on-demand from a reader function in fixed-size chunks. Chunks are
5//! cached for repeated access within the same region.
6//!
7//! Like the iterator adaptors documented in
8//! [`altrep_data`](crate::altrep_data), these implement only
9//! the data-level traits ([`AltrepLen`] + `Alt*Data`); to expose one to R,
10//! wrap it in a concrete `#[derive(Altrep*)]` + `#[altrep(manual)]` struct
11//! that delegates the data-trait methods (see the
12//! [`altrep_data`](crate::altrep_data) module docs for the pattern).
13
14use std::cell::RefCell;
15use std::collections::BTreeMap;
16
17use super::{AltIntegerData, AltRealData, AltrepLen};
18
19// region: StreamingRealData
20
21/// Streaming data adaptor for ALTREP real (f64) vectors.
22///
23/// Elements are loaded on-demand via a reader closure in fixed-size chunks.
24/// Chunks are cached in a `BTreeMap` for repeated access.
25///
26/// # Reader Contract
27///
28/// The reader `F(start, buf) -> count` fills `buf` with elements starting
29/// at index `start` and returns the number of elements actually written.
30/// Elements it does not write read as NA.
31///
32/// # Example
33///
34/// ```ignore
35/// use miniextendr_api::altrep_data::StreamingRealData;
36///
37/// let data = StreamingRealData::new(1000, 64, |start, buf| {
38///     let count = buf.len().min(1000 - start);
39///     for (i, slot) in buf[..count].iter_mut().enumerate() {
40///         *slot = (start + i) as f64 * 0.1;
41///     }
42///     count
43/// });
44/// ```
45pub struct StreamingRealData<F: Fn(usize, &mut [f64]) -> usize> {
46    len: usize,
47    reader: F,
48    cache: RefCell<BTreeMap<usize, Vec<f64>>>,
49    chunk_size: usize,
50}
51
52impl<F: Fn(usize, &mut [f64]) -> usize> StreamingRealData<F> {
53    /// Create a new streaming real data source.
54    ///
55    /// - `len`: total number of elements
56    /// - `chunk_size`: number of elements per cache chunk
57    /// - `reader`: closure that fills a buffer starting at a given index
58    pub fn new(len: usize, chunk_size: usize, reader: F) -> Self {
59        Self {
60            len,
61            reader,
62            cache: RefCell::new(BTreeMap::new()),
63            chunk_size: chunk_size.max(1),
64        }
65    }
66
67    /// Load a chunk into the cache if not already present.
68    fn ensure_chunk(&self, chunk_idx: usize) {
69        let mut cache = self.cache.borrow_mut();
70        if cache.contains_key(&chunk_idx) {
71            return;
72        }
73        let start = chunk_idx * self.chunk_size;
74        let count = self.chunk_size.min(self.len.saturating_sub(start));
75        if count == 0 {
76            return;
77        }
78        let mut buf = ::alloc::vec::from_elem(0.0f64, count)vec![0.0f64; count];
79        let written = (self.reader)(start, &mut buf);
80        buf.truncate(written);
81        cache.insert(chunk_idx, buf);
82    }
83}
84
85impl<F: Fn(usize, &mut [f64]) -> usize> AltrepLen for StreamingRealData<F> {
86    fn altrep_len(&self) -> usize {
87        self.len
88    }
89}
90
91impl<F: Fn(usize, &mut [f64]) -> usize> AltRealData for StreamingRealData<F> {
92    fn elt(&self, i: usize) -> f64 {
93        if i >= self.len {
94            return crate::altrep_traits::NA_REAL;
95        }
96        let chunk_idx = i / self.chunk_size;
97        self.ensure_chunk(chunk_idx);
98        let cache = self.cache.borrow();
99        let offset = i % self.chunk_size;
100        cache
101            .get(&chunk_idx)
102            .and_then(|chunk| chunk.get(offset).copied())
103            .unwrap_or(crate::altrep_traits::NA_REAL)
104    }
105
106    fn get_region(&self, start: usize, len: usize, buf: &mut [f64]) -> usize {
107        let count = len.min(self.len.saturating_sub(start)).min(buf.len());
108        if count == 0 {
109            return 0;
110        }
111        // R's ITERATE_BY_REGION reads all `count` slots whatever we return, so
112        // the elements the reader did not write read as NA, as through `elt`.
113        let written = (self.reader)(start, &mut buf[..count]).min(count);
114        buf[written..count].fill(crate::altrep_traits::NA_REAL);
115        count
116    }
117}
118
119// endregion
120
121// region: StreamingIntData
122
123/// Streaming data adaptor for ALTREP integer (i32) vectors.
124///
125/// Elements are loaded on-demand via a reader closure in fixed-size chunks.
126/// Chunks are cached in a `BTreeMap` for repeated access.
127///
128/// # Reader Contract
129///
130/// The reader `F(start, buf) -> count` fills `buf` with elements starting
131/// at index `start` and returns the number of elements actually written.
132/// Elements it does not write read as NA.
133///
134/// # Example
135///
136/// ```ignore
137/// use miniextendr_api::altrep_data::StreamingIntData;
138///
139/// let data = StreamingIntData::new(1000, 64, |start, buf| {
140///     let count = buf.len().min(1000 - start);
141///     for (i, slot) in buf[..count].iter_mut().enumerate() {
142///         *slot = (start + i) as i32;
143///     }
144///     count
145/// });
146/// ```
147pub struct StreamingIntData<F: Fn(usize, &mut [i32]) -> usize> {
148    len: usize,
149    reader: F,
150    cache: RefCell<BTreeMap<usize, Vec<i32>>>,
151    chunk_size: usize,
152}
153
154impl<F: Fn(usize, &mut [i32]) -> usize> StreamingIntData<F> {
155    /// Create a new streaming integer data source.
156    ///
157    /// - `len`: total number of elements
158    /// - `chunk_size`: number of elements per cache chunk
159    /// - `reader`: closure that fills a buffer starting at a given index
160    pub fn new(len: usize, chunk_size: usize, reader: F) -> Self {
161        Self {
162            len,
163            reader,
164            cache: RefCell::new(BTreeMap::new()),
165            chunk_size: chunk_size.max(1),
166        }
167    }
168
169    /// Load a chunk into the cache if not already present.
170    fn ensure_chunk(&self, chunk_idx: usize) {
171        let mut cache = self.cache.borrow_mut();
172        if cache.contains_key(&chunk_idx) {
173            return;
174        }
175        let start = chunk_idx * self.chunk_size;
176        let count = self.chunk_size.min(self.len.saturating_sub(start));
177        if count == 0 {
178            return;
179        }
180        let mut buf = ::alloc::vec::from_elem(0i32, count)vec![0i32; count];
181        let written = (self.reader)(start, &mut buf);
182        buf.truncate(written);
183        cache.insert(chunk_idx, buf);
184    }
185}
186
187impl<F: Fn(usize, &mut [i32]) -> usize> AltrepLen for StreamingIntData<F> {
188    fn altrep_len(&self) -> usize {
189        self.len
190    }
191}
192
193impl<F: Fn(usize, &mut [i32]) -> usize> AltIntegerData for StreamingIntData<F> {
194    fn elt(&self, i: usize) -> i32 {
195        if i >= self.len {
196            return crate::altrep_traits::NA_INTEGER;
197        }
198        let chunk_idx = i / self.chunk_size;
199        self.ensure_chunk(chunk_idx);
200        let cache = self.cache.borrow();
201        let offset = i % self.chunk_size;
202        cache
203            .get(&chunk_idx)
204            .and_then(|chunk| chunk.get(offset).copied())
205            .unwrap_or(crate::altrep_traits::NA_INTEGER)
206    }
207
208    fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
209        let count = len.min(self.len.saturating_sub(start)).min(buf.len());
210        if count == 0 {
211            return 0;
212        }
213        // R's ITERATE_BY_REGION reads all `count` slots whatever we return, so
214        // the elements the reader did not write read as NA, as through `elt`.
215        let written = (self.reader)(start, &mut buf[..count]).min(count);
216        buf[written..count].fill(crate::altrep_traits::NA_INTEGER);
217        count
218    }
219}
220
221// endregion