64

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?

Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
NathanD
  • 8,061
  • 7
  • 30
  • 26

4 Answers4

62

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.

Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
  • 2
    "Generators do not have any such information." That's kind of the point. One of the primary use cases of generators is to be able to do operations on infinite sequences without having to hold the infinite sequence in memory. This kind of lazy computing has many applications. So is there any way to avoid converting it back to an iterator? – zer0fool Aug 28 '20 at 01:54
  • @zer0fool The conversion to an iterator is not a problem here. It will still work with infinite sequences. The iterator is like a generator without the syntactical sugar of yield. – Karsten Aug 29 '20 at 12:30
  • 1
    @zer0fool standard Rust iterators *already* allow infinite sequences without holding the entire sequence. See [`iter::repeat`](https://doc.rust-lang.org/std/iter/fn.repeat.html) for an example. Generators do not add that ability. Iterators have **additional** information that generators do not that allows optimizations. – Shepmaster Aug 31 '20 at 14:10
  • @zer0fool You can avoid converting to an iterator by using the generator directly, as the body of the iterator adapter shows. You just will not be able to use it where iterators are expected. – Shepmaster Aug 31 '20 at 14:10
  • 1
    It is worth noting that new Rust versions have the unstable [`std::iter::from_generator()`](https://doc.rust-lang.org/stable/std/iter/fn.from_generator.html), so no need to write the iterator adapter yourself. https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=142447451f94ed2435c35da3e73c9026. – Chayim Friedman Jun 22 '23 at 09:25
42

As of Rust 1.34 stable, you have convenient std::iter::from_fn utility. It is not a coroutine (i.e. you have to return and finish execution for each element), 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| {
            let next = num + 1;
            if next < n {
                Some(next)
            } 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 configuration and Announcement or Rust By Example for explanation)

snipsnipsnip
  • 2,268
  • 2
  • 33
  • 34
  • It works but not always. It's not the best alternative to proper generators/yield – Konrad Jun 09 '21 at 23:22
  • Imagine you need to return someting and after returning you need to sleep for 1 second. You can't do that with `std::iter::from_fn` nor `std::iter::successors` – Konrad Jun 09 '21 at 23:23
  • @Konrad: You can sleep before returning, or upon reentry... How is sleeping for one second before returning distinguishable from sleeping one second *immediately* after returning? Are you just saying you don't want to sleep before producing the very first item, only for all subsequent items? – ShadowRanger Jun 21 '23 at 13:31
19

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).

Chris Rudd
  • 709
  • 7
  • 13
Lucian
  • 564
  • 4
  • 7
  • 4
    Depending on the exact sequence, there are several build-in iterators that could be used e.g. [`Unfoldr`](http://static.rust-lang.org/doc/core/iterator.html#struct-unfoldriterator), or [`Counter`](http://static.rust-lang.org/doc/core/iterator.html#struct-counter) plus a [`scan`](http://static.rust-lang.org/doc/core/iterator.html#method-scan) (and/or any of the other combining functions. (Unfortunately there's no documentation for those other than the types yet.) – huon May 08 '13 at 14:13
  • I wonder if it's possible to write a macro that lets you define custom iterator structs but with a bit less boilerplate. – eremzeit Jan 23 '17 at 11:09
  • 1
    This answer is *very* out of date. See [Shepmaster's answer](https://stackoverflow.com/a/30279122/1350209) – Cole Tobin May 27 '20 at 03:25
5

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();
Shepmaster
  • 388,571
  • 95
  • 1,107
  • 1,366
Xudong Huang
  • 141
  • 2
  • 3