Compiler Explorer

Source code

const std = @import("std");
const mem = std.mem;
const math = std.math;

const warn = std.debug.warn;
const assert = std.debug.assert;

const ArrayList = std.ArrayList;

fn Bucket(comptime T: type, comptime items_per_bucket: i32) type {
    return struct {
        const Self = @This();

data: [items_per_bucket]T,
        occupied: [items_per_bucket]bool,
        index: i32,
        count: i32,
    };
}

const BucketLocator = struct {
    bucket_index: i32,
    slot_index: usize,
};

fn indexOf(comptime T: type, haystack: var, needle: T) ?usize {
    var itr = haystack.iterator();
    var i: usize = 0;
    while (itr.next()) |e| : (i += 1) {
        if (e == needle)  return i;
    }
    return null;
}

// TODO: Add shrink to free empty buckets.
fn BucketArray(comptime T: type, comptime items_per_bucket: i32) type {
    return struct {
        const MyBucket = Bucket(T, items_per_bucket);
        const Self = @This();

allocator: *mem.Allocator,
        count: i64,

all_buckets: ArrayList(*MyBucket),
        unfull_buckets: ArrayList(*MyBucket),

pub fn init(allocator: *mem.Allocator) Self {
            var a = Self {
                .all_buckets = ArrayList(*MyBucket).init(allocator),
                .unfull_buckets = ArrayList(*MyBucket).init(allocator),
                .count = 0,
                .allocator = allocator,
            };
            return a;
        }

pub fn deinit(self: *Self) void {
            var itr = self.all_buckets.iterator();
            while (itr.next()) |bucket| {
                self.allocator.destroy(bucket);
            }
            self.unfull_buckets.shrink(0);
            self.all_buckets.shrink(0);
        }

fn addBucket(self: *Self) error{OutOfMemory}!void {
            var b = try self.allocator.create(MyBucket);
            b.index = @intCast(i32, self.all_buckets.len);
            b.count = 0;
            assert(b.index == @intCast(i32, self.all_buckets.len)); // NOTE: Asserts if overflowed.

try self.all_buckets.append(b);
            try self.unfull_buckets.append(b);
        }

pub fn append(self: *Self, item: T) error{OutOfMemory}!BucketLocator {
            if (self.unfull_buckets.len == 0)  try addBucket(self);
            assert(self.unfull_buckets.len != 0);

bucket.occupied[index.?] = true;
            bucket.count += 1;
            assert(bucket.count <= items_per_bucket);
            
            self.count += 1;

if (bucket.count == items_per_bucket) {
                var itr = self.unfull_buckets.iterator();
                var i: usize = 0;
                var removed = false;
                while (itr.next()) |value| : (i += 1) {
                    if (value == bucket) {
                        _ = self.unfull_buckets.swapRemove(i);
                        removed = true;
                        break;
                    }
                }
                assert(removed);
            }

var loc = BucketLocator{
                .bucket_index = bucket.index,
                .slot_index = index.?,
            };

bucket.data[index.?] = item;
            return loc;
        }

pub fn remove(self: *Self, loc: BucketLocator) !T {
            var bucket = self.all_buckets.at(@intCast(usize, loc.bucket_index));
            assert(bucket.occupied[loc.slot_index] == true);

const wasFull = (bucket.count == items_per_bucket);

bucket.occupied[loc.slot_index] = false;
            bucket.count -= 1;
            self.count -= 1;

if (wasFull) {
                assert(indexOf(*MyBucket, self.unfull_buckets, bucket) == null);
                try self.unfull_buckets.append(bucket);
            }

return bucket.data[loc.slot_index];
        }

pub fn at(self: *Self, loc: BucketLocator) T {
            var bucket = self.all_buckets.at(@intCast(usize, loc.bucket_index));
            assert(bucket.occupied[loc.slot_index] == true);
            var result = bucket.data[@intCast(usize, loc.slot_index)];
            return result;
        }

pub fn iterator(self: *Self) BucketIterator {
            return BucketIterator{
                .array = self,
                .locator = BucketLocator{.bucket_index = 0, .slot_index = 0},
                .index = -1,
            };
        }

const BucketIterator = struct {
            array: *Self,
            index: i64,
            locator: BucketLocator,

fn next(self: *BucketIterator) ?T {
                if (self.index >= self.array.count)  return null;

self.index += 1;
                var loc = self.locator;

if (self.index > 0) {
                    while (true) {
                        loc.slot_index += 1;
                        if (loc.slot_index >= @intCast(usize, items_per_bucket)) {
                            loc.slot_index = 0;
                            loc.bucket_index += 1;
                        }

if (@intCast(usize, loc.bucket_index) >= self.array.all_buckets.len)  return null;

var bucket = self.array.all_buckets.at(@intCast(usize, loc.bucket_index));
                        if (bucket.occupied[@intCast(usize, loc.slot_index)])  break;
                    }
                    self.locator = loc;
                }

return self.array.all_buckets.at(@intCast(usize, loc.bucket_index)).data[@intCast(usize, loc.slot_index)];
            }
        };
    };
}

pub fn main() !void {
    var arena = std.heap.ArenaAllocator.init(&std.heap.DirectAllocator.init().allocator);
    defer arena.deinit();

{
        var itr = ba.iterator();
        var j: i32 = 0;
        warn("{{ ");
        while (itr.next()) |value| : (j += 1) {
            if (j > 0) warn(", ");
            warn("{}", value);
            assert(value == math.pow(i32, 2, j));
        }
        warn(" }}\n");
    }

{
        const removed = try ba.remove(BucketLocator{.bucket_index = 0, .slot_index = 2});
        assert(removed == 4);
        assert(ba.count == countAfterAdding-1);
    }

{
        var itr = ba.iterator();
        var i: i32 = 0;
        warn("{{ ");
        while (itr.next()) |value| : (i += 1) {
            if (i > 0) warn(", ");
            warn("{}", value);

assert(value != 4); // Since we removed it.
        }
        warn(" }}\n");
    }
}