How to implement multi-valued iterator pattern in Rust?

Question

I have an iterator-type-object that can return zero, one or more items each time it's called. I want to implement a standard Iter API, i.e. next returns Option<Self::Item>, so it can be consumed item by item.

In Clojure I would probably do this with mapcat ("map and concatenate").

My current solution (thanks to @Ryan) uses flat_map but still requires a lot of allocation:

// Desired input:
// A stateful object that implements an iterator which returns a number of results each time.
// The real code is a bit more complicated, this is the minimal example.
struct MyThing {
    counter: i32,
}

impl Iterator for MyThing {
    type Item = Vec<String>;

    fn next(&mut self) -> Option<Vec<String>> {
        self.counter += 1;
        if self.counter == 4 {
            self.counter = 1;
        }

        match self.counter {
            1 => Some(vec!["One".to_string()]),
            2 => Some(vec!["One".to_string(), "Two".to_string()]),
            3 => Some(vec![
                "One".to_string(),
                "Two".to_string(),
                "Three".to_string(),
            ]),
            _ => Some(vec![]),
        }
    }
}

fn main() {
    let things = MyThing { counter: 0 };

    // Missing piece, though the following line does the job:
    let flattened = things.flat_map(|x| x);

    // However this requires a heap allocation at each loop.

    // Desired output: I can iterate, item by item.
    for item in flattened {
        println!("{:?}", item);
    }
}

Given the innovative things I have seen, I wonder if there's a more idiomatic, less expensive way of accomplishing this pattern.

Answer 1

If you know how to generate the "inner" values programmatically, replace Vec<String> with a struct you define that implements Iterator<Item = String>. (Technically only IntoIterator is necessary, but Iterator is sufficient.)

struct Inner {
    index: usize,
    stop: usize,
}

impl Inner {
    fn new(n: usize) -> Self {
        Inner { index: 0, stop: n }
    }
}

impl Iterator for Inner {
    type Item = String;

    fn next(&mut self) -> Option<String> {
        static WORDS: [&str; 3] = ["One", "Two", "Three"];
        let result = if self.index < self.stop {
            WORDS.get(self.index).map(|r| r.to_string())
        } else {
            None
        };
        self.index += 1;
        result
    }
}

Because Inner implements Iterator<Item = String>, it can be iterated over much like Vec<String>. But Inner does not have to pre-allocate a Vec and consume items one by one; it can lazily create each String on demand.

The "outer" iterator is just a struct that implements Iterator<Item = Inner>, likewise constructing each Inner lazily:

struct Outer {
    counter: i32,
}

impl Iterator for Outer {
    type Item = Inner;

    fn next(&mut self) -> Option<Inner> {
        self.counter = 1 + self.counter % 3;

        Some(Inner::new(self.counter as usize))
    }
}

As you know, Iterator::flat_map flattens nested structure, so something like the following works:

let things = Outer { counter: 0 };

for item in things.flat_map(|x| x).take(100) {
    println!("{:?}", item);
}

In real-life code, Inner and Outer are probably pretty different from this example most of the time. For example, it's not necessarily possible to write Inner without doing the equivalent of allocating a Vec. So the precise shapes and semantics of these iterators depend on concrete information about the use case.

---

The above assumes that Inner is somehow useful, or easier to implement on its own. You could easily write a single struct that iterates over the sequence without needing to be flattened, but you have to also put the inner iterator state (the index field) into Outer:

struct Outer {
    index: usize,
    counter: i32,
}

impl Iterator for Outer {
    type Item = String;

    fn next(&mut self) -> Option<String> {
        static WORDS: [&str; 3] = ["One", "Two", "Three"];
        let result = WORDS.get(self.index).map(|r| r.to_string());
        self.index += 1;
        if self.index >= self.counter as usize {
            self.counter = 1 + self.counter % 3;
            self.index = 0;
        };
        result
    }
}

fn main() {
    let things = Outer { counter: 1, index: 0 };

    for item in things.take(100) {
        println!("{:?}", item);
    }
}