harn_vm/vm/
iter.rs

1//! Lazy iterator protocol for the Harn VM.
2//!
3//! `VmIter` is the backing enum for `VmValue::Iter`. It's a single-pass, fused
4//! iterator; once `next` returns `None` the variant is replaced with
5//! `Exhausted`. Step (a) only introduces source variants (Vec, Dict, Chars,
6//! Gen, Chan, Exhausted) and wires them into the for-loop driver. Combinator
7//! variants (`Map`, `Filter`, `Take`, ...) and sink builtins land in later
8//! steps per the plan.
9
10use std::cell::RefCell;
11use std::collections::{BTreeMap, VecDeque};
12use std::rc::Rc;
13
14use crate::value::{VmChannelHandle, VmError, VmGenerator, VmValue};
15
16/// Backing enum for `VmValue::Iter`. See module docs.
17#[derive(Debug)]
18pub enum VmIter {
19    /// Step through a lazy integer range without materializing.
20    /// `next` is the value to emit on the next call; `stop` is the
21    /// first value that terminates the iteration (one past the end).
22    Range { next: i64, stop: i64 },
23    /// Snapshot over a shared list / set backing store.
24    Vec { items: Rc<Vec<VmValue>>, idx: usize },
25    /// Snapshot over a dict; yields one-key `{key, value}` dicts for now.
26    /// Step (b) swaps these for `VmValue::Pair` when the Pair variant lands.
27    Dict {
28        entries: Rc<BTreeMap<String, VmValue>>,
29        keys: Vec<String>,
30        idx: usize,
31    },
32    /// Unicode scalar iteration over a string.
33    Chars { s: Rc<str>, byte_idx: usize },
34    /// Drains a generator's yield channel.
35    Gen { gen: VmGenerator },
36    /// Reads from a channel handle.
37    Chan { handle: VmChannelHandle },
38    /// Maps each item through a closure.
39    Map {
40        inner: Rc<RefCell<VmIter>>,
41        f: VmValue,
42    },
43    /// Keeps only items for which the predicate is truthy.
44    Filter {
45        inner: Rc<RefCell<VmIter>>,
46        p: VmValue,
47    },
48    /// Maps each item to an iterable and flattens one level.
49    FlatMap {
50        inner: Rc<RefCell<VmIter>>,
51        f: VmValue,
52        cur: Option<Rc<RefCell<VmIter>>>,
53    },
54    /// Yields up to `remaining` items from `inner`, then becomes Exhausted.
55    Take {
56        inner: Rc<RefCell<VmIter>>,
57        remaining: usize,
58    },
59    /// Skips the first `remaining` items from `inner` on the first call, then
60    /// forwards. `remaining == 0` is the sentinel for "already primed".
61    Skip {
62        inner: Rc<RefCell<VmIter>>,
63        remaining: usize,
64    },
65    /// Yields items from `inner` while the predicate is truthy; after the
66    /// first falsy predicate or inner exhaustion, becomes Exhausted.
67    TakeWhile {
68        inner: Rc<RefCell<VmIter>>,
69        p: VmValue,
70        done: bool,
71    },
72    /// Discards items while the predicate is truthy; after the first falsy
73    /// item, forwards that item and all subsequent items from `inner`.
74    SkipWhile {
75        inner: Rc<RefCell<VmIter>>,
76        p: VmValue,
77        primed: bool,
78    },
79    /// Advances two inner iters in lockstep; yields `Pair(a, b)` until either
80    /// side is exhausted.
81    Zip {
82        a: Rc<RefCell<VmIter>>,
83        b: Rc<RefCell<VmIter>>,
84    },
85    /// Yields `Pair(i, item)` starting at `i = 0`.
86    Enumerate { inner: Rc<RefCell<VmIter>>, i: i64 },
87    /// Concatenates two iters: drains `a` first, then `b`.
88    Chain {
89        a: Rc<RefCell<VmIter>>,
90        b: Rc<RefCell<VmIter>>,
91        on_a: bool,
92    },
93    /// Yields `VmValue::List` batches of up to `n` items from `inner`.
94    /// The final batch may be shorter; empty input yields no batches.
95    Chunks {
96        inner: Rc<RefCell<VmIter>>,
97        n: usize,
98    },
99    /// Yields sliding windows of exactly `n` items from `inner` as `VmValue::List`.
100    /// If the input has fewer than `n` items total, no windows are yielded.
101    Windows {
102        inner: Rc<RefCell<VmIter>>,
103        n: usize,
104        buf: VecDeque<VmValue>,
105    },
106    /// Terminal state: `next` always returns `None`.
107    Exhausted,
108}
109
110impl VmIter {
111    /// Produce the next value, or `None` when exhausted.
112    ///
113    /// Combinator variants (`Map`, `Filter`, `FlatMap`) invoke user-provided
114    /// closures through the `vm` parameter.
115    pub fn next<'a>(
116        &'a mut self,
117        vm: &'a mut crate::vm::Vm,
118    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Option<VmValue>, VmError>> + 'a>>
119    {
120        Box::pin(async move { self.next_impl(vm).await })
121    }
122
123    async fn next_impl(&mut self, vm: &mut crate::vm::Vm) -> Result<Option<VmValue>, VmError> {
124        match self {
125            VmIter::Exhausted => Ok(None),
126            VmIter::Range { next, stop } => {
127                if *next < *stop {
128                    let v = *next;
129                    *next += 1;
130                    Ok(Some(VmValue::Int(v)))
131                } else {
132                    *self = VmIter::Exhausted;
133                    Ok(None)
134                }
135            }
136            VmIter::Vec { items, idx } => {
137                if *idx < items.len() {
138                    let v = items[*idx].clone();
139                    *idx += 1;
140                    Ok(Some(v))
141                } else {
142                    *self = VmIter::Exhausted;
143                    Ok(None)
144                }
145            }
146            VmIter::Dict { entries, keys, idx } => {
147                if *idx < keys.len() {
148                    let k = &keys[*idx];
149                    let v = entries.get(k).cloned().unwrap_or(VmValue::Nil);
150                    *idx += 1;
151                    Ok(Some(VmValue::Pair(Rc::new((
152                        VmValue::String(Rc::from(k.as_str())),
153                        v,
154                    )))))
155                } else {
156                    *self = VmIter::Exhausted;
157                    Ok(None)
158                }
159            }
160            VmIter::Chars { s, byte_idx } => {
161                if *byte_idx >= s.len() {
162                    *self = VmIter::Exhausted;
163                    return Ok(None);
164                }
165                let rest = &s[*byte_idx..];
166                if let Some(c) = rest.chars().next() {
167                    *byte_idx += c.len_utf8();
168                    Ok(Some(VmValue::String(Rc::from(c.to_string().as_str()))))
169                } else {
170                    *self = VmIter::Exhausted;
171                    Ok(None)
172                }
173            }
174            VmIter::Gen { gen } => {
175                if gen.done.get() {
176                    *self = VmIter::Exhausted;
177                    return Ok(None);
178                }
179                let rx = gen.receiver.clone();
180                let mut guard = rx.lock().await;
181                match guard.recv().await {
182                    Some(v) => Ok(Some(v)),
183                    None => {
184                        gen.done.set(true);
185                        drop(guard);
186                        *self = VmIter::Exhausted;
187                        Ok(None)
188                    }
189                }
190            }
191            VmIter::Map { inner, f } => {
192                let f = f.clone();
193                let item = next_handle(inner, vm).await?;
194                match item {
195                    None => {
196                        *self = VmIter::Exhausted;
197                        Ok(None)
198                    }
199                    Some(v) => {
200                        let out = vm.call_callable_value(&f, &[v]).await?;
201                        Ok(Some(out))
202                    }
203                }
204            }
205            VmIter::Filter { inner, p } => {
206                let p = p.clone();
207                loop {
208                    let item = next_handle(inner, vm).await?;
209                    match item {
210                        None => {
211                            *self = VmIter::Exhausted;
212                            return Ok(None);
213                        }
214                        Some(v) => {
215                            let keep = vm.call_callable_value(&p, &[v.clone()]).await?;
216                            if keep.is_truthy() {
217                                return Ok(Some(v));
218                            }
219                        }
220                    }
221                }
222            }
223            VmIter::FlatMap { inner, f, cur } => {
224                let f = f.clone();
225                loop {
226                    if let Some(cur_iter) = cur.clone() {
227                        let item = next_handle(&cur_iter, vm).await?;
228                        if let Some(v) = item {
229                            return Ok(Some(v));
230                        }
231                        *cur = None;
232                    }
233                    let item = next_handle(inner, vm).await?;
234                    match item {
235                        None => {
236                            *self = VmIter::Exhausted;
237                            return Ok(None);
238                        }
239                        Some(v) => {
240                            let result = vm.call_callable_value(&f, &[v]).await?;
241                            let lifted = iter_from_value(result)?;
242                            if let VmValue::Iter(h) = lifted {
243                                *cur = Some(h);
244                            } else {
245                                return Err(VmError::TypeError(
246                                    "flat_map: expected iterable result".to_string(),
247                                ));
248                            }
249                        }
250                    }
251                }
252            }
253            VmIter::Take { inner, remaining } => {
254                if *remaining == 0 {
255                    *self = VmIter::Exhausted;
256                    return Ok(None);
257                }
258                let item = next_handle(inner, vm).await?;
259                match item {
260                    None => {
261                        *self = VmIter::Exhausted;
262                        Ok(None)
263                    }
264                    Some(v) => {
265                        *remaining -= 1;
266                        if *remaining == 0 {
267                            *self = VmIter::Exhausted;
268                        }
269                        Ok(Some(v))
270                    }
271                }
272            }
273            VmIter::Skip { inner, remaining } => {
274                while *remaining > 0 {
275                    let item = next_handle(inner, vm).await?;
276                    match item {
277                        None => {
278                            *self = VmIter::Exhausted;
279                            return Ok(None);
280                        }
281                        Some(_) => {
282                            *remaining -= 1;
283                        }
284                    }
285                }
286                let item = next_handle(inner, vm).await?;
287                match item {
288                    None => {
289                        *self = VmIter::Exhausted;
290                        Ok(None)
291                    }
292                    Some(v) => Ok(Some(v)),
293                }
294            }
295            VmIter::TakeWhile { inner, p, done } => {
296                if *done {
297                    return Ok(None);
298                }
299                let p = p.clone();
300                let item = next_handle(inner, vm).await?;
301                match item {
302                    None => {
303                        *self = VmIter::Exhausted;
304                        Ok(None)
305                    }
306                    Some(v) => {
307                        let keep = vm.call_callable_value(&p, &[v.clone()]).await?;
308                        if keep.is_truthy() {
309                            Ok(Some(v))
310                        } else {
311                            *self = VmIter::Exhausted;
312                            Ok(None)
313                        }
314                    }
315                }
316            }
317            VmIter::SkipWhile { inner, p, primed } => {
318                if *primed {
319                    let item = next_handle(inner, vm).await?;
320                    return match item {
321                        None => {
322                            *self = VmIter::Exhausted;
323                            Ok(None)
324                        }
325                        Some(v) => Ok(Some(v)),
326                    };
327                }
328                let p = p.clone();
329                loop {
330                    let item = next_handle(inner, vm).await?;
331                    match item {
332                        None => {
333                            *self = VmIter::Exhausted;
334                            return Ok(None);
335                        }
336                        Some(v) => {
337                            let drop_it = vm.call_callable_value(&p, &[v.clone()]).await?;
338                            if !drop_it.is_truthy() {
339                                *primed = true;
340                                return Ok(Some(v));
341                            }
342                        }
343                    }
344                }
345            }
346            VmIter::Zip { a, b } => {
347                let ia = next_handle(a, vm).await?;
348                let x = match ia {
349                    None => {
350                        *self = VmIter::Exhausted;
351                        return Ok(None);
352                    }
353                    Some(v) => v,
354                };
355                let ib = next_handle(b, vm).await?;
356                let y = match ib {
357                    None => {
358                        *self = VmIter::Exhausted;
359                        return Ok(None);
360                    }
361                    Some(v) => v,
362                };
363                Ok(Some(VmValue::Pair(Rc::new((x, y)))))
364            }
365            VmIter::Enumerate { inner, i } => {
366                let item = next_handle(inner, vm).await?;
367                match item {
368                    None => {
369                        *self = VmIter::Exhausted;
370                        Ok(None)
371                    }
372                    Some(v) => {
373                        let idx = *i;
374                        *i += 1;
375                        Ok(Some(VmValue::Pair(Rc::new((VmValue::Int(idx), v)))))
376                    }
377                }
378            }
379            VmIter::Chain { a, b, on_a } => {
380                if *on_a {
381                    let item = next_handle(a, vm).await?;
382                    if let Some(v) = item {
383                        return Ok(Some(v));
384                    }
385                    *on_a = false;
386                }
387                let item = next_handle(b, vm).await?;
388                match item {
389                    None => {
390                        *self = VmIter::Exhausted;
391                        Ok(None)
392                    }
393                    Some(v) => Ok(Some(v)),
394                }
395            }
396            VmIter::Chunks { inner, n } => {
397                let n = *n;
398                let mut batch: Vec<VmValue> = Vec::with_capacity(n);
399                for _ in 0..n {
400                    let item = next_handle(inner, vm).await?;
401                    match item {
402                        Some(v) => batch.push(v),
403                        None => break,
404                    }
405                }
406                if batch.is_empty() {
407                    *self = VmIter::Exhausted;
408                    Ok(None)
409                } else {
410                    Ok(Some(VmValue::List(Rc::new(batch))))
411                }
412            }
413            VmIter::Windows { inner, n, buf } => {
414                let n = *n;
415                if buf.is_empty() {
416                    while buf.len() < n {
417                        let item = next_handle(inner, vm).await?;
418                        match item {
419                            Some(v) => buf.push_back(v),
420                            None => {
421                                *self = VmIter::Exhausted;
422                                return Ok(None);
423                            }
424                        }
425                    }
426                } else {
427                    let item = next_handle(inner, vm).await?;
428                    match item {
429                        Some(v) => {
430                            buf.pop_front();
431                            buf.push_back(v);
432                        }
433                        None => {
434                            *self = VmIter::Exhausted;
435                            return Ok(None);
436                        }
437                    }
438                }
439                let snapshot: Vec<VmValue> = buf.iter().cloned().collect();
440                Ok(Some(VmValue::List(Rc::new(snapshot))))
441            }
442            VmIter::Chan { handle } => {
443                let is_closed = handle.closed.load(std::sync::atomic::Ordering::Relaxed);
444                let rx = handle.receiver.clone();
445                let mut guard = rx.lock().await;
446                let item = if is_closed {
447                    guard.try_recv().ok()
448                } else {
449                    guard.recv().await
450                };
451                match item {
452                    Some(v) => Ok(Some(v)),
453                    None => {
454                        drop(guard);
455                        *self = VmIter::Exhausted;
456                        Ok(None)
457                    }
458                }
459            }
460        }
461    }
462}
463
464/// Advance a handle without holding a `RefCell` borrow across the await.
465///
466/// Swaps the iter state out into a local owned value (replacing it with
467/// `Exhausted`), runs `next` on the owned state, then swaps it back. This
468/// avoids `clippy::await_holding_refcell_ref` while preserving single-pass
469/// semantics: a nested `next` call on the same handle during the await would
470/// see `Exhausted` (the iter protocol doesn't permit re-entrant stepping of
471/// the same handle anyway).
472pub async fn next_handle(
473    handle: &Rc<RefCell<VmIter>>,
474    vm: &mut crate::vm::Vm,
475) -> Result<Option<VmValue>, VmError> {
476    let mut state = std::mem::replace(&mut *handle.borrow_mut(), VmIter::Exhausted);
477    let result = state.next(vm).await;
478    // Restore state unless the inner call replaced it with Exhausted.
479    *handle.borrow_mut() = state;
480    result
481}
482
483/// Fully consume an iter handle into a Vec of values.
484pub async fn drain(
485    handle: &Rc<RefCell<VmIter>>,
486    vm: &mut crate::vm::Vm,
487) -> Result<Vec<VmValue>, VmError> {
488    let mut out = Vec::new();
489    loop {
490        let v = next_handle(handle, vm).await?;
491        match v {
492            Some(v) => out.push(v),
493            None => break,
494        }
495    }
496    Ok(out)
497}
498
499/// Convenience: wrap a source value into a `VmValue::Iter`. Used by the
500/// `iter()` builtin and by combinator/sink implementations in later steps.
501pub fn iter_from_value(v: VmValue) -> Result<VmValue, VmError> {
502    let inner = match v {
503        VmValue::Iter(h) => return Ok(VmValue::Iter(h)),
504        VmValue::Range(r) => {
505            let stop = if r.inclusive {
506                r.end.saturating_add(1)
507            } else {
508                r.end
509            };
510            VmIter::Range {
511                next: r.start,
512                stop,
513            }
514        }
515        VmValue::List(items) => VmIter::Vec { items, idx: 0 },
516        VmValue::Set(items) => VmIter::Vec { items, idx: 0 },
517        VmValue::Dict(entries) => {
518            let keys: Vec<String> = entries.keys().cloned().collect();
519            VmIter::Dict {
520                entries,
521                keys,
522                idx: 0,
523            }
524        }
525        VmValue::String(s) => VmIter::Chars { s, byte_idx: 0 },
526        VmValue::Generator(gen) => VmIter::Gen { gen },
527        VmValue::Channel(handle) => VmIter::Chan { handle },
528        other => {
529            return Err(VmError::TypeError(format!(
530                "iter: value of type {} is not iterable",
531                other.type_name()
532            )))
533        }
534    };
535    Ok(VmValue::Iter(Rc::new(RefCell::new(inner))))
536}
harn_vm/vm/iter.rs

harn_vm/vm/
iter.rs
