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.

Answer 2

As of Rust 1.34 stable, you have convenient std::iter::from_fn utility. It is not a coroutine (i.e. you still have to return each time), but at least it saves you from defining another struct.

from_fn accepts a closure FnMut() -> Option<T> and repeatedly calls it to create an Iterator<T>. In pseudo-Python, def from_fn(f): while (val := f()) is not None: yield val.

// -> Box<dyn std::iter::Iterator<Item=u64>> in Rust 2015
fn firstn(n: u64) -> impl std::iter::Iterator<Item = u64> {
    let mut num = 0;
    std::iter::from_fn(move || {
        let result;
        if num < n {
            result = Some(num);
            num += 1
        } else {
            result = None
        }
        result
    })
}

fn main() {
  let sum_of_first_n = firstn(1000000).sum::<u64>();
  println!("sum(0 to 999999): {}", sum_of_first_n);
}

std::iter::successors is also available. It is less general but might be a bit easier to use since you just pass around the seed value explicitly. In pseudo-Python: def successors(seed, f): while seed is not None: yield seed; seed = f(seed).

fn firstn(n: u64) -> impl std::iter::Iterator<Item = u64> {
    std::iter::successors(
        Some(0),
        move |&num| {
            if num + 1 < n {
                Some(num + 1)
            } else {
                None
            }
        },
    )
}

However, Shepmaster's note applies to these utility too. (tldr: often hand-rolled Iterators are more memory efficient)

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.

(Note: returning impl is a Rust 2018 feature. See the Edition Guide for details)

Answer 3

Rust 1.0 does not have generator functions, so you'd have to do it manually with explicit iterators.

First, rewrite your Python example as a class with a next() method, since that is closer to the model you're likely to get in Rust. Then you can rewrite it in Rust with a struct that implements the Iterator trait.

You might also be able to use a function that returns a closure to achieve a similar result, but I don't think it would be possible to have that implement the Iterator trait (since it would require being called to generate a new result).

Answer 4

You can use my stackful Rust generator library which supports stable Rust:

#[macro_use]
extern crate generator;
use generator::{Generator, Gn};

fn firstn(n: usize) -> Generator<'static, (), usize> {
    Gn::new_scoped(move |mut s| {
        let mut num = 0;
        while num < n {
            s.yield_(num);
            num += 1;
        }
        done!();
    })
}

fn main() {
    let sum_of_first_n: usize = firstn(1000000).sum();
    println!("sum ={}", sum_of_first_n);
}

or more simply:

let n = 100000;
let range = Gn::new_scoped(move |mut s| {
    let mut num = 0;
    while num < n {
        s.yield_(num);
        num += 1;
    }
    done!();
});

let sum: usize = range.sum();