Skip to main content

jaq_core/
box_iter.rs

1//! Boxed iterators.
2
3use alloc::boxed::Box;
4
5/// A boxed iterator.
6pub type BoxIter<'a, T> = Box<dyn Iterator<Item = T> + 'a>;
7
8/// A boxed iterator over `Result`s.
9pub type Results<'a, T, E> = BoxIter<'a, Result<T, E>>;
10
11/// Return a boxed iterator that yields a single element.
12pub fn box_once<'a, T: 'a>(x: T) -> BoxIter<'a, T> {
13    Box::new(core::iter::once(x))
14}
15
16/// If `x` is an `Err`, return it as iterator, else apply `f` to `x` and return its output.
17pub fn then<'a, T, U: 'a, E: 'a>(
18    x: Result<T, E>,
19    f: impl FnOnce(T) -> Results<'a, U, E>,
20) -> Results<'a, U, E> {
21    x.map_or_else(|e| box_once(Err(e)), f)
22}
23
24/// Return next element if iterator returns at most one element, else `None`.
25///
26/// This is one of the most important functions for performance in jaq.
27/// It enables optimisations for the case when a filter yields exactly one output,
28/// which is very common in typical jq programs.
29///
30/// For example, for the filter `f | g`, this function is called on the outputs of `f`.
31/// If `f` yields a single output `y`, then the output of `f | g` is `y` applied to `g`.
32/// This has two beneficial consequences:
33///
34/// 1. We can estimate the number of outputs of `f | g` by estimating the outputs of `g`.
35///    (In general, we cannot do this even when we know the number of outputs of `f`,
36///    because `g` might yield a different number of elements for each output of `f`.)
37/// 2. We do not need to clone the context when passing it to `g`,
38///    because we know that `g` is only called once.
39///
40/// This optimisation applies to many other filters as well.
41///
42/// To see the impact of this function, you can replace its implementation with just `None`.
43/// This preserves correctness, but can result in severely degraded performance.
44fn next_if_one<T>(iter: &mut impl Iterator<Item = T>) -> Option<T> {
45    if iter.size_hint().1 == Some(1) {
46        let ly = iter.next()?;
47        // the Rust documentation states that
48        // "a buggy iterator may yield [..] more than the upper bound of elements",
49        // but so far, it seems that all iterators here are not buggy :)
50        debug_assert!(iter.next().is_none());
51        return Some(ly);
52    }
53    None
54}
55
56/// For every element `y` returned by `l`, return the output of `r(y, x)`.
57///
58/// In case that `l` returns only a single element, this does not clone `x`.
59pub fn map_with<'a, T: Clone + 'a, U: 'a, V: 'a>(
60    mut l: impl Iterator<Item = U> + 'a,
61    x: T,
62    r: impl Fn(U, T) -> V + 'a,
63) -> BoxIter<'a, V> {
64    match next_if_one(&mut l) {
65        Some(ly) => box_once(r(ly, x)),
66        None => Box::new(l.map(move |ly| r(ly, x.clone()))),
67    }
68}
69
70/// For every element `y` returned by `l`, return the outputs of `r(y, x)`.
71///
72/// In case that `l` returns only a single element, this does not clone `x`.
73pub fn flat_map_with<'a, T: Clone + 'a, U: 'a, V: 'a>(
74    mut l: impl Iterator<Item = U> + 'a,
75    x: T,
76    r: impl Fn(U, T) -> BoxIter<'a, V> + 'a,
77) -> BoxIter<'a, V> {
78    match next_if_one(&mut l) {
79        Some(ly) => Box::new(r(ly, x)),
80        None => Box::new(l.flat_map(move |ly| r(ly, x.clone()))),
81    }
82}
83
84/// Combination of [`Iterator::flat_map`] and [`then`].
85pub fn flat_map_then<'a, T: 'a, U: 'a, E: 'a>(
86    mut l: impl Iterator<Item = Result<T, E>> + 'a,
87    r: impl Fn(T) -> Results<'a, U, E> + 'a,
88) -> Results<'a, U, E> {
89    match next_if_one(&mut l) {
90        Some(ly) => then(ly, r),
91        None => Box::new(l.flat_map(move |y| then(y, |y| r(y)))),
92    }
93}
94
95/// Combination of [`flat_map_with`] and [`then`].
96pub fn flat_map_then_with<'a, T: Clone + 'a, U: 'a, V: 'a, E: 'a>(
97    l: impl Iterator<Item = Result<U, E>> + 'a,
98    x: T,
99    r: impl Fn(U, T) -> Results<'a, V, E> + 'a,
100) -> Results<'a, V, E> {
101    flat_map_with(l, x, move |y, x| then(y, |y| r(y, x)))
102}