iterator.zig

const std = @import("std");
const assert = std.debug.assert;

const Iterator_Buffer_Size = 24;

// Returns a function pointer to 'ImplType.next: fn(&ImplType) ?T'
// This could be generalized by taking a comptime name: []const u8 and looking up the function.
// Also, for nullable function pointers, with some metaprogramming we could determine wether the
// function is present in ImplType and only included it then.
fn get_next_fn(comptime ImplType: type, comptime T: type) fn(usize)?T {
    return struct {
        fn next(self_type_erased: usize) ?T {
            const self = @intToPtr(&ImplType, self_type_erased);
            return @inlineCall(self.next);
        }
    }.next;
}

fn IteratorVTable(comptime T: type) type {
    return struct {
        const Self = this;

        // Self arguments turn into usize, we go back and forth with
        // @ptrToInt, @intToPtr
        next_fn: fn(self_type_erased: usize) ?T,
    
        fn init(comptime ImplType: type) Self {
            return Self {
                .next_fn = get_next_fn(ImplType, T),
            };
        }
    };
}

// The vtable for any implementation type is generated at comptime using 'comptime IteratorVTable(T).init(ImplType)'
// This forces the vtable to live in static memory of our binary, so we just take a pointer to it.
// This leads us to the expression '&comptime IteratorVTable(T).init(ImplType)'
// This is essentially what languages with classes and virtual functions do but implemented in user code rather than compiler
// code.
pub fn Iterator(comptime T: type) type {
    return struct {
        const Self = this;
        const Item = T;

        vtable: &const IteratorVTable(T),
        // extern union guarantees zig will not add a tag to the union.
        // we want this behaviour since we keep a flag byte ourselves.
        // The packed structs guarantee that zig will not add any fields,
        // that the fields will be in the correct order and that no padding
        // bytes will be added.
        // These guarantees are checked in the init methods with comptime asserts.
        data: extern union {
            small_buffer: packed struct {
                mem: [Iterator_Buffer_Size - 1]u8,
                flag_byte: u8,
            },
            heap_ptr: packed struct {
                ptr: &u8,
                alloc: &std.mem.Allocator,
                unused_mem: [Iterator_Buffer_Size - @sizeOf(&u8) - @sizeOf(&std.mem.Allocator) - 1]u8,
                flag_byte: u8,
            },
        },

        // For now, any value of the flag byte except 0 is used to indicate the object lives in the small buffer.
        // We could use a single bit for this information and use the 7 remaining bits for context specific
        // (heap/small buffer) flags.
        pub fn is_stored_inline(self: &const Self) bool {
            return self.data.small_buffer.flag_byte != 0;
        }

        // Use with care.
        // Initializes the vtable pointer and returns a pointer into the small buffer of the type requested.
        // You should use this functions to initialize undefined interfaces. 
        pub fn init_inplace_ptr(self: &Self, comptime ImplType: type) &align(1) ImplType {
            comptime assert(@sizeOf(@typeOf(self.data)) == Iterator_Buffer_Size * @sizeOf(u8));
            comptime assert(@sizeOf(ImplType) < Iterator_Buffer_Size);

            self.vtable = &comptime IteratorVTable(T).init(ImplType);
            self.data.small_buffer.flag_byte = 1;
            return @ptrCast(&align(1) ImplType, &self.data.small_buffer.mem[0]);
        }

        // Copy the bytes from 'obj' into the interface's inline buffer and sets the vtable pointer. (comptime checked) 
        pub fn init_small(obj: var) Self {
            const ImplType = @typeOf(obj).Child;
            comptime assert(@sizeOf(@typeOf(self.data)) == Iterator_Buffer_Size * @sizeOf(u8));
            comptime assert(@sizeOf(ImplType) < Iterator_Buffer_Size);

            // Just copy data on our buffer.
            var self: Self = undefined;
            self.vtable = &comptime IteratorVTable(T).init(ImplType);
            self.data.small_buffer.flag_byte = 1;
            std.mem.copy(u8, self.data.small_buffer.mem[0..], @ptrCast(&const u8, obj)[0..@sizeOf(ImplType)]);
            return self;
        }

        // Copy the bytes from 'obj' into the interface, allocating if necessary and sets the vtable pointer.
        pub fn init(alloc: &std.mem.Allocator, obj: var) !Self {
            const ImplType = @typeOf(obj).Child;

            var self: Self = undefined;
            self.vtable = &comptime IteratorVTable(T).init(ImplType);

            comptime assert(@sizeOf(@typeOf(self.data)) == Iterator_Buffer_Size * @sizeOf(u8));

            if (@sizeOf(ImplType) >= Iterator_Buffer_Size) {
                self.data.heap_ptr.alloc = alloc;
                // Heap allocate space for our object.
                self.data.heap_ptr.ptr = @ptrCast(&u8, try alloc.create(ImplType));
                // Copy memory into the heap.
                std.mem.copy(u8, self.data.heap_ptr.ptr[0..@sizeOf(ImplType)], @ptrCast(&const u8, obj)[0..@sizeOf(ImplType)]);
                self.data.heap_ptr.flag_byte = 0;
                return self;
            }

            // Just copy data on our buffer.
            self.data.small_buffer.flag_byte = 1;
            std.mem.copy(u8, self.data.small_buffer.mem[0..], @ptrCast(&const u8, obj)[0..@sizeOf(ImplType)]);
            return self;
        }

        // TODO: We could add a deinit function pointer to our vtable (even make it optional)
        // We would call that function here before we deallocated our heap memory (if we have it).
        pub fn deinit(self: &Self) void {
            if (!self.is_stored_inline()) {
                self.data.heap_ptr.alloc.destroy(self.data.heap_ptr.ptr);
            }
        }

        // Get our instance pointer from the small buffer or thea heap, call into the vtable
        // with a type erased pointer.
        pub fn next(self: &Self) ?T {
            if (!self.is_stored_inline()) {
                return self.vtable.next_fn(@ptrToInt(self.data.heap_ptr.ptr));
            }

            return self.vtable.next_fn(@ptrToInt(&self.data.small_buffer.mem[0]));
        }

        // TODO: Rest of methods
    };
}

