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> {
34 len: usize,
35 iter: RefCell<Option<I>>,
36 consumed: RefCell<usize>,
37 window: RefCell<Vec<T>>,
38 window_start: RefCell<usize>,
39 window_size: usize,
40 materialized: OnceLock<Vec<T>>,
41}
42
43impl<I, T> WindowedIterState<I, T>
44where
45 I: Iterator<Item = T>,
46 T: Copy,
47{
48 pub fn new(iter: I, len: usize, window_size: usize) -> Self {
50 let window_size = window_size.max(1);
51 Self {
52 len,
53 iter: RefCell::new(Some(iter)),
54 consumed: RefCell::new(0),
55 window: RefCell::new(Vec::with_capacity(window_size)),
56 window_start: RefCell::new(0),
57 window_size,
58 materialized: OnceLock::new(),
59 }
60 }
61
62 pub fn get_element(&self, i: usize) -> Option<T> {
64 if i >= self.len {
65 return None;
66 }
67
68 if let Some(vec) = self.materialized.get() {
70 return vec.get(i).copied();
71 }
72
73 let window_start = *self.window_start.borrow();
74 let window = self.window.borrow();
75
76 if i >= window_start && i < window_start + window.len() {
78 return Some(window[i - window_start]);
79 }
80 drop(window);
81
82 let consumed = *self.consumed.borrow();
84 if i >= consumed {
85 self.advance_to(i);
87 let window = self.window.borrow();
88 let window_start = *self.window_start.borrow();
89 if i >= window_start && i < window_start + window.len() {
90 return Some(window[i - window_start]);
91 }
92 return None; }
94
95 {
::core::panicking::panic_fmt(format_args!("windowed iterator ALTREP: element {0} left the window, which now starts at element {1}; a windowed vector can only be read forward",
i + 1, *self.window_start.borrow() + 1));
};panic!(
97 "windowed iterator ALTREP: element {} left the window, which now starts at element {}; \
98 a windowed vector can only be read forward",
99 i + 1,
100 *self.window_start.borrow() + 1
101 );
102 }
103
104 fn advance_to(&self, i: usize) {
106 let mut iter_opt = self.iter.borrow_mut();
107 let iter = match iter_opt.as_mut() {
108 Some(it) => it,
109 None => return,
110 };
111
112 let mut consumed = self.consumed.borrow_mut();
113 let mut window = self.window.borrow_mut();
114 let mut window_start = self.window_start.borrow_mut();
115
116 let target_window_start = if i >= self.window_size {
118 i - self.window_size + 1
119 } else {
120 0
121 };
122
123 while *consumed < target_window_start {
125 if iter.next().is_some() {
126 *consumed += 1;
127 } else {
128 iter_opt.take();
130 window.clear();
131 *window_start = *consumed;
132 return;
133 }
134 }
135
136 window.clear();
138 *window_start = *consumed;
139
140 while window.len() < self.window_size && *consumed < self.len {
141 if let Some(elem) = iter.next() {
142 window.push(elem);
143 *consumed += 1;
144 } else {
145 iter_opt.take();
147 break;
148 }
149 if *consumed > i + 1 && window.len() >= self.window_size {
151 break;
152 }
153 }
154 }
155
156 pub fn materialize_all(&self) -> &[T] {
163 if let Some(vec) = self.materialized.get() {
164 return vec;
165 }
166
167 let mut iter_opt = self.iter.borrow_mut();
168 let mut window = self.window.borrow_mut();
169 let mut consumed = self.consumed.borrow_mut();
170 if !(*consumed == window.len()) {
{
::core::panicking::panic_fmt(format_args!("windowed iterator ALTREP: the first {0} elements left the window and cannot be read again",
*consumed - window.len()));
}
};assert!(
171 *consumed == window.len(),
172 "windowed iterator ALTREP: the first {} elements left the window and cannot be read again",
173 *consumed - window.len()
174 );
175
176 if let Some(iter) = iter_opt.take() {
179 window.extend(iter.take(self.len - *consumed));
180 *consumed = window.len();
181 }
182 if !(window.len() == self.len) {
{
::core::panicking::panic_fmt(format_args!("ALTREP iterator ended before element {0} of its declared length {1}",
window.len() + 1, self.len));
}
};assert!(
183 window.len() == self.len,
184 "ALTREP iterator ended before element {} of its declared length {}",
185 window.len() + 1,
186 self.len
187 );
188
189 let result = std::mem::take(&mut *window);
190 drop(consumed);
191 drop(window);
192 drop(iter_opt);
193 self.materialized.get_or_init(|| result)
194 }
195
196 pub fn as_materialized(&self) -> Option<&[T]> {
198 self.materialized.get().map(|v| v.as_slice())
199 }
200
201 pub fn len(&self) -> usize {
203 self.len
204 }
205
206 pub fn is_empty(&self) -> bool {
208 self.len == 0
209 }
210}
211
212impl<I, T> WindowedIterState<I, T>
213where
214 I: ExactSizeIterator<Item = T>,
215 T: Copy,
216{
217 pub fn from_exact_size(iter: I, window_size: usize) -> Self {
219 let len = iter.len();
220 Self::new(iter, len, window_size)
221 }
222}
223pub struct WindowedIterIntData<I: Iterator<Item = i32>> {
233 state: WindowedIterState<I, i32>,
234}
235
236impl<I: Iterator<Item = i32>> WindowedIterIntData<I> {
237 pub fn from_iter(iter: I, len: usize, window_size: usize) -> Self {
239 Self {
240 state: WindowedIterState::new(iter, len, window_size),
241 }
242 }
243}
244
245impl<I: ExactSizeIterator<Item = i32>> WindowedIterIntData<I> {
246 pub fn from_exact_iter(iter: I, window_size: usize) -> Self {
248 Self {
249 state: WindowedIterState::from_exact_size(iter, window_size),
250 }
251 }
252}
253
254impl<I: Iterator<Item = i32>> AltrepLen for WindowedIterIntData<I> {
255 fn altrep_len(&self) -> usize {
256 self.state.len()
257 }
258}
259
260impl<I: Iterator<Item = i32>> AltIntegerData for WindowedIterIntData<I> {
261 fn elt(&self, i: usize) -> i32 {
262 self.state
263 .get_element(i)
264 .unwrap_or(crate::altrep_traits::NA_INTEGER)
265 }
266
267 fn as_slice(&self) -> Option<&[i32]> {
268 self.state.as_materialized()
269 }
270
271 fn get_region(&self, start: usize, len: usize, buf: &mut [i32]) -> usize {
272 fill_region(start, len, self.altrep_len(), buf, |idx| self.elt(idx))
273 }
274}
275
276pub struct WindowedIterRealData<I: Iterator<Item = f64>> {
282 state: WindowedIterState<I, f64>,
283}
284
285impl<I: Iterator<Item = f64>> WindowedIterRealData<I> {
286 pub fn from_iter(iter: I, len: usize, window_size: usize) -> Self {
288 Self {
289 state: WindowedIterState::new(iter, len, window_size),
290 }
291 }
292}
293
294impl<I: ExactSizeIterator<Item = f64>> WindowedIterRealData<I> {
295 pub fn from_exact_iter(iter: I, window_size: usize) -> Self {
297 Self {
298 state: WindowedIterState::from_exact_size(iter, window_size),
299 }
300 }
301}
302
303impl<I: Iterator<Item = f64>> AltrepLen for WindowedIterRealData<I> {
304 fn altrep_len(&self) -> usize {
305 self.state.len()
306 }
307}
308
309impl<I: Iterator<Item = f64>> AltRealData for WindowedIterRealData<I> {
310 fn elt(&self, i: usize) -> f64 {
311 self.state
312 .get_element(i)
313 .unwrap_or(crate::altrep_traits::NA_REAL)
314 }
315
316 fn as_slice(&self) -> Option<&[f64]> {
317 self.state.as_materialized()
318 }
319
320 fn get_region(&self, start: usize, len: usize, buf: &mut [f64]) -> usize {
321 fill_region(start, len, self.altrep_len(), buf, |idx| self.elt(idx))
322 }
323}
324
325#[cfg(test)]
328mod tests {
329 use super::WindowedIterState;
330 use std::panic::{AssertUnwindSafe, catch_unwind};
331
332 #[test]
333 fn a_read_behind_the_window_is_an_error() {
334 let state = WindowedIterState::new(0..10, 10, 3);
335 assert_eq!(state.get_element(0), Some(0));
336 assert_eq!(state.get_element(5), Some(5));
337 assert!(catch_unwind(AssertUnwindSafe(|| state.get_element(1))).is_err());
338 assert_eq!(state.get_element(4), Some(4));
339 assert_eq!(state.get_element(6), Some(6));
340 }
341
342 #[test]
343 fn materialization_before_the_window_moves_reads_exactly_len() {
344 let source = || (0..4).chain(std::iter::once_with(|| panic!("read beyond the length")));
345 let state = WindowedIterState::new(source(), 4, 2);
346 assert_eq!(state.materialize_all(), [0, 1, 2, 3]);
347
348 let state = WindowedIterState::new(source(), 4, 2);
349 assert_eq!(state.get_element(1), Some(1));
350 assert_eq!(state.materialize_all(), [0, 1, 2, 3]);
351 assert_eq!(state.get_element(3), Some(3));
352 }
353
354 #[test]
355 fn materialization_after_the_window_moves_is_an_error() {
356 let state = WindowedIterState::new(0..6, 6, 2);
357 assert_eq!(state.get_element(4), Some(4));
358 assert!(catch_unwind(AssertUnwindSafe(|| state.materialize_all())).is_err());
359 assert!(state.as_materialized().is_none());
360 assert_eq!(state.get_element(4), Some(4));
361 assert_eq!(state.get_element(5), Some(5));
362 }
363
364 #[test]
365 fn early_exhaustion_never_publishes_a_partial_buffer() {
366 let state = WindowedIterState::new(0..2, 3, 4);
367 for _ in 0..2 {
368 assert!(catch_unwind(AssertUnwindSafe(|| state.materialize_all())).is_err());
369 assert!(state.as_materialized().is_none());
370 assert_eq!(state.get_element(0), Some(0));
371 assert_eq!(state.get_element(1), Some(1));
372 assert_eq!(state.get_element(2), None);
373 }
374 }
375
376 #[test]
377 fn an_exhausted_source_is_never_polled_again() {
378 struct Flaky(u8);
380 impl Iterator for Flaky {
381 type Item = i32;
382 fn next(&mut self) -> Option<i32> {
383 self.0 += 1;
384 match self.0 {
385 1 => Some(0),
386 3 => Some(1),
387 _ => None,
388 }
389 }
390 }
391
392 let state = WindowedIterState::new(Flaky(0), 3, 2);
394 assert_eq!(state.get_element(0), Some(0));
395 assert_eq!(state.get_element(1), None);
396 assert_eq!(state.get_element(1), None);
397
398 let state = WindowedIterState::new(Flaky(0), 5, 1);
400 assert_eq!(state.get_element(2), None);
401 assert_eq!(state.get_element(2), None);
402 assert!(catch_unwind(AssertUnwindSafe(|| state.get_element(0))).is_err());
403 }
404}