Rust for C-Programmers

13.3 Creating Custom Iterators

While standard library iterators cover many use cases, you’ll often need to make your own data structures iterable. When creating custom iterators, there are generally two structural approaches:

1. The Type is the Iterator: For simple cases, the type itself can hold the necessary iteration state (like a current index or value) and directly implement the Iterator trait, including the next() method. Instances of this type can then be used directly in loops or iterator chains. We will see this pattern with a Counter example.
2. The Type Produces an Iterator: More commonly, especially for types acting as collections, the type itself doesn’t implement Iterator. Instead, it implements the IntoIterator trait. Its into_iter() method constructs and returns a separate iterator struct (which holds the iteration state and implements Iterator with the next() logic). This is the pattern used by standard collections like Vec and the one we’ll initially demonstrate for a custom Pixel struct.

A key benefit of implementing the Iterator trait (either directly on your type or on a separate iterator struct) is that you automatically gain access to a wide array of powerful adapter and consumer methods defined directly on the trait itself (like map, filter, fold, sum, collect, and many others shown in Section 13.2). These methods have default implementations written in terms of the required next() method. Therefore, by simply providing the core next() logic for your specific type, you enable users to immediately leverage the entire rich ecosystem of standard iterator operations on your custom iterator, just like they would with standard library iterators.

Let’s illustrate these approaches with examples.

13.3.1 Example 1: Iterating Over Struct Fields (Manual Implementation)

This approach follows the second pattern mentioned above: the Pixel struct implements IntoIterator to produce separate iterator structs (PixelIter, PixelIterMut, etc.) which implement Iterator. This is general but can involve boilerplate code.

#[derive(Debug, Clone, Copy)] // Added derives for easier use later
struct Pixel {
    r: u8,
    g: u8,
    b: u8,
}

// --- Consuming Iterator (Yields owned u8) ---

struct PixelIntoIterator {
    pixel: Pixel, // Owns the pixel data
    index: u8,    // State: which component is next (0=r, 1=g, 2=b)
}

impl Iterator for PixelIntoIterator {
    type Item = u8; // Yields owned u8 values

    fn next(&mut self) -> Option<Self::Item> {
        let result = match self.index {
            0 => Some(self.pixel.r),
            1 => Some(self.pixel.g),
            2 => Some(self.pixel.b),
            _ => None, // Sequence exhausted
        };
        self.index = self.index.wrapping_add(1); // Use wrapping_add for safety
        result
    }
}

// Implement IntoIterator for Pixel to enable `for val in pixel`
impl IntoIterator for Pixel {
    type Item = u8;
    type IntoIter = PixelIntoIterator;

    fn into_iter(self) -> Self::IntoIter {
        PixelIntoIterator { pixel: self, index: 0 }
    }
}

// --- Immutable Reference Iterator (Yields &u8) ---

// Lifetime 'a ensures the iterator doesn't outlive the borrowed Pixel
struct PixelIter<'a> {
    pixel: &'a Pixel, // Holds an immutable reference
    index: u8,
}

impl<'a> Iterator for PixelIter<'a> {
    type Item = &'a u8; // Yields immutable references

    fn next(&mut self) -> Option<Self::Item> {
        let result = match self.index {
            0 => Some(&self.pixel.r),
            1 => Some(&self.pixel.g),
            2 => Some(&self.pixel.b),
            _ => None,
        };
        self.index = self.index.wrapping_add(1);
        result
    }
}

// Implement IntoIterator for &Pixel to enable `for val_ref in &pixel`
impl<'a> IntoIterator for &'a Pixel {
    type Item = &'a u8;
    type IntoIter = PixelIter<'a>;

    fn into_iter(self) -> Self::IntoIter {
        PixelIter { pixel: self, index: 0 }
    }
}

// --- Mutable Reference Iterator (Yields &mut u8) ---

struct PixelIterMut<'a> {
    pixel: &'a mut Pixel, // Holds a mutable reference
    index: u8,
}

impl<'a> Iterator for PixelIterMut<'a> {
    type Item = &'a mut u8; // Yields mutable references

