Closures and Iterators · about 40 min

Writing Your Own Iterator

Why this matters

The standard iterators cover collections, ranges and text. Your own data structures and generators deserve the same treatment: a tree that can be walked with for, a tokenizer that yields tokens, a sequence generator. Implement one method, next, and your type gains every adapter and consumer from the last lesson. This lesson shows how, including iterators that borrow from data, and the shortcuts for when a full struct is overkill.

Implementing Iterator

An iterator is a struct holding its state, plus next to advance it:

struct Countdown {
    remaining: u32,
}

impl Iterator for Countdown {
    type Item = u32;

    fn next(&mut self) -> Option<u32> {
        if self.remaining == 0 {
            None
        } else {
            self.remaining -= 1;
            Some(self.remaining + 1)
        }
    }
}

fn main() {
    for n in (Countdown { remaining: 3 }) {
        print!("{n} ");
    }
    println!("liftoff");
    let total: u32 = Countdown { remaining: 10 }.filter(|n| n % 2 == 0).sum();
    let v: Vec<u32> = Countdown { remaining: 5 }.skip(1).step_by(2).collect();
    println!("{total} {:?}", v);
}

3 2 1 liftoff
30 [4, 2]

type Item = u32; declares what the iterator yields.
next(&mut self) updates the state and returns Some(item), or None when finished.
Everything else (filter, sum, skip, step_by, collect, use in for) comes from the trait's provided methods.

Once next has returned None, well-behaved iterators keep returning None; adapters like fuse() enforce it for iterators that do not.

Infinite iterators

Nothing requires an iterator to end. Infinite iterators are fine as long as the consumer stops: take, take_while, find:

struct Collatz {
    n: u64,
}

impl Iterator for Collatz {
    type Item = u64;
    fn next(&mut self) -> Option<u64> {
        let current = self.n;
        self.n = if current % 2 == 0 { current / 2 } else { 3 * current + 1 };
        Some(current)
    }
}

fn main() {
    let seq: Vec<u64> = Collatz { n: 6 }.take_while(|&x| x != 1).collect();
    println!("{:?}", seq);
    let steps = Collatz { n: 27 }.position(|x| x == 1).unwrap();
    println!("27 reaches 1 after {steps} steps");
}

[6, 3, 10, 5, 16, 8, 4, 2]
27 reaches 1 after 111 steps

Laziness is what makes this work: values are produced only when asked for.

Iterating over your own collection

To make for x in &my_collection work, implement IntoIterator for a reference to your type. The simplest way is to hand out an iterator you already have, such as the one of an inner Vec:

struct Playlist {
    songs: Vec<String>,
}

impl<'a> IntoIterator for &'a Playlist {
    type Item = &'a String;
    type IntoIter = std::slice::Iter<'a, String>;
    fn into_iter(self) -> Self::IntoIter {
        self.songs.iter()
    }
}

fn main() {
    let p = Playlist { songs: vec![String::from("Intro"), String::from("Outro")] };
    for (i, s) in (&p).into_iter().enumerate() {
        println!("{}: {s}", i + 1);
    }
    for s in &p {
        print!("{} ", s.len());
    }
    println!();
}

1: Intro
2: Outro
5 5

for x in EXPR calls IntoIterator::into_iter(EXPR). Every iterator is trivially IntoIterator, and collections implement it for T, &T and &mut T, which is how the three loop forms work.

Iterators that borrow

An iterator over borrowed data carries a lifetime, like any struct holding references (module 10). Here is an iterator yielding the lines of a text that are not comments, as slices of the original:

struct CodeLines<'a> {
    inner: std::str::Lines<'a>,
}

impl<'a> Iterator for CodeLines<'a> {
    type Item = &'a str;
    fn next(&mut self) -> Option<&'a str> {
        loop {
            let line = self.inner.next()?.trim();
            if !line.is_empty() && !line.starts_with('#') {
                return Some(line);
            }
        }
    }
}

fn code_lines(text: &str) -> CodeLines<'_> {
    CodeLines { inner: text.lines() }
}

