Lazy sequence generation in Rust

Question

How can I create what other languages call a lazy sequence or a "generator" function?

In Python, I can use yield as in the following example (from Python's docs) to lazily generate a sequence that is iterable in a way that does not use the memory of an intermediary list:

# a generator that yields items instead of returning a list
def firstn(n):
    num = 0
    while num < n:
        yield num
        num += 1

sum_of_first_n = sum(firstn(1000000))

How can I do something similar in Rust?

Answer 1

Rust does have generators, but they are highly experimental and not currently available in stable Rust.

Works in stable Rust 1.0 and above

Range handles your concrete example. You can use it with the syntactical sugar of ..:

fn main() {
    let sum: u64 = (0..1_000_000).sum();
    println!("{}", sum)
}

What if Range didn't exist? We can create an iterator that models it!

struct MyRange {
    start: u64,
    end: u64,
}

impl MyRange {
    fn new(start: u64, end: u64) -> MyRange {
        MyRange {
            start: start,
            end: end,
        }
    }
}

impl Iterator for MyRange {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        if self.start == self.end {
            None
        } else {
            let result = Some(self.start);
            self.start += 1;
            result
        }
    }
}

fn main() {
    let sum: u64 = MyRange::new(0, 1_000_000).sum();
    println!("{}", sum)
}

The guts are the same, but more explicit than the Python version. Notably, Python's generators keep track of the state for you. Rust prefers explicitness, so we have to create our own state and update it manually. The important part is the implementation of the Iterator trait. We specify that the iterator yields values of a specific type (type Item = u64) and then deal with stepping each iteration and how to tell we have reached the end of iteration.

This example is not as powerful as the real Range, which uses generics, but shows an example of how to go about it.

Works in nightly Rust

Nightly Rust does have generators, but they are highly experimental. You need to bring in a few unstable features to create one. However, it looks pretty close to the Python example, with some Rust-specific additions:

// 1.43.0-nightly (2020-02-09 71c7e149e42cb0fc78a8)
#![feature(generators, generator_trait)]

use std::{
    ops::{Generator, GeneratorState},
    pin::Pin,
};

fn firstn(n: u64) -> impl Generator<Yield = u64, Return = ()> {
    move || {
        let mut num = 0;
        while num < n {
            yield num;
            num += 1;
        }
    }
}

Since everything in current Rust operates on iterators, we create an adapter that converts a generator into an iterator in order to play with the broader ecosystem. I'd expect that such an adapter would be present in the standard library eventually:

struct GeneratorIteratorAdapter<G>(Pin<Box<G>>);

impl<G> GeneratorIteratorAdapter<G>
where
    G: Generator<Return = ()>,
{
    fn new(gen: G) -> Self {
        Self(Box::pin(gen))
    }
}

impl<G> Iterator for GeneratorIteratorAdapter<G>
where
    G: Generator<Return = ()>,
{
    type Item = G::Yield;

    fn next(&mut self) -> Option<Self::Item> {
        match self.0.as_mut().resume(()) {
            GeneratorState::Yielded(x) => Some(x),
            GeneratorState::Complete(_) => None,
        }
    }
}

Now we can use it:

fn main() {
    let generator_iterator = GeneratorIteratorAdapter::new(firstn(1_000_000));
    let sum: u64 = generator_iterator.sum();
    println!("{}", sum);
}

What's interesting about this is that it's less powerful than an implementation of Iterator. For example, iterators have the size_hint method, which allows consumers of the iterator to have an idea of how many elements are remaining. This allows optimizations when collecting into a container. Generators do not have any such information.