Demystifying Asynchronous Rust

1.3. Iterators are closures are iterators

And now for the fun part.

Let's recap the first two sections of this chapter, in as many sentences:

• Iterators are state-machines that are implemented using a structure to keep their state and a .next() method that takes a mutable reference to said state (&mut self) in order to move the machine forward.
• Closures are state-machines that are implemented using a structure to hold their state (i.e. their captured environment) and a call_once/call_mut/call method that takes said state by move/mutable reference/shared reference (respectively) in order to drive the machine forward.

If you're thinking to yourself that these two sound similar, that's because they are.
In this section we're going to have us a little fun by digressing into these similarities.

1.3.a. Iterators as closures

Consider our Range iterator from earlier:

#![allow(unused_variables)]
fn main() {
pub struct Range<T> {
    cur: T,
    end: T,
    incr: T,
}

impl<T> Range<T> {
    pub fn new(start: T, end: T, incr: T) -> Self {
        Self {
            cur: start,
            end,
            incr,
        }
    }
}

impl<T> Iterator for Range<T>
where
    T: std::ops::AddAssign + PartialOrd + Clone,
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        match &self.cur {
            v if *v < self.end => {
                let ret = self.cur.clone();
                self.cur += self.incr.clone();
                ret.into()
            }
            _ => None,
        }
    }
}
}

that we used like this:

#![allow(unused_variables)]
fn main() {
let mut j = 1;
for i in Range::new(1usize, 4, 1) {
    assert_eq!(j, i);
    j += 1;
}
}

Could we express Range in terms of a closure instead?
Well of course we can, what's an iterator but a FnMut() -> Option<T> closure, really?

#![allow(unused_variables)]
fn main() {
pub mod range_fn {
    pub fn new<T>(mut start: T, end: T, incr: T) -> impl FnMut() -> Option<T>
    where
        T: std::ops::AddAssign + PartialOrd + Clone,
    {
        move || {
            if start < end {
                let ret = start.clone();
                start += incr.clone();
                return ret.into();
            }
            None
        }
    }
}
}

which we can basically use the same way as an iterator:

#![allow(unused_variables)]
fn main() {
let mut f = range_fn::new(1, 4, 1);
assert_eq!(Some(1), f());
assert_eq!(Some(2), f());
assert_eq!(Some(3), f());
assert_eq!(None, f());
}

But what about combinators, you say?

1.3.b. Closure combinators

Remember our Range iterator could be combined with a Bounds iterator, allowing us to express something like the following:

#![allow(unused_variables)]
fn main() {
let mut it = Bounds::new(Range::new(1usize, 20, 1), 5, 8);
assert_eq!(Some(5), it.next());
assert_eq!(Some(6), it.next());
assert_eq!(Some(7), it.next());
assert_eq!(None, it.next());
}

We can apply the same pattern to closures: by moving ownership of the wrappee inside the state of the wrapper, we can delegate the state-machinery from the wrapper and into the wrappee.

#![allow(unused_variables)]
fn main() {
pub mod bounds_fn {
    pub fn new<T, F>(mut inner: F, min: T, max: T) -> impl FnMut() -> Option<T>
    where
        T: PartialOrd,
        F: FnMut() -> Option<T>,
    {
        move || loop {
            match inner() {
                Some(v) if v >= min && v < max => return v.into(),
                Some(_) => {}
                None => return None,
            }
        }
    }
}
}

Again, using our closure combinator is pretty similar to using our iterator combinator:

#![allow(unused_variables)]
fn main() {
let mut f = bounds_fn::new(range_fn::new(1usize, 20, 1), 5, 8);
assert_eq!(Some(5), f());
assert_eq!(Some(6), f());
assert_eq!(Some(7), f());
assert_eq!(None, f());
}

And, as we'd expect based on everything we've learned so far, what gets generated here is pretty much the same thing, both from a memory representation and execution model standpoints, as what got generated for the equivalent iterator combinator.

$ cargo rustc --lib -- --test -Zprint-type-sizes
# [...]
## Iterator combinator
`Bounds<Range<usize>, usize>`: 40 bytes, alignment: 8 bytes
## Closure combinator
`[closure<f:[closure<start:usize, end:usize, incr:usize>], min:usize, max:usize>]`: 40 bytes, alignment: 8 bytes

What about extensions, though? Those were the true killer feature of iterator combinators!

1.3.c. Closure extensions

Remember how we were able to do this?

#![allow(unused_variables)]
fn main() {
let mut it = Range::new(1usize, 20, 1).bounds(1, 20).bounds(3, 13).bounds(5, 8);
assert_eq!(Some(5), it.next());
assert_eq!(Some(6), it.next());
assert_eq!(Some(7), it.next());
assert_eq!(None, it.next());
}

Well closures are a trait too, aren't they? Surely we can extend it!

#![allow(unused_variables)]
fn main() {
trait BoundsExtFn<'a, T>: FnMut() -> Option<T>
where
    Self: 'a + Sized,
    T: 'a + std::cmp::PartialOrd,
{
    fn bounds(self, min: T, max: T) -> Box<dyn FnMut() -> Option<T> + 'a> {
        Box::new(bounds_fn::new(self, min, max))
    }
}

