perf(stdlib): make enumerate return a lazy iterator 🦥

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

Author: timfennis

ndc_stdlib/src/sequence.rs

@@ -5,7 +5,7 @@ use ndc_core::compare::FallibleOrd;
 use ndc_macros::export_module;
 use ndc_vm::VmCallable;
 use ndc_vm::value::{Object, SeqValue, Value};
-use ndc_vm::{CombinationsIter, TakeIter};
+use ndc_vm::{CombinationsIter, EnumerateIter, TakeIter};
 use std::cmp::Ordering;
 
 fn try_sort_by<E>(
@@ -334,16 +334,12 @@ mod inner {
         }
     }
 
-    /// Enumerates the given sequence returning a list of tuples where the first element of the tuple is the index of the element in the input sequence.
-    #[function(return_type = Vec<_>)]
+    /// Returns a lazy iterator yielding `(index, value)` tuples for each element of `seq`.
+    #[function(return_type = Iterator<Value>)]
     pub fn enumerate(seq: SeqValue) -> anyhow::Result<Value> {
-        Ok(Value::list(
-            seq.try_into_iter()
-                .ok_or_else(|| anyhow!("enumerate requires a sequence"))?
-                .enumerate()
-                .map(|(i, v)| Value::tuple(vec![Value::Int(i as i64), v]))
-                .collect(),
-        ))
+        let iter =
+            EnumerateIter::new(seq).ok_or_else(|| anyhow!("enumerate requires a sequence"))?;
+        Ok(Value::iterator(iter.into_shared()))
     }
 
     /// Reduces/folds the given sequence using the given combining function and a custom initial value.

ndc_vm/src/iterator.rs

@@ -574,6 +574,58 @@ impl VmIterator for TakeIter {
     }
 }
 
+/// Pairs each upstream element with its 0-based index, lazily.
+///
+/// `deep_copy` is supported when the source is a `Shared` iterator. It returns
+/// `None` for eager sources (list/deque/map `IntoIter`s) since those lack a
+/// copy mechanism — same constraint as [`TakeIter`].
+pub struct EnumerateIter {
+    source: ValueIter,
+    index: usize,
+}
+
+impl EnumerateIter {
+    /// Returns `None` if `value` is not iterable.
+    pub fn new(value: Value) -> Option<Self> {
+        Some(Self {
+            source: value.try_into_iter()?,
+            index: 0,
+        })
+    }
+
+    pub fn into_shared(self) -> SharedIterator {
+        Rc::new(RefCell::new(self))
+    }
+}
+
+impl VmIterator for EnumerateIter {
+    fn next(&mut self) -> Option<Value> {
+        let v = self.source.next()?;
+        let i = self.index;
+        self.index += 1;
+        Some(Value::tuple(vec![Value::Int(i as i64), v]))
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        match &self.source {
+            ValueIter::Shared(s) => s.borrow().size_hint(),
+            ValueIter::List(it) => it.size_hint(),
+            ValueIter::Deque(it) => it.size_hint(),
+            ValueIter::Map(it) => it.size_hint(),
+        }
+    }
+
+    fn deep_copy(&self) -> Option<SharedIterator> {
+        let ValueIter::Shared(shared) = &self.source else {
+            return None;
+        };
+        Some(Rc::new(RefCell::new(Self {
+            source: ValueIter::Shared(shared.borrow().deep_copy()?),
+            index: self.index,
+        })))
+    }
+}
+
 /// Iterates over string characters, yielding each as a single-char string
 pub struct StringIter {
     string: Rc<RefCell<String>>,

ndc_vm/src/lib.rs

@@ -13,7 +13,9 @@ pub use vm::*;
 
 pub use compiler::CompileError;
 pub use error::VmError;
-pub use iterator::{CombinationsIter, RepeatIter, SharedIterator, TakeIter, VmIterator};
+pub use iterator::{
+    CombinationsIter, EnumerateIter, RepeatIter, SharedIterator, TakeIter, VmIterator,
+};
 pub use value::*;
 
 #[cfg(test)]

tests/functional/programs/603_stdlib_seq/002_enumerate.ndc

@@ -1,6 +1,12 @@
 let my_list = ["Foo", "Bar", "Baz"];
-let out = my_list.enumerate();
+let out = my_list.enumerate().list();
 
 assert_eq(out[0], (0, "Foo"));
 assert_eq(out[1], (1, "Bar"));
 assert_eq(out[2], (2, "Baz"));
+
+// Laziness: enumerating an unbounded range and taking is finite.
+assert_eq((0..).enumerate().take(3).list(), [(0, 0), (1, 1), (2, 2)]);
+
+// Enumeration composes with iterator combinators.
+assert_eq([10, 20, 30].enumerate().map(fn(p) { p[0] + p[1] }), [10, 21, 32]);