fn main() {
    let src = "# config\nport = 80\n\n  # comment\nhost = example\n";
    let v: Vec<&str> = code_lines(src).collect();
    println!("{:?}", v);
}

["port = 80", "host = example"]

Item = &'a str ties each yielded slice to the original text, not to the iterator, so the results stay valid after the iterator is gone. The ? on self.inner.next() ends our iterator when the inner one ends: wrapping and filtering another iterator is a very common shape.

Shortcuts: from_fn and successors

For one-off iterators, closures avoid defining a struct:

fn main() {
    let mut count = 0;
    let evens: Vec<u32> = std::iter::from_fn(|| {
        count += 2;
        if count <= 10 { Some(count) } else { None }
    })
    .collect();
    println!("{:?}", evens);

    let powers: Vec<u64> = std::iter::successors(Some(1u64), |&x| x.checked_mul(10)).take(5).collect();
    println!("{:?}", powers);

    let ones: Vec<i32> = std::iter::repeat(1).take(3).collect();
    println!("{:?}", ones);
}

[2, 4, 6, 8, 10]
[1, 10, 100, 1000, 10000]
[1, 1, 1]

from_fn turns a closure returning Option into an iterator (the closure's captured variables are the state). successors builds each item from the previous one. repeat, once and empty round out the toolbox.

Common mistakes

Forgetting type Item: E0046 "not all trait items implemented, missing: Item".
Returning items that borrow from the iterator itself (&self.buffer): not possible with Iterator, because the signature next(&mut self) -> Option<Self::Item> cannot tie the item to that borrow. Borrow from the underlying data ('a) instead, or yield owned values.
Infinite iterator with an unbounded consumer: Collatz { n: 6 }.collect::<Vec<_>>() never ends. Always bound it with take, take_while or find.
for x in my_struct without IntoIterator: E0277 "Playlist is not an iterator". Implement IntoIterator (or iterate a field).

Recap

Implement Iterator with type Item and fn next(&mut self) -> Option<Item>; state lives in the struct.
All adapters and consumers come for free; rev additionally needs DoubleEndedIterator.
Infinite iterators are fine with bounded consumers.
IntoIterator for &YourType makes for x in &value work.
Borrowing iterators carry a lifetime and yield &'a T tied to the data.
iter::from_fn, successors, repeat create quick iterators from closures.

Going further: implement DoubleEndedIterator for Countdown (add next_back) so that .rev() works, and read the std::iter module documentation.

Exercises

Exercise 1 · write code

A Fibonacci iterator

Implement struct Fib { a: u64, b: u64 } with fn new() -> Fib and impl Iterator for Fib yielding 0, 1, 1, 2, 3, 5, 8, ... forever (it may stop by returning None once the next value would overflow u64; use checked_add). The given main reads n and prints the first n numbers, then the first Fibonacci number above 1000, then the sum of the even ones below 100:

0 1 1 2 3 5 8 13 21 34
first above 1000: 1597
even sum below 100: 44

for input 10.

File: main.rs

Starting point

use std::io;

fn main() {
    let mut s = String::new();
    io::stdin().read_line(&mut s).unwrap();
    let n: usize = s.trim().parse().unwrap();
    let first: Vec<String> = Fib::new().take(n).map(|x| x.to_string()).collect();
    println!("{}", first.join(" "));
    println!("first above 1000: {}", Fib::new().find(|&x| x > 1000).unwrap());
    let even: u64 = Fib::new().take_while(|&x| x < 100).filter(|x| x % 2 == 0).sum();
    println!("even sum below 100: {even}");
}

Example

input

expected output

0 1 1 2 3 5 8 13 21 34
first above 1000: 1597
even sum below 100: 44

Example

input

expected output

0
first above 1000: 1597
even sum below 100: 44

Hint 1

impl Iterator for Fib { type Item = u64; fn next(&mut self) -> Option<u64> { ... } }

Hint 2

Return the current a, then advance: new a is old b, new b is a + b.

Hint 3

The struct may get an extra field (for example done: bool) to stop cleanly when checked_add returns None.

Reference solution

use std::io;

struct Fib {
    a: u64,
    b: u64,
    done: bool,
}

impl Fib {
    fn new() -> Fib {
        Fib { a: 0, b: 1, done: false }
    }
}

impl Iterator for Fib {
    type Item = u64;
    fn next(&mut self) -> Option<u64> {
        if self.done {
            return None;
        }
        let current = self.a;
        match self.a.checked_add(self.b) {
            Some(next) => {
                self.a = self.b;
                self.b = next;
            }
            None => self.done = true,
        }
        Some(current)
    }
}

fn main() {
    let mut s = String::new();
    io::stdin().read_line(&mut s).unwrap();
    let n: usize = s.trim().parse().unwrap();
    let first: Vec<String> = Fib::new().take(n).map(|x| x.to_string()).collect();
    println!("{}", first.join(" "));
    println!("first above 1000: {}", Fib::new().find(|&x| x > 1000).unwrap());
    let even: u64 = Fib::new().take_while(|&x| x < 100).filter(|x| x % 2 == 0).sum();
    println!("even sum below 100: {even}");
}

Exercise 2 · write code

An iterator over borrowed text

Implement struct Tokens<'a> { rest: &'a str } whose iterator yields &'a str tokens: runs of ASCII digits, runs of ASCII letters, and every other non-space character on its own; spaces are skipped. Add fn tokens(s: &str) -> Tokens<'_>. The given main prints each line's tokens with {:?}:

["x1", "=", "42", "+", "y"]
["(", "a", ")", "*", "10"]

for the lines x1 = 42+y and (a)*10. Note: a letter run may continue with digits (x1 is one token), but a digit run stops at a letter.

File: main.rs

Starting point

use std::io;

fn main() {
    for line in io::stdin().lines() {
        let line = line.unwrap();
        let t: Vec<&str> = tokens(&line).collect();
        println!("{:?}", t);
    }
}

Example

input

x1 = 42+y
(a)*10

expected output

["x1", "=", "42", "+", "y"]
["(", "a", ")", "*", "10"]

Example

input

12ab

expected output

["12", "ab"]
[]

Hint 1

In next, skip leading spaces, look at the first char (? returns None at the end), and compute the token length.

Hint 2

self.rest.find(|c: char| !c.is_ascii_digit()).unwrap_or(self.rest.len()) is the length of a digit run.

Hint 3

let (tok, rest) = self.rest.split_at(len); then store rest and return Some(tok).

Reference solution

use std::io;

struct Tokens<'a> {
    rest: &'a str,
}

