Make new module into_iter.

01mf02 · Jan 17, 2024 · 2f5799f · 2f5799f
1 parent d028aee
commit 2f5799f
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 56 deletions.
diff --git a/jaq-interpret/src/filter.rs b/jaq-interpret/src/filter.rs
@@ -225,60 +225,6 @@ impl<'a> FilterT<'a> for &'a Owned {
     }
 }
 
-#[derive(Clone)]
-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)]
-enum Either<L, R> {
-    L(L),
-    R(R),
-}
-
-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()),
-        })
-    }
-}
-
-fn bla<I: Iterator, F: FnOnce() -> I + Clone>(f: F) -> Either<Vec<I::Item>, Delay<F>> {
-    let iter = f.clone()();
-    match iter.size_hint().1 {
-        Some(_) => Either::L(iter.collect::<Vec<_>>()),
-        None => Either::R(Delay(f)),
-    }
-}
-
 impl<'a> FilterT<'a> for Ref<'a> {
     fn run(self, cv: Cv<'a>) -> ValRs<'a> {
         use core::iter::{once, once_with};
@@ -335,7 +281,7 @@ impl<'a> FilterT<'a> for Ref<'a> {
                 use crate::path::{self, Path};
                 let run = |i| {
                     let cv = cv.clone();
-                    bla(move || w(i).run(cv))
+                    crate::into_iter::collect_if_once(move || w(i).run(cv))
                 };
                 let path = path.0.iter();
                 let path = path.map(move |(part, opt)| (part.as_ref().map(run), *opt));
@@ -440,7 +386,7 @@ impl<'a> FilterT<'a> for Ref<'a> {
                 use crate::path::{self, Path};
                 let run = |i| {
                     let cv = cv.clone();
-                    bla(move || w(i).run(cv))
+                    crate::into_iter::collect_if_once(move || w(i).run(cv))
                 };
                 let path = path.0.iter();
                 let path = path.map(move |(part, opt)| (part.as_ref().map(run), *opt));

diff --git a/jaq-interpret/src/into_iter.rs b/jaq-interpret/src/into_iter.rs
@@ -0,0 +1,60 @@
+//! Functions and type to deal with `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))
+}
diff --git a/jaq-interpret/src/lib.rs b/jaq-interpret/src/lib.rs
@@ -51,6 +51,7 @@ pub mod error;
 mod filter;
 mod hir;
 mod lir;
+mod into_iter;
 mod mir;
 mod path;
 mod rc_iter;