Rust and Data Structures

Iterator Support for Custom Collections

To make a custom collection feel like a first-class citizen in Rust's for loops and functional pipelines, it should provide three distinct iterator modes corresponding to Rust's three ownership forms (&self, &mut self, and self), wired to both inherent methods (.iter(), .iter_mut()) and the IntoIterator trait.

The Three Standard Iteration Modes for `Collection<T>`:

  1. Shared Borrow (`for x in &col`     === `col.iter()`)     --> Yields `&'a T`
  2. Mutable Borrow (`for x in &mut col` === `col.iter_mut()`) --> Yields `&'a mut T`
  3. Consuming Move (`for x in col`      === `col.into_iter()`)--> Yields owned `T`

Implementing All Three IntoIterator Modes

When you write for item in &collection, the compiler desugars the loop into IntoIterator::into_iter(&collection). Therefore, implementing IntoIterator for:

  1. Collection<T> (consuming)
  2. &'a Collection<T> (delegating to self.iter())
  3. &'a mut Collection<T> (delegating to self.iter_mut())

makes all three for loop forms work seamlessly!

Whenever your custom collection wraps a standard slice or collection internally, you can either delegate to std::slice::Iter<'a, T> or write custom iterator state structs that also implement DoubleEndedIterator (enabling .rev()) and ExactSizeIterator (enabling .len()):

pub struct RingPair<T> {
    items: [T; 2],
}

impl<T> RingPair<T> {
    pub const fn new(first: T, second: T) -> Self {
        Self {
            items: [first, second],
        }
    }

    pub fn iter(&self) -> std::slice::Iter<'_, T> {
        self.items.iter()
    }

    pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
        self.items.iter_mut()
    }
}

/// 1. Consuming `IntoIterator` (`for x in pair`): yields owned `T` values!
impl<T> IntoIterator for RingPair<T> {
    type Item = T;
    type IntoIter = std::array::IntoIter<T, 2>;

    fn into_iter(self) -> Self::IntoIter {
        self.items.into_iter()
    }
}

/// 2. Shared Reference `IntoIterator` (`for x in &pair`): yields `&'a T`!
impl<'a, T> IntoIterator for &'a RingPair<T> {
    type Item = &'a T;
    type IntoIter = std::slice::Iter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

/// 3. Mutable Reference `IntoIterator` (`for x in &mut pair`): yields `&'a mut T`!
impl<'a, T> IntoIterator for &'a mut RingPair<T> {
    type Item = &'a mut T;
    type IntoIter = std::slice::IterMut<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter_mut()
    }
}

fn main() {
    let mut pair = RingPair::new(String::from("alpha"), String::from("beta"));

    // 1. Mutate in-place via `for s in &mut pair`:
    for s in &mut pair {
        s.make_ascii_uppercase();
    }

    // 2. Read without consuming via `for s in &pair`:
    for s in &pair {
        println!("Borrowed: {s}");
    }

    // 3. Consume by value in reverse via `pair.into_iter().rev()`:
    let reversed: Vec<String> = pair.into_iter().rev().collect();
    assert_eq!(reversed, vec!["BETA", "ALPHA"]);
    println!("Consumed in reverse: {reversed:?}");
}
Borrowed: ALPHA
Borrowed: BETA
Consumed in reverse: ["BETA", "ALPHA"]

Writing a Custom Mutable Iterator (IterMut) Safely

Why is writing a custom IterMut<'a, T> from scratch trickier than writing a shared Iter<'a, T>?

  • In Iter<'a, T>, holding slice: &'a [T] and an index pos: usize works easily because you can copy a shared reference &'a T out of &self.slice[pos] inside fn next(&mut self) -> Option<&'a T>.
  • In IterMut<'a, T>, fn next(&mut self) only has a short borrow &'next mut self, yet it must return a longer-lived mutable reference &'a mut T—and the borrow checker must be convinced that two calls to next() can never yield aliasing &'a mut T references to the same element!

To implement a custom IterMut<'a, T> in 100% safe Rust, store remaining: &'a mut [T] and use std::mem::take(&mut self.remaining) followed by split_first_mut():

pub struct SafeIterMut<'a, T> {
    remaining: &'a mut [T],
}

impl<'a, T> Iterator for SafeIterMut<'a, T> {
    type Item = &'a mut T;

    fn next(&mut self) -> Option<Self::Item> {
        // `mem::take` temporarily swaps `self.remaining` with `&mut []`,
        // giving us full `'a` ownership of the slice to split off the first element safely!
        let slice = std::mem::take(&mut self.remaining);
        let (first, rest) = slice.split_first_mut()?;
        self.remaining = rest;
        Some(first)
    }
}

fn main() {
    let mut nums = [1, 2, 3];
    for x in (SafeIterMut { remaining: &mut nums }) {
        *x *= 10;
    }
    println!("Mutated via 100% safe custom IterMut: {nums:?}");
}
Mutated via 100% safe custom IterMut: [10, 20, 30]