Rust: Implementing the Iterator Trait

Written: September 23, 2024

In my [previous blog post], I introduced a C++ iterator to elegantly extract packets from a contiguous sequence of bytes. If you haven’t read it, I recommend going through the first section to familiarize yourself with the practical example. We’ll be tackling the same problem today, but this time using Rust. So, buckle up!

To recap briefly, the problem involves a contiguous sequence of bytes where a protocol exists to locate the next packet within that sequence. Our goal is to create an interface that extracts slices from this sequence based on a user-provided protocol. A typical use case might look like this:

for packet in PacketIterator(&buffer, protocol) {
    // Process slices of packets
}

Unlike C++, where anything that behaves like an iterator is considered one, Rust requires us to explicitly define a struct that implements the Iterator trait. Let’s dive into defining the necessary states for our iterator.

Defining the iterator

First, the iterator needs to keep track of the current packet's starting point and the available length, which essentially is a slice of the original buffer. Additionally, we need a way to determine the size of the next packet, which will be specified by a user-provided protocol. This protocol will know how to extract the size from the packet header and will be passed as a function pointer that accepts a slice of the buffer and returns the size of the next packet, or zero if no packet is found.

To ensure flexibility, we’ll allow the underlying buffer type to be generic. This leads us to the following struct:

pub struct PacketIterator<'a, T, F'> where
    F: Fn(&[T]) -> usize {
    buffer: &'a [T],
    get_size_func: F,
}

Additionally, we'll provide a `new()` method allowing initialization:

impl<'a, T, F'> PacketIterator<'a, T, F'> where
    F: Fn(&[T]) -> usize {
    pub fn new(buffer: &'a [T], protocol: F) -> PacketIterator<T, F> {
        PacketIterator {
            buffer: buffer,
            get_size_func: protocol,
        }
    }
}

Implementing the Iterator trait

Now that we have our struct, the next step is to implement the Iterator trait. To do this, we need to define two things:

• type Item: Specifies the type of element being iterated over, which, in our case, is a slice.
• fn next(&mut self) -> Option<Self::Item>: This function advances the iterator and returns the next value, or `None` if the iteration is complete.

We’ll consider the iteration complete if the remaining buffer is empty, or if the next packet size is zero or exceeds the available buffer length. With that in mind, here’s our implementation:

impl<'a, T, F> Iterator for PacketIterator<'a, T, F> where
    F: Fn(&[T]) -> usize {
    type Item = &'a [T];

    fn next(&mut self) -> Option<Self::Item> {
        if self.buffer.is_empty() {
            return None
        }
        let next_packet_size = (self.get_size_func)(self.buffer);
        if next_packet_size == 0 || next_packet_size > self.buffer.len() {
            return None
        }
        let next_packet = &self.buffer[0..next_packet_size];
        self.buffer = &self.buffer[next_packet_size..];
        Some(next_packet)
    }
}

And there you have it! The iterator is fully set up, ready to extract packet slices.

Use case example

Let’s consider a simple protocol where the size of each packet is stored in the first integer of the packet. Using this protocol, we can iterate over packets like so:

let buffer = vec![2, 0, 4, 1, 0, 0, 3, 0, 8];
// Packet 1: [2, 0]
// Packet 2: [4, 1, 0, 0]
// Packet 3: [3, 0, 8]

let protocol = |buffer: &[i32]| -> usize {
    if buffer.is_empty() {
        return 0
    }
    buffer[0] as usize
};

for packet in PacketIterator::new(&buffer, protocol) {
    // Packet 1: [2, 0]
    // Packet 2: [4, 1, 0, 0]
    // Packet 3: [3, 0, 8]
}

For greater flexibility and reusability, we can alternatively encapsulate the protocol in a function and return the iterator like this:

pub fn my_packet_iterator(buffer: &[i32]) 
    -> PacketIterator<i32, impl Fn(&[i32]) -> usize> {
    let protocol = |buffer: &[i32]| -> usize {
        if buffer.is_empty() {
            return 0
        }
        buffer[0] as usize
    };
    PacketIterator::new(buffer, protocol)
}

In that case extracting packets is as simple as:

let buffer = vec![2, 0, 4, 1, 0, 0, 3, 0, 8];

for packet in my_packet_iterator(&buffer) {
    // Packet 1: [2, 0]
    // Packet 2: [4, 1, 0, 0]
    // Packet 3: [3, 0, 8]
}

Summary

By designing a custom iterator, we've shown how Rust enables elegant, efficient iteration, balancing flexibility and performance. Whether for legacy or modern systems, a well-constructed iterator is invaluable. Enjoy writing clean, idiomatic Rust!

(You can find the source code for this iterator in my repository.)