Merge pull request #226 from 01mf02/flat-map-then-strict

Make `flat_map_then` strict.

jaq-core/src/box_iter.rs

@@ -21,6 +21,38 @@ pub fn then<'a, T, U: 'a, E: 'a>(
     x.map_or_else(|e| box_once(Err(e)), f)
 }
 
+/// Return next element if iterator returns at most one element, else `None`.
+///
+/// This is one of the most important functions for performance in jaq.
+/// It enables optimisations for the case when a filter yields exactly one output,
+/// which is very common in typical jq programs.
+///
+/// For example, for the filter `f | g`, this function is called on the outputs of `f`.
+/// If `f` yields a single output `y`, then the output of `f | g` is `y` applied to `g`.
+/// This has two beneficial consequences:
+///
+/// 1. We can estimate the number of outputs of `f | g` by estimating the outputs of `g`.
+///    (In general, we cannot do this even when we know the number of outputs of `f`,
+///    because `g` might yield a different number of elements for each output of `f`.)
+/// 2. We do not need to clone the context when passing it to `g`,
+///    because we know that `g` is only called once.
+///
+/// This optimisation applies to many other filters as well.
+///
+/// To see the impact of this function, you can replace its implementation with just `None`.
+/// This preserves correctness, but can result in severely degraded performance.
+fn next_if_one<T>(iter: &mut impl Iterator<Item = T>) -> Option<T> {
+    if iter.size_hint().1 == Some(1) {
+        let ly = iter.next()?;
+        // the Rust documentation states that
+        // "a buggy iterator may yield [..] more than the upper bound of elements",
+        // but so far, it seems that all iterators here are not buggy :)
+        debug_assert!(iter.next().is_none());
+        return Some(ly);
+    }
+    None
+}
+
 /// For every element `y` returned by `l`, return the output of `r(y, x)`.
 ///
 /// In case that `l` returns only a single element, this does not clone `x`.
@@ -29,17 +61,10 @@ pub fn map_with<'a, T: Clone + 'a, U: 'a, V: 'a>(
     x: T,
     r: impl Fn(U, T) -> V + 'a,
 ) -> BoxIter<'a, V> {
-    // this special case is to avoid cloning `x`
-    if l.size_hint().1 == Some(1) {
-        if let Some(ly) = l.next() {
-            // the Rust documentation states that
-            // "a buggy iterator may yield [..] more than the upper bound of elements",
-            // but so far, it seems that all iterators here are not buggy :)
-            assert!(l.next().is_none());
-            return box_once(r(ly, x));
-        }
+    match next_if_one(&mut l) {
+        Some(ly) => box_once(r(ly, x)),
+        None => Box::new(l.map(move |ly| r(ly, x.clone()))),
     }
-    Box::new(l.map(move |ly| r(ly, x.clone())))
 }
 
 /// For every element `y` returned by `l`, return the outputs of `r(y, x)`.
@@ -50,25 +75,21 @@ pub fn flat_map_with<'a, T: Clone + 'a, U: 'a, V: 'a>(
     x: T,
     r: impl Fn(U, T) -> BoxIter<'a, V> + 'a,
 ) -> BoxIter<'a, V> {
-    // this special case is to avoid cloning `x`
-    if l.size_hint().1 == Some(1) {
-        if let Some(ly) = l.next() {
-            // the Rust documentation states that
-            // "a buggy iterator may yield [..] more than the upper bound of elements",
-            // but so far, it seems that all iterators here are not buggy :)
-            assert!(l.next().is_none());
-            return Box::new(r(ly, x));
-        }
+    match next_if_one(&mut l) {
+        Some(ly) => Box::new(r(ly, x)),
+        None => Box::new(l.flat_map(move |ly| r(ly, x.clone()))),
     }
-    Box::new(l.flat_map(move |ly| r(ly, x.clone())))
 }
 
 /// Combination of [`Iterator::flat_map`] and [`then`].
 pub fn flat_map_then<'a, T: 'a, U: 'a, E: 'a>(
-    l: impl Iterator<Item = Result<T, E>> + 'a,
+    mut l: impl Iterator<Item = Result<T, E>> + 'a,
     r: impl Fn(T) -> Results<'a, U, E> + 'a,
 ) -> Results<'a, U, E> {
-    Box::new(l.flat_map(move |y| then(y, |y| r(y))))
+    match next_if_one(&mut l) {
+        Some(ly) => then(ly, r),
+        None => Box::new(l.flat_map(move |y| then(y, |y| r(y)))),
+    }
 }
 
 /// Combination of [`flat_map_with`] and [`then`].

jaq-core/tests/tests.rs

@@ -359,7 +359,7 @@ const ACKERMANN: &str = "def ack($m; $n):
   else ack($m-1; ack($m; $n-1))
   end;";
 
-yields!(ackermann, &(ACKERMANN.to_owned() + "ack(3; 4)"), 125);
+yields!(ackermann, &(ACKERMANN.to_owned() + "ack(3; 3)"), 61);
 
 #[test]
 fn reduce() {

jaq-std/tests/funs.rs

@@ -121,6 +121,8 @@ fn limit() {
     give(json!(null), "[limit(-1; 0, 1)]", json!([]));
 }
 
+yields!(limit_overflow, "[limit(0; def f: f | .; f)]", json!([]));
+
 yields!(
     math_0_argument_scalar_filters,
     "[-2.2, -1.1, 0, 1.1, 2.2 | sin as $s | cos as $c | $s * $s + $c * $c]",