fn tokens(s: &str) -> Tokens<'_> {
    Tokens { rest: s }
}

impl<'a> Iterator for Tokens<'a> {
    type Item = &'a str;
    fn next(&mut self) -> Option<&'a str> {
        self.rest = self.rest.trim_start_matches(' ');
        let first = self.rest.chars().next()?;
        let len = if first.is_ascii_alphabetic() {
            self.rest.find(|c: char| !c.is_ascii_alphanumeric()).unwrap_or(self.rest.len())
        } else if first.is_ascii_digit() {
            self.rest.find(|c: char| !c.is_ascii_digit()).unwrap_or(self.rest.len())
        } else {
            first.len_utf8()
        };
        let (tok, rest) = self.rest.split_at(len);
        self.rest = rest;
        Some(tok)
    }
}

fn main() {
    for line in io::stdin().lines() {
        let line = line.unwrap();
        let t: Vec<&str> = tokens(&line).collect();
        println!("{:?}", t);
    }
}

Exercise 3 · quiz

What do you get for free?

After writing impl Iterator for MyIter { type Item = u32; fn next(&mut self) -> Option<u32> { ... } }, which of these can you call on a MyIter without writing anything else? Select all that apply.

Answer

map and filter
sum
collect::<Vec<u32>>()
zip

All of Iterator's provided methods are built on next, so you get map, filter, sum, collect, zip, take, fold and dozens more. rev() additionally requires DoubleEndedIterator (a next_back method), because walking backwards cannot be derived from walking forwards.