Custom iterator over some struct fields

mvolkmann · January 1, 2021, 11:05pm

I'm trying to learn how to implement a custom iterator.
In this example I have a Color struct with r, g, and b fields.
I want to be able to iterate over those values.
The code below is my start, but I'm not sure how to finish it.

use std::iter::Iterator;

struct Color {
    r: u8,
    g: u8,
    b: u8
}

impl Iterator for Color {
    type Item = u8;
    
    // How can I keep track of the last value returned?
    
    fn next(&mut self) -> Option<Self::Item> {
        // Return one of these in this sequence:
        // Some(self.r)
        // Some(self.g)
        // Some(self.b)
        // None
    }
}

fn main() {
    let color = Color {r: 100, g: 0, b: 150};
    for v in color.iter() {
        println!("{}", v);
    }
}

steffahn · January 1, 2021, 11:12pm

The trick is that, usually, the Iterator is its own custom type distinct from the datatype. E.g. Vec<T> does not implement Iterator, instead it implements IntoIterator, by means of which it can be turned into an std::vec::IntoIter<T>. Only that IntoIter type implements the Iterator trait.

Furthermore in order to also iterate over a Vec<T> when you only have a &Vec<T> or a &mut Vec<T>, there are two more IntoIterator implementations: for &Vec<T> and for &mut Vec<T>, allowing those to be turned into std::slice::Iter<'_, T> and std::slice::IterMut<'_, T>, respectively.

vitalyd · January 2, 2021, 12:04am

Here's a simple version using std::iter::once chaining: playground

mvolkmann · January 2, 2021, 12:24am

I've never encountered a return type beginning with impl before. Either I missed that when I read "The Rust Programming Language Book" or the book doesn't contain an example like that. Is that shorthand for defining that return type in its own impl statement?

How would you implement IntoIterator in your playground so line 21 could be this:

    for v in color {

steffahn · January 2, 2021, 12:44am

As an exercise of how a custom (by-immutable-reference) iterator type could work, try completing this example [this code is a bit contrived since this approach aims to avoid copying the u8s; but it generalizes to cases where the contents are not copyable]:

pub struct Color {
    r: u8,
    g: u8,
    b: u8
}
impl Color {
    fn iter(&self) -> color::Iter<'_> {
        self.into_iter()
    }
}

pub mod color {
    use std::iter::Iterator;
    use super::Color;
    enum State {
        FirstField,
        SecondField,
        ThirdField,
        Done,
    }
    use State::*;
    
    pub struct Iter<'a> {
        color: &'a Color,
        state: State,
    }
    
    impl<'a> IntoIterator for &'a Color {
        // TODO add missing items
    }
    
    impl<'a> Iterator for Iter<'a> {
        type Item = &'a u8;
        // TODO add missing items
    }
}

fn main() {
    let color = Color {r: 100, g: 0, b: 150};
    // alternatively, you can also write
    // `for v in &color {`
    for v in color.iter() {
        println!("{}", v);
    }
}

(playground)

(partial solution)

(complete solution)

quinedot · January 2, 2021, 12:50am

It means "I'm returning a concrete type that implements this trait, but I'm not telling you what that concrete type is". Note that it's still a single, concrete type (e.g. if returning a match expression, the arms would need to unify). The compiler infers the type. It gives you the flexibility to change the type without breaking compatibility, and you don't have to type out complicated types (as nested iterators can be). It can also involve closures, which you can't name. It takes away the ability to name the type (e.g. as the caller), and thus to do various type-specific things.