impl<'a, F, T> BoundsExtFn<'a, T> for F
where
    F: 'a + FnMut() -> Option<T>,
    T: 'a + std::cmp::PartialOrd,
{
}
}

Ta-daaaa!

#![allow(unused_variables)]
fn main() {
let mut f = range_fn::new(1usize, 20, 1).bounds(1, 20).bounds(3, 13).bounds(5, 8);
assert_eq!(Some(5), f());
assert_eq!(Some(6), f());
assert_eq!(Some(7), f());
assert_eq!(None, f());
}

Ok, that's not really "ta-da" worthy, actually. I lied.

While what we've created here does indeed provide similar functional behavior as our hand-crafted combinator defined above, it has a completely different memory representation and execution model (not to mention that the code itself looks way more complex).
And by different, I actually mean worse in every way.
We've just brought heap allocations and pointer indirections upon ourselves. Oh noes.

$ cargo rustc --lib -- --test -Zprint-type-sizes
# [...]
`[closure<f:Box<dyn FnMut() -> Option<usize>>, min:usize, max:usize>]`: 32 bytes, alignment: 8 bytes

All of our issues stem from the use of a Box there (Box<dyn FnMut() -> Option<T> + 'a>), which begs the question: why did we reach for a Box in the first place?

The reason is actually a limitation in Rust's type system, namely the lack of Generic Associated Types, which prevents us from expressing a trait method that returns a impl FnMut() -> Option<T>, i.e. an unconstrained generic type (GATs are in fact a limited form of HKTs which, once again, are a topic for another day).

But wait a minute, why didn't we face this issue back when we implemented BoundsExt for iterators?

#![allow(unused_variables)]
fn main() {
pub trait BoundsExt: Iterator
where
    Self: Sized,
{
    fn bounds<T>(self, min: T, max: T) -> Bounds<Self, T> {
        Bounds::new(self, min, max)
    }
}

impl<I: Iterator> BoundsExt for I {}
}

That right here is the magical part: Bounds<Self, T>.
I.e. we never had the problem before because we were actually capable of referring to a Bounds<Self, T> by its name.

Unfortunately, one of the first thing we've learned about closures is that we cannot name them; they're all different, and they're all anonymous. Thus we have to return a generic type here, and it certainly cannot be constrained by the caller, since they couldn't possibly name the type of a closure that doesn't even yet exist!

Therefore, what we're left to work with is an unconstrained generic type to return, and an unfortunate solution to the problem: boxing the return value into a trait object, which is exactly what we did there.
Of course, in doing so we've also sacrificed any chance at monomorphization and all the good things that can come out of it: inlining, code elimination, recursive optimizations, etc...

1.3.d. Back and forth

We've shown that we can express iterators as closures and vice-versa, thus we should be able to freely turn one into the other and back again on the fly, shouldn't we?

Turning an iterator into a closure is just a matter of implementing FnMut for I: Iterator, where the implementation does nothing but delegate the call to Iterator::next:

#![allow(unused_variables)]
fn main() {
pub fn iter_to_closure<I: Iterator>(inner: I) -> impl FnMut() -> Option<I::Item> {
    // We cannot implement Fn* traits directly on `I: Iterator` because of
    // coherence.
    struct Iter<I>(I);

    impl<I> FnOnce<()> for Iter<I>
    where
        I: Iterator,
    {
        type Output = Option<I::Item>;

        extern "rust-call" fn call_once(mut self, _args: ()) -> Self::Output {
            Iterator::next(&mut self.0)
        }
    }
    impl<I> FnMut<()> for Iter<I>
    where
        I: Iterator,
    {
        extern "rust-call" fn call_mut(&mut self, _args: ()) -> Self::Output {
            Iterator::next(&mut self.0)
        }
    }

    Iter(inner)
}
}

Going the other way is even more straightforward:

#![allow(unused_variables)]
fn main() {
pub fn closure_to_iter<T, F: FnMut() -> Option<T>>(inner: F) -> impl Iterator<Item = T> {
    struct Iter<F>(F);

    impl<F, T> Iterator for Iter<F>
    where
        F: FnMut() -> Option<T>,
    {
        type Item = T;

        fn next(&mut self) -> Option<Self::Item> {
            self.0()
        }
    }

    Iter(inner)
}
}

And in fact, it seems so natural to want to express an iterator as a closure that libcore provides the exact same code via core::iter::from_fn.

And, voila! From iterators to closures and back again.
No magic tricks, no hidden allocs, no nothing.

#![allow(unused_variables)]
fn main() {
let mut it = Range::new(1usize, 20, 1).into_iter().bounds(5, 14);
assert_eq!(Some(5), it.next());
assert_eq!(Some(6), it.next());
assert_eq!(Some(7), it.next());

let mut f = iter_to_closure(it);
assert_eq!(Some(8), f());
assert_eq!(Some(9), f());
assert_eq!(Some(10), f());

let mut it = closure_to_iter(f);
assert_eq!(Some(11), it.next());
assert_eq!(Some(12), it.next());
assert_eq!(Some(13), it.next());

assert_eq!(None, it.next());
}