miniextendr_api/altrep_data/iter/
windowed.rs1use std::cell::RefCell;
12use std::sync::OnceLock;
13
14use crate::altrep_data::{AltIntegerData, AltRealData, AltrepLen, fill_region};
15
16pub struct WindowedIterState<I, T> {
32 len: usize,
33 iter: RefCell<Option<I>>,
34 consumed: RefCell<usize>,
35 window: RefCell<Vec<T>>,
36 window_start: RefCell<usize>,
37 window_size: usize,
38 materialized: OnceLock<Vec<T>>,
39}
40
41impl<I, T> WindowedIterState<I, T>
42where
43 I: Iterator<Item = T>,
44 T: Copy,
45{
46 pub fn new(iter: I, len: usize, window_size: usize) -> Self {
48 let window_size = window_size.max(1);
49 Self {
50 len,
51 iter: RefCell::new(Some(iter)),
52 consumed: RefCell::new(0),
53 window: RefCell::new(Vec::with_capacity(window_size)),
54 window_start: RefCell::new(0),
55 window_size,
56 materialized: OnceLock::new(),
57 }
58 }
59
60 pub fn get_element(&self, i: usize) -> Option<T> {
62 if i >= self.len {
63 return None;
64 }
65
66 if let Some(vec) = self.materialized.get() {
68 return vec.get(i).copied();
69 }
70
71 let window_start = *self.window_start.borrow();
72 let window = self.window.borrow();
73
74 if i >= window_start && i < window_start + window.len() {
76 return Some(window[i - window_start]);
77 }
78 drop(window);
79
80 let consumed = *self.consumed.borrow();
82 if i >= consumed {
83 self.advance_to(i);
85 let window = self.window.borrow();
86 let window_start = *self.window_start.borrow();
87 if i >= window_start && i < window_start + window.len() {
88 return Some(window[i - window_start]);
89 }
90 return None; }
92
93 self.materialize_all();
95 self.materialized.get().and_then(|v| v.get(i).copied())
96 }
97
98 fn advance_to(&self, i: usize) {
100 let mut iter_opt = self.iter.borrow_mut();
101 let iter = match iter_opt.as_mut() {
102 Some(it) => it,
103 None => return,
104 };
105
106 let mut consumed = self.consumed.borrow_mut();
107 let mut window = self.window.borrow_mut();
108 let mut window_start = self.window_start.borrow_mut();
109
110 let target_window_start = if i >= self.window_size {
112 i - self.window_size + 1
113 } else {
114 0
115 };
116
117 while *consumed < target_window_start {
119 if iter.next().is_some() {
120 *consumed += 1;
121 } else {
122 return;
123 }
124 }
125
126 window.clear();
128 *window_start = *consumed;
129
130 while window.len() < self.window_size && *consumed < self.len {
131 if let Some(elem) = iter.next() {
132 window.push(elem);
133 *consumed += 1;
134 } else {
135 break;
136 }
137 if *consumed > i + 1 && window.len() >= self.window_size {
139 break;
140 }
141 }
142 }
143
144 pub fn materialize_all(&self) -> &[T] {
146 if let Some(vec) = self.materialized.get() {
147 return vec;
148 }
149
150 let mut iter_opt = self.iter.borrow_mut();
154 let window = self.window.borrow();
155 let window_start = *self.window_start.borrow();
156
157 let mut result = Vec::with_capacity(self.len);
158
159 if let Some(iter) = iter_opt.take() {
162 for _ in 0..window_start {
168 result.push(window.first().copied().unwrap_or_else(|| {
170 unsafe { std::mem::zeroed() }
172 }));
173 }
174
175 result.extend_from_slice(&window);
177
178 for elem in iter {
180 if result.len() >= self.len {
181 break;
182 }
183 result.push(elem);
184 }
185 }
186
187 drop(window);
188 drop(iter_opt);
189
190 if result.len() < self.len {
191 eprintln!(
192 "[miniextendr warning] windowed iterator ALTREP: could only recover {}/{} elements on materialization",
193 result.len(),
194 self.len
195 );
196 }
197
198 self.materialized.get_or_init(|| result)
199 }
200
201 pub fn as_materialized(&self) -> Option<&[T]> {
203 self.materialized.get().map(|v| v.as_slice())
204 }
205
206 pub fn len(&self) -> usize {
208 self.len
209 }
210
211 pub fn is_empty(&self) -> bool {
213 self.len == 0
214 }
215}
216
217impl<I, T> WindowedIterState<I, T>
218where
219 I: ExactSizeIterator<Item = T>,
220 T: Copy,
221{
222 pub fn from_exact_size(iter: I, window_size: usize) -> Self {
224 let len = iter.len();
225 Self::new(iter, len, window_size)
226 }
227}
228pub struct WindowedIterIntData<I: Iterator<Item = i32>> {
238 state: WindowedIterState<I, i32>,
239}
240
241impl<I: Iterator<Item = i32>> WindowedIterIntData<I> {
242 pub fn from_iter(iter: I, len: usize, window_size: usize) -> Self {
244 Self {
245 state: WindowedIterState::new(iter, len, window_size),
246 }
247 }
248}
249
250impl<I: ExactSizeIterator<Item = i32>> WindowedIterIntData<I> {
251 pub fn from_exact_iter(iter: I, window_size: usize) -> Self {
253 Self {
254 state: WindowedIterState::from_exact_size(iter, window_size),
255 }
256 }
257}
258
259impl<I: Iterator<Item = i32>> AltrepLen for WindowedIterIntData<I> {
260 fn len(&self) -> usize {
261 self.state.len()
262 }
263}
264
265impl<I: Iterator<Item = i32>> AltIntegerData for WindowedIterIntData<I> {
266 fn elt(&self, i: usize) -> i32 {
267 self.state
268 .get_element(i)
269 .unwrap_or(crate::altrep_traits::NA_INTEGER)
270 }
271
272 fn as_slice(&self) -> Option<&[i32]> {
273 self.state.as_materialized()
274 }
275
276 fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
277 fill_region(start, len, self.len(), buf, |idx| self.elt(idx))
278 }
279}
280
281pub struct WindowedIterRealData<I: Iterator<Item = f64>> {
287 state: WindowedIterState<I, f64>,
288}
289
290impl<I: Iterator<Item = f64>> WindowedIterRealData<I> {
291 pub fn from_iter(iter: I, len: usize, window_size: usize) -> Self {
293 Self {
294 state: WindowedIterState::new(iter, len, window_size),
295 }
296 }
297}
298
299impl<I: ExactSizeIterator<Item = f64>> WindowedIterRealData<I> {
300 pub fn from_exact_iter(iter: I, window_size: usize) -> Self {
302 Self {
303 state: WindowedIterState::from_exact_size(iter, window_size),
304 }
305 }
306}
307
308impl<I: Iterator<Item = f64>> AltrepLen for WindowedIterRealData<I> {
309 fn len(&self) -> usize {
310 self.state.len()
311 }
312}
313
314impl<I: Iterator<Item = f64>> AltRealData for WindowedIterRealData<I> {
315 fn elt(&self, i: usize) -> f64 {
316 self.state.get_element(i).unwrap_or(f64::NAN)
317 }
318
319 fn as_slice(&self) -> Option<&[f64]> {
320 self.state.as_materialized()
321 }
322
323 fn get_region(&self, start: usize, len: usize, buf: &mut [f64]) -> usize {
324 fill_region(start, len, self.len(), buf, |idx| self.elt(idx))
325 }
326}
327
328