    // Returning mutable references from `next` when iterating over mutable
    // fields of a struct borrowed mutably can be tricky for the borrow checker.
    // Using raw pointers temporarily inside `next` is one pattern to handle this,
    // though it requires `unsafe`. It bypasses the borrow checker's static
    // analysis for this specific, localized operation, relying on the programmer
    // to ensure safety (which holds here as we access distinct fields per index).
    fn next(&mut self) -> Option<Self::Item> {
        let pixel_ptr: *mut Pixel = self.pixel; // Get raw pointer to the mutable pixel
        let result = match self.index {
          // Safety: `pixel_ptr` is valid, and index ensures we access distinct fields
          // mutably within the lifetime 'a.
            0 => Some(unsafe { &mut (*pixel_ptr).r }),
            1 => Some(unsafe { &mut (*pixel_ptr).g }),
            2 => Some(unsafe { &mut (*pixel_ptr).b }),
            _ => None,
        };
        self.index = self.index.wrapping_add(1);
        result
    }
}

// Implement IntoIterator for &mut Pixel to enable `for val_mut in &mut pixel`
impl<'a> IntoIterator for &'a mut Pixel {
    type Item = &'a mut u8;
    type IntoIter = PixelIterMut<'a>;

    fn into_iter(self) -> Self::IntoIter {
        PixelIterMut { pixel: self, index: 0 }
    }
}

// Optional: Add convenience methods like standard collections
impl Pixel {
    fn iter(&self) -> PixelIter<'_> {
        self.into_iter() // Equivalent to (&*self).into_iter()
    }

    fn iter_mut(&mut self) -> PixelIterMut<'_> {
        self.into_iter() // Equivalent to (&mut *self).into_iter()
    }
}

fn main() {
    let pixel1 = Pixel { r: 255, g: 0, b: 128 };

    println!("Iterating by value (consumes pixel1):");
    // Note: pixel1 cannot be used after this loop because it's moved
    for val in pixel1 {
        println!(" - Value: {}", val);
    }
    // println!("{:?}", pixel1); // Error: use of moved value

    let pixel2 = Pixel { r: 10, g: 20, b: 30 };
    println!("\nIterating by immutable reference:");
    for val_ref in pixel2.iter() { // or `for val_ref in &pixel2`
        println!(" - Ref: {}", val_ref); // *val_ref is u8
    }
    println!("Pixel 2 after iter: {:?}", pixel2); // pixel2 is still usable

    let mut pixel3 = Pixel { r: 100, g: 150, b: 200 };
    println!("\nIterating by mutable reference:");
    for val_mut in pixel3.iter_mut() { // or `for val_mut in &mut pixel3`
        *val_mut = val_mut.saturating_add(10); // Modify value safely
        println!(" - Mut Ref: {}", *val_mut);
    }
    println!("Pixel 3 after iter_mut: {:?}", pixel3);

    let pixel4 = Pixel { r: 2, g: 3, b: 4 };
    // Using methods inherited from the Iterator trait:
    let sum: u16 = pixel4.iter().map(|&v| v as u16).sum();
    println!("\nSum using iter(): {}", sum); // Output: 9

    let product: u32 = pixel4.into_iter().map(|v| v as u32).product();
    println!("Product using into_iter(): {}", product); // Output: 24
    // pixel4 is consumed here
}

Key points from this example:

