const std = @import("std");
const testing = std.testing;
const assert = std.debug.assert;
const Allocator = std.mem.Allocator;
const DirectAllocator = std.heap.DirectAllocator;

fn max(comptime T: type, a: T, b: T) T {
    return if (a > b) a else b;
}

pub fn FlatMap(
    comptime K: type,
    comptime V: type,
    comptime hash: fn (K) u64,
    comptime key_equals: fn (K, K) bool,
) type {
    return struct {
        const Self = @This();

        const KV = struct {
            value: V,
            key: K,
            distance: i8, // -1 == empty, -2 == deleted, otherwise represents probe distance
        };

        slots: []KV,
        size: usize,
        maxLoadFactor: f32,
        longestProbe: i8,
        allocator: *Allocator,

        const EMPTY: i8 = -1;
        const TOMBSTONE: i8 = -1;

        pub fn init(allocator: *Allocator) !Self {
            var new_map = Self{
                .slots = undefined,
                .size = 0,
                .maxLoadFactor = 0.6,
                .longestProbe = 0,
                .allocator = allocator,
            };

            new_map.slots = try allocator.alloc(KV, 2);
            return new_map;
        }

        pub fn deinit(self: *Self) void {
            self.allocator.free(self.slots);
        }

        pub inline fn capacity(self: *const Self) usize {
            return self.slots.len;
        }

        pub inline fn loadFactor(self: *const Self) f32 {
            return @intToFloat(f32, self.size) / @intToFloat(f32, self.capacity());
        }

        pub fn insert(self: *Self, key: K, value: V) !void {
            if (self.loadFactor() > self.maxLoadFactor) {
                try self.reserve(max(usize, 2, self.capacity() * 2));
            }

            const N = self.capacity() - 1;
            var idx: usize = hash(key) & N;

            var entry = KV{ .key = key, .value = value, .distance = 0 };

            while (true) {
                var probed = &self.slots[idx];

                if (probed.distance < 0) { // encountered EMPTY/TOMBSTONE
                    probed.* = entry;
                    self.size += 1;
                    self.longestProbe = max(i8, entry.distance, self.longestProbe);
                    return;
                } else if (key_equals(probed.key, entry.key)) {
                    probed.value = entry.value;
                    return;
                } else if (probed.distance < entry.distance) {
                    self.longestProbe = max(i8, entry.distance, self.longestProbe);
                    std.mem.swap(KV, probed, &entry);
                }

                idx = (idx + 1) & N;
                entry.distance += 1;
            }
        }

        pub fn lookup(self: *Self, lookup_key: K) ?*V {
            const N = self.capacity() - 1;
            var idx: usize = hash(lookup_key) & N;
            var distance: i8 = 0;

            while (true) {
                var probed = &self.slots[idx];

                if (key_equals(probed.key, lookup_key)) {
                    return &probed.value;
                } else if (probed.distance == EMPTY) {
                    return null;
                }

                idx = (idx + 1) & N;
                distance += 1;

                if (distance > self.longestProbe) {
                    return null;
                }
            }
        }

        pub fn contains(self: *Self, key: K) bool {
            return self.lookup(key) != null;
        }

        pub fn remove(self: *Self, key_to_delete: K) bool {
            const N = self.capacity() - 1;
            var idx: usize = hash(key) & N;
            var distance: i8 = 0;

            while (true) {
                var probed = &self.slots[idx];

                if (key_equals(probed.key, key_to_delete)) {
                    probed.distance = TOMBSTONE;
                    self.size -= 1;
                    return true;
                } else if (probed.distance == EMPTY) {
                    return false;
                }

                idx = (idx + 1) & N;
                distance += 1;

                if (distance > self.longestProbe) {
                    return false;
                }
            }
        }

        pub fn reserve(self: *Self, newCapacity: usize) !void {
            assert(newCapacity > self.capacity());
            assert(newCapacity & (newCapacity - 1) == 0); // demand capacity is power of two
            var oldSlots = self.slots;
            self.slots = try self.allocator.alloc(KV, newCapacity);

            for (self.slots) |*entry| {
                entry.distance = -1;
            }

            self.size = 0;
            self.longestProbe = 0;
            for (oldSlots) |entry| {
                if (entry.distance >= 0) {
                    _ = self.insert(entry.key, entry.value);
                }
            }

            self.allocator.free(oldSlots);
        }
    };
}

fn internal_hash(x: u32) u64 {
    var y: u64 = @intCast(u64, x);

    y = (y + 0x7ed55d16) + (y << 12);
    y = (y ^ 0xc761c23c) ^ (y >> 19);
    y = (y + 0x165667b1) + (y << 5);
    y = (y + 0xd3a2646c) ^ (y << 9);
    y = (y + 0xfd7046c5) + (y << 3);
    y = (y ^ 0xb55a4f09) ^ (y >> 16);

    return y;
}

fn internal_equals(x: u32, y: u32) bool {
    return x == y;
}

test "FlatMap" {
    var dalloc = DirectAllocator.init();
    defer dalloc.deinit();

    var map = try FlatMap(u32, u32, internal_hash, internal_equals).init(&dalloc.allocator);
    defer map.deinit();

    var i: u32 = 0;
    while (i < 10000) {
        try map.insert(i, i * 2);
        testing.expect(map.size == i + 1);
        testing.expect(map.contains(i));
        i += 1;
    }

    i = 0;
    while (i < 300) {
        if (map.lookup(i)) |x| {
            testing.expect(x.* == i * 2);
        } else {
            testing.expect(false);
        }
        i += 1;
    }
}