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