• Separate iterator structs (PixelIntoIterator, PixelIter, PixelIterMut) manage state and hold either owned data, an immutable reference, or a mutable reference.
• Implementing IntoIterator for Pixel, &Pixel, and &mut Pixel makes the struct work seamlessly with for loops in all three modes.
• Lifetimes ('a) are crucial for the reference iterators.
• The unsafe block in PixelIterMut::next demonstrates a pattern sometimes needed to safely return mutable references to different fields across calls, bypassing borrow checker limitations within the method body.
• Crucially, even though we only implemented next(), we could still call .map() and .sum() or .product() because those methods are provided by the Iterator trait itself.

13.3.2 Example 2: A Simple Self-Contained Iterator (Counter)

Sometimes, the iterator is the primary object, holding its own state directly, rather than iterating over a separate collection. This follows the first structural pattern mentioned earlier: the type implements Iterator directly.

// An iterator that counts from 'start' up to 'end' (inclusive).
struct Counter {
    current: u32,
    end: u32,
}

impl Counter {
    fn new(start: u32, end: u32) -> Self {
        Counter { current: start, end }
    }
}

impl Iterator for Counter {
    type Item = u32;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current <= self.end {
            let value = self.current;
            // Use saturating_add for safety against overflow, though unlikely here
            self.current = self.current.saturating_add(1);
            Some(value)
        } else {
            None // Signal the end of the iteration
        }
    }
}

fn main() {
    println!("Counting from 1 to 5:");
    // The Counter struct itself implements Iterator
    let counter1 = Counter::new(1, 5);
    for count in counter1 { // `for` loop works directly on an Iterator
        println!(" - {}", count);
    }

    // Iterator methods like `sum` can be called directly on Counter
    // because it implements Iterator.
    let sum_of_range: u32 = Counter::new(10, 15).sum();
    println!("\nSum of range 10 to 15: {}", sum_of_range); // 10+..+15 = 75

    let mut counter2 = Counter::new(1, 3);
    assert_eq!(counter2.next(), Some(1));
    assert_eq!(counter2.next(), Some(2));
    assert_eq!(counter2.next(), Some(3));
    assert_eq!(counter2.next(), None);
    assert_eq!(counter2.next(), None); // Should remain None (FusedIterator behavior)
}

In this Counter example, we didn’t need IntoIterator. Why?

• The Counter struct itself implements the Iterator trait. It holds its own state (current, end).
• The for loop and methods like sum() are designed to work with any type that implements Iterator. If the type passed to them already is an iterator (like our counter1 variable), they use it directly.
• If, however, the type used in a for loop (like a Vec or our Pixel struct) is not an iterator itself, the loop requires that type to implement the IntoIterator trait. The loop then implicitly calls the .into_iter() method on that type to obtain the actual iterator it needs.
• Therefore, IntoIterator is primarily needed to define how to get an iterator from another type (like a collection). Since Counter is already the iterator, this step isn’t required for it.

13.3.3 Leveraging Array Iterators via Delegation

The manual implementation for Pixel in the first example works, but involves significant boilerplate code. If the data within your struct can be logically represented as a standard collection type, like an array or a slice, you can often simplify the implementation significantly by delegating to the standard library’s existing, optimized iterators.

This approach involves:

1. Storing the data internally in a standard collection (like an array).
2. Implementing IntoIterator for your type (and its references) by calling the corresponding into_iter(), .iter(), or .iter_mut() methods on the internal collection.

Let’s revise the Pixel struct to hold its components in an internal array [u8; 3] and see how this simplifies the iterator implementations.

use std::slice::{Iter, IterMut}; // Import slice iterators for type annotations

#[derive(Debug, Clone, Copy)]
struct PixelArray {
    // Store components in an array
    channels: [u8; 3], // [r, g, b]
}

impl PixelArray {
    fn new(r: u8, g: u8, b: u8) -> Self {
        PixelArray { channels: [r, g, b] }
    }

    // Convenience accessors (optional but helpful)
    fn r(&self) -> u8 { self.channels[0] }
    fn g(&self) -> u8 { self.channels[1] }
    fn b(&self) -> u8 { self.channels[2] }
}

// Implement IntoIterator for PixelArray (consuming iteration)
// This delegates to the array's consuming iterator.
impl IntoIterator for PixelArray {
    type Item = u8;
    // Delegate to the array's consuming iterator type: `std::array::IntoIter`
    type IntoIter = std::array::IntoIter<u8, 3>;

    fn into_iter(self) -> Self::IntoIter {
        // Arrays implement IntoIterator, so we just call it on the internal array
        self.channels.into_iter()
    }
}

