Zig Iterators
Why use an iterator vs a for loop is highly contextual. Rather than dive into when or why use an iterator, let’s focus on what an iterator is and how to write one for your collection. In the example below we’re creating a generic Slice with very basic features, enough to implement the focus of this post, the iterator.
const std = @import("std");
const print = std.debug.print;
const assert = std.debug.assert;
pub fn Slice(comptime T: type) type {
return struct {
const Self = @This();
allocator: std.mem.Allocator,
buffer: []T,
len: usize,
pub fn init(allocator: std.mem.Allocator) Self {
return Self{
.allocator = allocator,
.buffer = &[_]T{},
.len = 0,
};
}
pub fn deinit(self: *Self) void {
self.allocator.free(self.buffer);
}
fn resize(self: *Self, new_capacity: usize) !void {
const new_buffer = try self.allocator.alloc(T, new_capacity);
@memcpy(new_buffer[0..self.len], self.buffer[0..self.len]);
self.allocator.free(self.buffer);
self.buffer = new_buffer;
}
pub fn append(self: *Self, item: T) !void {
if (self.len >= self.buffer.len) {
try self.resize(if (self.buffer.len == 0) 1 else self.buffer.len * 2);
}
self.buffer[self.len] = item;
self.len += 1;
}
pub fn iterator(self: *Self) SliceIterator {
return SliceIterator{ .slice = self, .index = 0 };
}
pub const SliceIterator = struct {
slice: *Self,
index: usize,
pub fn next(self: *SliceIterator) ?T {
if (self.index >= self.slice.len) {
return null;
}
const item = self.slice.buffer[self.index];
self.index += 1;
return item;
}
};
};
}
Our slice is a generic function that takes in a type and has a few useful methods to do some basic slice stuff.
- An
init(allocator: std.mem.Allocator)method to initialize the slice. - A
deinit()method to deallocate our buffer. - A
resize(new_capacity: usize)method to increase the size of our buffer as we add items to ensure the capacity meets demand. - An
append(item: T)method to add new items to the slice.
The append(item: T) method will automatically scale the buffer by doubling the capacity when we need more room for items by calling the resize(new_capacity: usize) as needed.
However, the primary focus here is the nested SliceIterator struct and the iterator() method. An iterator is nothing more than a struct that holds our Slice and an index. It has a next() method that will return the item in the buffer at the index, and if that index is out-of-bounds it returns null, which means we’ve reached the end.
The iterator() method returns a new SliceIterator setting the index to 0.
pub fn iterator(self: *Self) SliceIterator {
return SliceIterator{ .slice = self, .index = 0 };
}
The SliceIterator struct contains the next() method that is called in a while loop. This method will check if we’ve reached the end of our slice, if not it’ll increment the index counter and return the current item. This mutating of the index counter is the reason an iterator is declared as a mutable variable instead of constant. If we reached the end of our slice, we return null. This tells the while loop to stop iterating.
pub fn next(self: *SliceIterator) ?T {
if (self.index >= self.slice.len) {
return null;
}
const item = self.slice.buffer[self.index];
self.index += 1;
return item;
}
Let’s write a quick test to see that in action.
test "slice iterator" {
const allocator = std.heap.page_allocator;
var slice = Slice(u32).init(allocator);
try slice.append(10);
try slice.append(20);
try slice.append(30);
try slice.append(40);
try slice.append(50);
assert(slice.len == 5);
var it = slice.iterator();
while (it.next()) |item| {
assert(item > 0);
}
}
Hopefully this helps demystify iterators.