zig 6.0 hashmap + smap benchmark

main.zig

const std = @import("std");
const smap = @import("smap.zig");
const AutoHashMap = std.AutoHashMap;
const time = std.time;

fn step(x: u64) u64 {
    if (x & 1 > 0) {
        return 3 * x + 1;
    } else {
        return x / 2;
    }
}

fn length(comptime cachetype: type, cache: *cachetype, x: u64) anyerror!u64 {
    if (x <= 1) return 0;
    if (cache.getValue(x)) |e| {
        return e;
    } else {
        const next = step(x);
        const len = 1 + try length(cachetype, cache, next);
        _ = try cache.put(x, len);
        return len;
    }
}

fn bench(comptime cachetype: type, n: u64) anyerror!void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;

    var cache = cachetype.init(alloc);
    defer cache.deinit();
    try cache.ensureCapacity(n * 2 + n / 100 * 5);

    var timer = try time.Timer.start();

    var x: u64 = 0;
    var maxx: u64 = 0;
    var maxl: u64 = 0;
    while (x < n) : (x += 1) {
        const l = try length(cachetype, &cache, x);
        if (l > maxl) {
            maxl = l;
            maxx = x;
        }
    }

    const nanosecs = timer.read();

    std.debug.warn("{s: >101} {}ms\n", .{ @typeName(cachetype), nanosecs / 1000000 });
}

pub fn main() anyerror!void {
    const groupsizes = [_]comptime_int{ 1, 2, 4, 8, 16, 32, 64, 128 };
    const hashfn = comptime smap.getAutoHashFn(u64);
    const eqlfn = comptime smap.getAutoEqlFn(u64);

    inline for (groupsizes) |groupsize| {
        const cachetype = smap.Map(u64, u64, hashfn, eqlfn, groupsize);
        try bench(cachetype, 1000000);
    }
    {
        const cachetype = AutoHashMap(u64, u64);
        try bench(cachetype, 1000000);
    }
}

smap.zig

const std = @import("std");
const math = std.math;
const assert = std.debug.assert;
const Allocator = std.mem.Allocator;
const testing = std.testing;
const Wyhash = std.hash.Wyhash;
const autoHash = std.hash.autoHash;
const meta = std.meta;