pub fn Once(comptime Item: type) type {
    const Iter = Iterator(Item);
    return struct {
        const Self = this;
        item: ?Item,

        // Initializes a Once iterator into an Iterator interface object.
        // We need no allocator since Once is guaranteed to fit in our iterator
        // small buffer.
        // Note that Once does not need to know about Iterator(T), we just choose
        // to provide a helper function to get an Iterato interface object directly.
        pub fn init_iter(item: Item) Iter {
            // We use init_inplace_ptr to avoid a copy.
            // We initialize self inside the iterator's small buffer directly,
            // instead of using init_small (in which case we would have to initialize
            // self on the stack then copy it into the iterator's small buffer).
            var iter: Iter = undefined;
            var self_ptr = iter.init_inplace_ptr(Self);
            self_ptr.item = item;
            return iter;
        }

        // Initializes a Once object.
        pub fn init(item: Item) Self {
            return Self {
                .item = item,
            };
        }

        pub fn next(self: &Self) ?Item {
            if (self.item != null) {
                var ret = self.item;
                self.item = null;
                return ret;
            }
            return null;
        }
    };
}

const BigTestIterator = struct {
    const Self = this;

    data: [25]u8,

    fn next(self: &Self) ?usize {
        return 42;
    }
};

test "Once" {
    var once = Once(usize).init_iter(10);

    assert(??once.next() == 10);

    var next_null = once.next();
    assert(next_null == null);

    var big_test_it: BigTestIterator = undefined;
    // This fails at comptime:
    // var it = Iterator(usize).init_small(big_test_it);

    var it = Iterator(usize).init(std.debug.global_allocator, big_test_it) catch unreachable;
    defer it.deinit();
    assert(??it.next() == 42);
}