Implementing simple recursion as an iterator in Zig 0.11.0.

recursive-iterator.zig

const std = @import("std");

const Buffer = std.ArrayList(u8);

const source = [_]u8{ 1, 2, 3, 4, 5 };
const expectedOrder = [_]u8{ 1, 2, 3, 4, 5, 5, 4, 3, 2, 1 };

fn recursive(read: []const u8, write: *Buffer) !void {
    if (read.len > 0) {
        try write.append(read[0]);
        try recursive(read[1..], write);
        try write.append(read[0]);
    }
}

test "recursive" {
    var buffer = Buffer.init(std.heap.page_allocator);
    try recursive(&source, &buffer);
    try std.testing.expect(std.mem.eql(u8, try buffer.toOwnedSlice(), &expectedOrder));
}

const Iterator = struct {
    frame: Frame,
    stack: Stack,

    const Stack = std.ArrayList(Frame);
    const Frame = struct {
        read: []const u8,
        entryPoint: EntryPoint,
    };
    const EntryPoint = enum { a, b, c, done };

    fn init(read: []const u8) Iterator {
        return Iterator{
            .frame = .{
                .read = read,
                .entryPoint = .a,
            },
            .stack = Stack.init(std.heap.page_allocator),
        };
    }

    fn recurse(this: *Iterator, read: []const u8) !void {
        try this.stack.append(this.frame);
        this.frame = .{
            .read = read,
            .entryPoint = .a,
        };
    }

    fn backtrack(this: *Iterator) void {
        this.frame = this.stack.pop();
    }

    fn next(this: *Iterator) !?u8 {
        while (true) {
            const entryPoint = this.frame.entryPoint;

            // If the current frame is done executing...
            if (entryPoint == .done) {
                // ...then backtrack to the previous frame on the stack, if there is one...
                if (this.stack.items.len > 0) {
                    this.backtrack();
                    continue;
                }
                // ...otherwise, we are all done.
                return null;
            }

            // Prepare frame for the next round.
            this.frame.entryPoint = @enumFromInt(@intFromEnum(entryPoint) + 1);

            // Essentially the same code as in the recursive version.
            if (this.frame.read.len > 0) {
                // Using a switch to allow different entry points into the function body.
                switch (entryPoint) {
                    .a => {
                        // "yield" out a value by returning it. The caller will call next() again, resuming at the next entry point.
                        return this.frame.read[0];
                    },
                    .b => {
                        // Push the current frame to the stack and restart with a new
                        // frame, just like when calling a function recursively.
                        try this.recurse(this.frame.read[1..]);
                        continue;
                    },
                    .c => {
                        return this.frame.read[0];
                    },
                    // .done is handled at the top of the loop. Not really part of the business logik.
                    .done => unreachable,
                }
            }
        }
    }
};

test "iterator" {
    var buffer = Buffer.init(std.heap.page_allocator);
    var iterator = Iterator.init(&source);
    while (try iterator.next()) |value| {
        try buffer.append(value);
    }
    try std.testing.expect(std.mem.eql(u8, try buffer.toOwnedSlice(), &expectedOrder));
}