pub fn Map(comptime K: type, comptime V: type, comptime hash: fn (key: K) u32, comptime eql: fn (a: K, b: K) bool, comptime groupsize: u32) type {
    return struct {
        meta: []Meta,
        entries: [][groupsize]KV,
        size: usize,
        allocator: *Allocator,

        pub const KV = struct {
            key: K,
            val: V,
        };

        const Self = @This();
        const Meta = @Vector(groupsize, i8);
        const bitmaskType = @Type(std.builtin.TypeInfo{ .Int = .{ .is_signed = false, .bits = groupsize } });

        const Ctrl = enum(i8) {
            Empty = -128,
            Deleted = -2,
            Sentinel = -1,
            _,
        };

        pub fn init(allocator: *Allocator) Self {
            return Self{
                .meta = &[_]Meta{},
                .entries = &[_][groupsize]KV{},
                .size = 0,
                .allocator = allocator,
            };
        }
        pub fn deinit(self: Self) void {
            self.allocator.free(self.meta);
            self.allocator.free(self.entries);
        }
        pub fn count(self: Self) usize {
            return self.size;
        }

        fn initCapacity(self: *Self, capacity: u32) !void {
            assert(math.isPowerOfTwo(capacity));
            self.meta = try self.allocator.alloc(Meta, capacity);
            self.entries = try self.allocator.alloc([groupsize]KV, capacity);

            for (self.meta) |*g| {
                g.* = mSplat(Ctrl.Empty);
            }
        }

        inline fn mSplat(c: Ctrl) Meta {
            return @splat(groupsize, @enumToInt(c));
        }
        inline fn mSplatInt(c: i8) Meta {
            return @splat(groupsize, c);
        }
        inline fn bitmask(v: @Vector(groupsize, bool)) bitmaskType {
            return @ptrCast(*const bitmaskType, &v).*;
        }
        const bitmaskIterType = struct {
            b: bitmaskType,
            inline fn next(self: *bitmaskIterType) ?usize {
                if (self.b == 0) return null;

                const ret: usize = @ctz(bitmaskType, self.b);
                self.b = self.b & (self.b - 1);
                return ret;
            }
        };
        inline fn bitmaskIter(b: bitmaskType) bitmaskIterType {
            return bitmaskIterType{ .b = b };
        }

        inline fn h1(h: u32) u32 {
            return h >> 7;
        }
        inline fn h2(h: u32) i8 {
            return @as(i8, @truncate(u7, h));
        }
        inline fn groupIdx(self: *Self, h: u32) u32 {
            return h1(h) % @truncate(u32, self.meta.len);
        }
        const groupIter = struct {
            idx: u32,
            i: u32,
            n: usize,

            inline fn next(self: *groupIter) ?u32 {
                if (self.i == self.n) {
                    return null;
                } else {
                    self.idx = (self.idx + self.i) % @truncate(u32, self.n);
                    self.i += 1;
                    return self.idx;
                }
            }
        };
        fn newGroupIter(self: *Self, h: u32) groupIter {
            return groupIter{ .idx = self.groupIdx(h), .i = 0, .n = self.meta.len };
        }

        const Slot = struct {
            group: usize,
            bit: usize,
        };

        fn getSlot(self: *Self, k: K, h: u32) ?Slot {
            if (self.entries.len == 0) return null;
            var grpiter = self.newGroupIter(h);
            while (grpiter.next()) |g| {
                const m = self.meta[g];
                var candidates = bitmaskIter(bitmask(m == mSplatInt(h2(h))));
                while (candidates.next()) |bit| {
                    if (eql(self.entries[g][bit].key, k)) {
                        return Slot{ .group = g, .bit = bit };
                    }
                }
                const empties = bitmask(m == mSplat(Ctrl.Empty));
                if (empties > 0) {
                    return null;
                }
            }
            return null;
        }
        fn getInternal(self: *Self, k: K, h: u32) ?*KV {
            if (self.getSlot(k, h)) |s| {
                return &self.entries[s.group][s.bit];
            }
            return null;
        }
        pub fn get(self: *Self, k: K) ?*KV {
            return self.getInternal(k, hash(k));
        }
        pub fn contains(self: *Self, k: K) bool {
            return self.getSlot(k, hash(k)) != null;
        }
        pub fn getValue(self: *Self, k: K) ?V {
            if (self.get(k)) |kv| {
                return kv.val;
            }
            return null;
        }
        fn remove(self: *Self, k: K) ?KV {
            if (self.getSlot(k, hash(k))) |s| {
                const kv = self.entries[s.group][s.bit];
                self.entries[s.group][s.bit] = undefined;
                self.meta[s.group][s.bit] = @enumToInt(Ctrl.Deleted);
                self.size -= 1;
                self.autoCapacity() catch unreachable;
                return kv;
            }
            return null;
        }
        pub fn putComposed(self: *Self, k: K, v: V) !?KV {
            const h = hash(k);
            if (self.getInternal(k, h)) |kv| {
                const ret = kv.*;
                kv.val = v;
                return ret;
            }

            try self.ensureCapacity(self.size + 1);
            self.putInternal(k, v, h);
            return null;
        }
        pub fn putAssumeCapacity(self: *Self, k: K, v: V) void {
            self.putInternal(k, v, hash(k));
        }
        pub fn put(self: *Self, k: K, v: V) !?KV {
            try self.ensureCapacity(self.size + 1);
            const h = hash(k);

            var grpiter = self.newGroupIter(h);
            outer: while (grpiter.next()) |g| {
                var m = &self.meta[g];
                var candidates = bitmaskIter(bitmask(m.* == mSplatInt(h2(h))));
                while (candidates.next()) |bit| {
                    const kv = &self.entries[g][bit];
                    if (eql(kv.key, k)) {
                        const ret = kv.*;
                        kv.key = k;
                        kv.val = v;
                        return ret;
                    }
                }
                candidates = bitmaskIter(bitmask(m.* < mSplatInt(0)));
                while (candidates.next()) |bit| {
                    const e = &self.entries[g][bit];
                    e.key = k;
                    e.val = v;
                    m.*[bit] = h2(h);
                    self.size += 1;
                    break :outer;
                }
                const empties = bitmask(m.* == mSplat(Ctrl.Empty));
                if (empties > 0) {
                    unreachable;
                }
            }
            while (grpiter.next()) |g| {
                const m = self.meta[g];
                var candidates = bitmaskIter(bitmask(m == mSplatInt(h2(h))));
                while (candidates.next()) |bit| {
                    const kv = &self.entries[g][bit];
                    if (eql(kv.key, k)) {
                        const ret = kv.*;
                        kv.* = undefined;
                        self.meta[g][bit] = @enumToInt(Ctrl.Deleted);
                        return ret;
                    }
                }
                const empties = bitmask(m == mSplat(Ctrl.Empty));
                if (empties > 0) {
                    return null;
                }
            }
            return null;
            //unreachable;
        }

        fn putInternal(self: *Self, k: K, v: V, h: u32) void {
            var grpiter = self.newGroupIter(h);
            while (grpiter.next()) |g| {
                var m = &self.meta[g];
                const b = bitmask(m.* < mSplatInt(0));
                if (b == 0) continue;

                var candidates = bitmaskIter(b);
                while (candidates.next()) |bit| {
                    const e = &self.entries[g][bit];
                    e.key = k;
                    e.val = v;
                    m.*[bit] = h2(h);
                    self.size += 1;
                    return;
                }
            }
            unreachable;
        }
        pub inline fn ensureCapacity(self: *Self, expected_count: usize) !void {
            const optcap = optimizedCapacity(expected_count);
            if (optcap > self.entries.len) {
                try self.resize(optcap);
            }
        }
        inline fn autoCapacity(self: *Self) !void {
            const optcap = optimizedCapacity(self.size);
            if (optcap != self.entries.len) {
                try self.resize(optcap);
            }
        }
        inline fn optimizedCapacity(size: usize) usize {
            var nextsize = math.ceilPowerOfTwo(usize, math.max(groupsize, size * 8 / 7)) catch unreachable;
            return nextsize / groupsize;
        }

        pub const Iterator = struct {
            m: *Self,
            g: usize,
            b: bitmaskIterType,

            pub inline fn next(it: *Iterator) ?*KV {
                while (true) {
                    if (it.b.next()) |bit| {
                        return &it.m.entries[it.g][bit];
                    }
                    it.g += 1;
                    if (it.g >= it.m.meta.len) {
                        return null;
                    }
                    it.b = bitmaskIter(bitmask(it.m.meta[it.g] >= mSplatInt(0)));
                }
                return null;
            }
        };
        pub inline fn iterator(self: *Self) Iterator {
            var b = bitmaskIter(0);
            if (self.meta.len > 0) {
                b = bitmaskIter(bitmask(self.meta[0] >= mSplatInt(0)));
            }
            return Iterator{ .m = self, .g = 0, .b = b };
        }
        fn copySizeExact(self: *Self, newcount: usize) !Self {
            var new = Self.init(self.allocator);
            try new.initCapacity(@truncate(u32, newcount));

            var it = self.iterator();
            var i: usize = 0;
            while (it.next()) |kv| {
                const h = hash(kv.key);
                new.putInternal(kv.key, kv.val, h);
                i += 1;
            }
            assert(new.size == self.size);
            return new;
        }
        fn copy(self: *Self) !Self {
            return self.copySizeExact(optimizedCapacity(self.size));
        }
        fn resize(self: *Self, newcount: usize) !void {
            var new = try self.copySizeExact(newcount);
            var old = self.*;
            self.* = new;
            old.deinit();
        }
    };
}