// --- We DO need explicit impl IntoIterator for &PixelArray ---
// To enable `for item in &my_pixel_array`, we must implement `IntoIterator`
// for the reference type `&PixelArray`. We achieve this easily by
// delegating to the `.iter()` method of the internal `channels` array,
// which returns an iterator yielding `&u8`.
impl<'a> IntoIterator for &'a PixelArray {
    type Item = &'a u8;
    // The iterator type yielded by `.iter()` on an array/slice is `std::slice::Iter`
    type IntoIter = Iter<'a, u8>;

    fn into_iter(self) -> Self::IntoIter {
        // Call `.iter()` on the internal array
        self.channels.iter()
    }
}

// --- We DO need explicit impl IntoIterator for &mut PixelArray ---
// Similarly, to enable `for item in &mut my_pixel_array`, we implement
// `IntoIterator` for `&mut PixelArray`. This implementation delegates
// to the internal array's `.iter_mut()` method, which returns an
// iterator yielding `&mut u8`.
impl<'a> IntoIterator for &'a mut PixelArray {
    type Item = &'a mut u8;
     // The type yielded by `.iter_mut()` on an array/slice is `std::slice::IterMut`
    type IntoIter = IterMut<'a, u8>;

    fn into_iter(self) -> Self::IntoIter {
         // Call `.iter_mut()` on the internal array
        self.channels.iter_mut()
    }
}

// By providing these implementations, we correctly leverage the standard
// library's efficient slice iterators (`slice::Iter` and `slice::IterMut`)
// for our custom type, without needing to rewrite the iteration logic itself.

// Optional convenience methods (often added for discoverability, mirroring std lib)
impl PixelArray {
    pub fn iter(&self) -> Iter<'_, u8> {
        self.channels.iter() // Delegate directly
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, u8> {
        self.channels.iter_mut() // Delegate directly
    }
}

fn main() {
    let pixel = PixelArray::new(255, 0, 128);

    println!("Iterating by value (consuming):");
    // `for val in pixel` calls `pixel.into_iter()`
    for val in pixel {
        println!(" - Value: {}", val);
    }
    // pixel is consumed here

    let pixel_ref = PixelArray::new(10, 20, 30);
    println!("\nIterating by immutable ref. (via impl IntoIterator for &PixelArray):");
    // `for val_ref in &pixel_ref` now correctly calls `(&pixel_ref).into_iter()`
    for val_ref in &pixel_ref { // This now works
        println!(" - Ref: {}", val_ref); // val_ref is &u8
    }
     // Example using the convenience method explicitly:
    // for val_ref in pixel_ref.iter() { println!(" - Ref: {}", val_ref); }

    let mut pixel_mut = PixelArray::new(100, 150, 200);
    println!("\nIterating by mutable ref. (via impl IntoIterator for &mut PixelArray):");
    // `for val_mut in &mut pixel_mut` now correctly calls `(&mut pixel_mut).into_iter()`
    for val_mut in &mut pixel_mut { // This now works
        *val_mut = val_mut.saturating_sub(10); // Modify
        println!(" - Mut Ref: {}", *val_mut); // val_mut is &mut u8
    }
    println!("Pixel after mut iteration: {:?}", pixel_mut);
    // Example using the convenience method explicitly:
    // for val_mut in pixel_mut.iter_mut() { *val_mut += 5;
    // println!(" - Mut Ref: {}", *val_mut); }

    // We can still use map, sum etc. because the iterators produced
    // (`std::array::IntoIter`, `slice::Iter`, `slice::IterMut`) implement Iterator.
    let pixel_sum = PixelArray::new(5, 6, 7);
    let sum: u16 = pixel_sum.iter().map(|&v| v as u16).sum();
     println!("\nSum using iter() on PixelArray: {}", sum); // Output: 18
}

This section demonstrates how to make a struct iterable in all three modes by containing a standard collection (an array in this case) and implementing the necessary IntoIterator traits via simple delegation. This is often much less work and less error-prone than implementing the next() logic manually, while also benefiting from the performance of the standard library’s iterators.