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}