pub fn AutoMap(comptime K: type, comptime V: type) type {
    return AutoMapSize(K, V, 16);
}
pub fn AutoMapSize(comptime K: type, comptime V: type, comptime groupsize: comptime_int) type {
    return Map(K, V, getAutoHashFn(K), getAutoEqlFn(K), groupsize);
}

test "1024" {
    var m = AutoMap(u32, u32).init(std.testing.allocator);
    const cap = 1024;

    var i: u32 = 0;
    while (i < cap) : (i += 1) {
        const old = try m.put(i, i * 10);
        testing.expectEqual(old, null);
        testing.expectEqual(@as(usize, i + 1), m.size);
    }
    i = 0;
    while (i < cap) : (i += 1) {
        const kv = m.get(i).?;
        testing.expectEqual(kv.key, i);
        testing.expectEqual(kv.val, i * 10);
        testing.expectEqual(@as(usize, cap), m.size);
    }
    i = 0;
    while (i < cap) : (i += 1) {
        const old = try m.put(i, i * 10);
        testing.expectEqual(old.?.key, i);
        testing.expectEqual(old.?.val, i * 10);
        testing.expectEqual(@as(usize, cap), m.size);
    }
    i = 0;
    while (i < cap) : (i += 1) {
        const deleted = m.remove(i);
        testing.expectEqual(deleted.?.key, i);
        testing.expectEqual(deleted.?.val, i * 10);
        testing.expectEqual(@as(usize, cap - i - 1), m.size);
    }

    testing.expectEqual(@as(usize, 0), m.size);

    i = 0;
    while (i < cap) : (i += 1) {
        const old = try m.put(i, i * 10);
        testing.expectEqual(old, null);
        testing.expectEqual(@as(usize, i + 1), m.size);
    }
    i = 0;
    while (i < cap) : (i += 1) {
        testing.expectEqual(m.contains(i), true);
    }
    i = 0;
    while (i < cap) : (i += 1) {
        testing.expectEqual(m.getValue(i), i * 10);
    }
    m.deinit();
    try testing.allocator_instance.validate();
}

pub fn getAutoHashFn(comptime K: type) (fn (K) u32) {
    return struct {
        fn hash(key: K) u32 {
            var hasher = Wyhash.init(0);
            autoHash(&hasher, key);
            return @truncate(u32, hasher.final());
        }
    }.hash;
}

pub fn getAutoEqlFn(comptime K: type) (fn (K, K) bool) {
    return struct {
        fn eql(a: K, b: K) bool {
            return meta.eql(a, b);
        }
    }.eql;
}