auto_promoting_array_list.zig

const std = @import("std");

const need_auto_promoting = @sizeOf(usize) > 4;
const SmallInt = if (need_auto_promoting) u32 else usize;

pub fn AutoPromotingArrayList(comptime T: type) type {
    return struct {
        pointer: [*]T,
        small_len: SmallInt,
        small_capacity: SmallInt,

        const Self = @This();

        const Header = packed struct {
            len: usize,
            capacity: usize,
        };

        pub fn len(self: Self) usize {
            if (!need_auto_promoting or self.small_len != 0 or self.small_capacity != 0) {
                return self.small_len;
            } else {
                return self.largeHeader().len;
            }
        }

        pub fn capacity(self: Self) usize {
            if (!need_auto_promoting or self.small_len != 0 or self.small_capacity != 0) {
                return self.small_capacity;
            } else {
                return self.largeHeader().capacity;
            }
        }

        pub fn items(self: *Self) []T {
            return self.pointer[0..self.len()];
        }

        pub fn append(self: *Self, item: T) !void {
            const new_len = self.len() + 1;
            try self.ensureCapacity(new_len);
            self.setLen(new_len);
        }

        pub fn largeHeader(self: Self) *Header {
            @setCold(true);
            return @ptrCast([*]Header, self.pointer) - 1;
        }

        pub fn ensureCapacity(self: *Self, new_capacity: usize) !void {
            if (self.capacity() < new_capacity) {
                const new_ptr = self.allocate(new_capacity);
                const old_items = self.items();
                std.mem.copy(T, new_ptr[0..new_capacity], old_items);
                std.testing.allocator.free(old_items);
                if (new_capacity > std.math.maxInt(SmallInt)) {
                    self.largeHeader().* = .{
                        .capacity = new_capacity,
                        .len = self.len(),
                    };
                } else {
                    self.small_capacity = @intCast(SmallInt, new_capacity);
                }
            }
        }

        fn allocate(self: *Self, new_capacity: usize) ![*]T {
            var n = std.math.mul(@sizeOf(T), new_capacity)
            const needed_header_bytes = if (need_auto_promoting) blk: {
                // We must ask for alignment of T, and therefore space for
                // the header must be a multiple of @alignOf(T).
                const needed_header_bytes = (@sizeOf(Header) + (@alignOf(T) - 1)) / @alignOf(T);
                // Integer overflow here means failure to allocate.
                n = std.math.add(n, needed_header_bytes) catch return error.OutOfMemory;
                break :blk needed_header_bytes;
            } else 0;
            // Allocate extra space to make it growable.
            n = std.math.add(n, n / 2) catch std.math.maxInt(usize);

            const bytes = try std.testing.allocator.alignedAlloc(u8, @alignOf(T), n);
            return @ptrCast([*]T, @alignCast(@alignOf(T), bytes.ptr + needed_header_bytes));
        }

        fn setLen(self: *Self, new_len: usize) void {
            if (new_len > std.math.maxInt(SmalInt)) {
                self.largeHeader().len = new_len;
            } else {
                self.small_len = @intCast(SmallInt, new_len);
            }
        }
    };
}