Lazy path execution (#150)

jaq-interpret/src/filter.rs

@@ -277,12 +277,18 @@ impl<'a> FilterT<'a> for Ref<'a> {
             Ast::Ite(if_, then_, else_) => w(if_).pipe(cv, move |cv, v| {
                 w(if v.as_bool() { then_ } else { else_ }).run(cv)
             }),
-            Ast::Path(f, path) => then(path.eval(|p| w(p).run(cv.clone())), |path| {
-                let outs = w(f).run(cv).map(move |i| path.collect(i?));
-                Box::new(
-                    outs.flat_map(|vals| then(vals, |vals| Box::new(vals.into_iter().map(Ok)))),
-                )
-            }),
+            Ast::Path(f, path) => {
+                let path = path.map_ref(|i| {
+                    let cv = cv.clone();
+                    crate::into_iter::collect_if_once(move || w(i).run(cv))
+                });
+                flat_map_with(w(f).run(cv), path, |y, path| {
+                    then(y, |y| {
+                        flat_map_with(path.explode(), y, |path, y| then(path, |path| path.run(y)))
+                    })
+                })
+            }
+
             Ast::Update(path, f) => w(path).update(
                 (cv.0.clone(), cv.1),
                 Box::new(move |v| w(f).run((cv.0.clone(), v))),
@@ -366,14 +372,17 @@ impl<'a> FilterT<'a> for Ref<'a> {
             Ast::Fold(..) => todo!(),
 
             Ast::Id => f(cv.1),
-            Ast::Path(l, path) => w(l).update(
-                (cv.0.clone(), cv.1),
-                Box::new(move |v| {
-                    then(path.eval(|i| w(i).run((cv.0.clone(), v.clone()))), |path| {
-                        path.update(v, &f)
-                    })
-                }),
-            ),
+            Ast::Path(l, path) => {
+                let path = path.map_ref(|i| {
+                    let cv = cv.clone();
+                    crate::into_iter::collect_if_once(move || w(i).run(cv))
+                });
+                let f = move |v| {
+                    let mut paths = path.clone().explode();
+                    box_once(paths.try_fold(v, |acc, path| path?.update(acc, &f)))
+                };
+                w(l).update(cv, Box::new(f))
+            }
             Ast::Pipe(l, false, r) => w(l).update(
                 (cv.0.clone(), cv.1),
                 Box::new(move |v| w(r).update((cv.0.clone(), v), f.clone())),

jaq-interpret/src/into_iter.rs

@@ -0,0 +1,60 @@
+//! Functions and types for `IntoIterator` and `FnOnce() -> Iterator`.
+
+#[derive(Clone)]
+pub struct Delay<F>(F);
+
+impl<I: Iterator, F: FnOnce() -> I> IntoIterator for Delay<F> {
+    type Item = I::Item;
+    type IntoIter = I;
+    fn into_iter(self) -> Self::IntoIter {
+        self.0()
+    }
+}
+
+#[derive(Clone)]
+pub enum Either<L, R> {
+    L(L),
+    R(R),
+}
+
+pub struct EitherIter<L, R>(Either<L, R>);
+
+impl<L: Iterator, R: Iterator<Item = L::Item>> Iterator for EitherIter<L, R> {
+    type Item = L::Item;
+    fn next(&mut self) -> Option<Self::Item> {
+        match &mut self.0 {
+            Either::L(l) => l.next(),
+            Either::R(r) => r.next(),
+        }
+    }
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        match &self.0 {
+            Either::L(l) => l.size_hint(),
+            Either::R(r) => r.size_hint(),
+        }
+    }
+}
+
+impl<L: IntoIterator, R: IntoIterator<Item = L::Item>> IntoIterator for Either<L, R> {
+    type Item = L::Item;
+    type IntoIter = EitherIter<L::IntoIter, R::IntoIter>;
+    fn into_iter(self) -> Self::IntoIter {
+        EitherIter(match self {
+            Self::L(l) => Either::L(l.into_iter()),
+            Self::R(r) => Either::R(r.into_iter()),
+        })
+    }
+}
+
+pub fn collect_if_once<I: Iterator, F: FnOnce() -> I + Clone>(
+    f: F,
+) -> Either<core::iter::Once<I::Item>, Delay<F>> {
+    let mut iter = f.clone()();
+    if iter.size_hint().1 == Some(1) {
+        if let Some(x) = iter.next() {
+            assert!(iter.next().is_none());
+            return Either::L(core::iter::once(x));
+        }
+    }
+    Either::R(Delay(f))
+}

jaq-interpret/src/lib.rs

@@ -50,6 +50,7 @@ mod box_iter;
 pub mod error;
 mod filter;
 mod hir;
+mod into_iter;
 mod lir;
 mod mir;
 mod path;