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:
Collection<T>(consuming)&'a Collection<T>(delegating toself.iter())&'a mut Collection<T>(delegating toself.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>, holdingslice: &'a [T]and an indexpos: usizeworks easily because you can copy a shared reference&'a Tout of&self.slice[pos]insidefn 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 tonext()can never yield aliasing&'a mut Treferences 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]