tiny.simd.cuckoo2x2

1const std = @import("std");2const hash_mod = @import("hash.zig");3const indexed = @import("indexed.zig");4const random = @import("random.zig");5const shift = @import("shift.zig");6const tag = @import("tag.zig");78pub const minimum_bucket_count: usize = 262_144;9pub const configuration_count: usize = 3;10pub const max_attempts: usize = 512;11pub const max_displacements: usize = 500;1213pub const Cuckoo2x2Config = struct {14    bucket_mask: u32 = 0,15    hash_key: u32 = 0,1617    const Self = @This();1819    pub fn init(num_buckets: usize, hash_key: u32) Self {20        std.debug.assert(num_buckets != 0);21        std.debug.assert(num_buckets <= std.math.maxInt(u32));22        std.debug.assert(std.math.isPowerOfTwo(num_buckets));23        return .{24            .bucket_mask = @intCast(num_buckets - 1),25            .hash_key = hash_key,26        };27    }2829    pub fn numBuckets(self: Self) usize {30        return @as(usize, self.bucket_mask) + 1;31    }32};3334pub const Cuckoo2x2Data = struct {35    config: Cuckoo2x2Config = .{},36    entries: []const u32 = &.{},37    config_idx: usize = 0,38    attempt_idx: usize = 0,39    num_primary: u32 = 0,40    num_secondary: u32 = 0,4142    const Self = @This();4344    pub fn isEmpty(self: Self) bool {45        return self.entries.len == 0;46    }4748    pub fn numBuckets(self: Self) usize {49        if (self.isEmpty()) return 0;50        std.debug.assert(self.entries.len == self.config.numBuckets());51        return self.entries.len;52    }5354    pub fn allocatedBytes(self: Self) usize {55        return self.entries.len * @sizeOf(u32);56    }57};5859pub const Cuckoo2x2 = struct {60    hash1: hash_mod.WeakTwoMul,61    data: Cuckoo2x2Data,6263    const Self = @This();6465    pub fn init(data: Cuckoo2x2Data) Self {66        std.debug.assert(!data.isEmpty());67        std.debug.assert(data.numBuckets() >= minimum_bucket_count);68        return .{69            .hash1 = hash_mod.WeakTwoMul.initKey(data.config.hash_key),70            .data = data,71        };72    }7374    pub fn getData(self: Self) Cuckoo2x2Data {75        return self.data;76    }7778    pub fn contains(self: Self, key: u32) bool {79        const h1 = self.hash1.hash(key);80        const b1 = h1 & self.data.config.bucket_mask;81        const fingerprint = h1 >> 18;82        const b2 = b1 ^ (fingerprint + 1);83        const primary: u16 = @truncate(fingerprint | 0x4000);84        const secondary: u16 = @truncate(fingerprint | 0x8000);85        const e1 = self.data.entries[b1];86        const e2 = self.data.entries[b2];87        return primary == @as(u16, @truncate(e1)) or88            primary == @as(u16, @truncate(e1 >> 16)) or89            secondary == @as(u16, @truncate(e2)) or90            secondary == @as(u16, @truncate(e2 >> 16));91    }9293    pub fn query(self: Self, comptime D: type, keys: D.Vector) D.Mask {94        requireTag(D);95        const h1 = self.hash1.oneVec(D, keys);96        const b1 = h1 & @as(D.Vector, @splat(self.data.config.bucket_mask));97        const fingerprint = shift.shiftRight(D, 18, h1);98        const b2 = b1 ^ (fingerprint +% @as(D.Vector, @splat(1)));99        const primary = fingerprint | @as(D.Vector, @splat(0x4000));100        const secondary = fingerprint | @as(D.Vector, @splat(0x8000));101        const e1 = indexed.gatherIndex(D, self.data.entries, b1);102        const e2 = indexed.gatherIndex(D, self.data.entries, b2);103        const low_mask: D.Vector = @splat(0xffff);104        const found = (e1 & low_mask == primary) |105            (shift.shiftRight(D, 16, e1) == primary) |106            (e2 & low_mask == secondary) |107            (shift.shiftRight(D, 16, e2) == secondary);108        return !found;109    }110111    pub fn containsVec(self: Self, comptime D: type, keys: D.Vector) D.Mask {112        return !self.query(D, keys);113    }114};115116pub const Cuckoo2x2BuildError = error{117    BuildFailed,118    CapacityExceeded,119    DuplicateKey,120    EntriesTooSmall,121    PlanMismatch,122    ScratchTooSmall,123};124125pub const Cuckoo2x2Plan = struct {126    key_count: usize,127    bucket_counts: [configuration_count]usize,128    max_bucket_count: usize,129    entries_len: usize,130    scratch_len: usize,131132    const Self = @This();133134    pub fn inspect(key_count: usize) Cuckoo2x2BuildError!Self {135        if (key_count > std.math.maxInt(u32)) return error.CapacityExceeded;136        const quarter = key_count / 4;137        const base = if (quarter <= 1)138            @as(usize, 1)139        else140            std.math.ceilPowerOfTwo(usize, quarter) catch return error.CapacityExceeded;141        var bucket_counts: [configuration_count]usize = undefined;142        inline for (0..configuration_count) |index| {143            const multiplier = @as(usize, 1) << index;144            const candidate = std.math.mul(usize, base, multiplier) catch145                return error.CapacityExceeded;146            const num_buckets = @max(candidate, minimum_bucket_count);147            if (num_buckets > std.math.maxInt(u32)) return error.CapacityExceeded;148            bucket_counts[index] = num_buckets;149        }150        const max_bucket_count = bucket_counts[configuration_count - 1];151        const bucket_key_count = std.math.mul(usize, max_bucket_count, 2) catch152            return error.CapacityExceeded;153        const byte_count = std.math.add(usize, key_count, max_bucket_count) catch154            return error.CapacityExceeded;155        const byte_words = std.math.divCeil(usize, byte_count, @sizeOf(u32)) catch156            unreachable;157        const hash_and_keys = std.math.add(usize, key_count, bucket_key_count) catch158            return error.CapacityExceeded;159        const scratch_len = std.math.add(usize, hash_and_keys, byte_words) catch160            return error.CapacityExceeded;161        return .{162            .key_count = key_count,163            .bucket_counts = bucket_counts,164            .max_bucket_count = max_bucket_count,165            .entries_len = max_bucket_count,166            .scratch_len = scratch_len,167        };168    }169170    pub fn build(171        self: Self,172        scratch: []u32,173        entries: []u32,174        keys: []const u32,175    ) Cuckoo2x2BuildError!Cuckoo2x2Data {176        if (keys.len != self.key_count) return error.PlanMismatch;177        if (scratch.len < self.scratch_len) return error.ScratchTooSmall;178        if (entries.len < self.entries_len) return error.EntriesTooSmall;179        const parts = self.sections(scratch);180        const engine = random.AesCtrEngine.initDeterministic();181        const distinct_hash = hash_mod.WeakTwoMul.initSeed(&engine, 0);182        hash_mod.hashArray(tag.ScalableTag(u32), distinct_hash, keys, parts.hashes);183        std.mem.sort(u32, parts.hashes, {}, std.sort.asc(u32));184        const sorted_hashes = parts.hashes;185        if (sorted_hashes.len > 1) {186            const previous_hashes = sorted_hashes[0 .. sorted_hashes.len - 1];187            for (sorted_hashes[1..], previous_hashes) |current, previous| {188                if (current == previous) return error.DuplicateKey;189            }190        }191192        for (self.bucket_counts, 0..) |num_buckets, config_idx| {193            for (0..max_attempts) |attempt_idx| {194                const hash_key: u32 = @truncate(engine.generate(attempt_idx, 0));195                const config = Cuckoo2x2Config.init(num_buckets, hash_key);196                const hash1 = hash_mod.WeakTwoMul.initKey(hash_key);197                @memset(parts.counts[0..num_buckets], 0);198                @memset(entries[0..num_buckets], 0);199                hash_mod.hashArray(tag.ScalableTag(u32), hash1, keys, parts.hashes);200                if (!cuckooAssign(201                    config,202                    parts.hashes,203                    parts.choices,204                    parts.counts[0..num_buckets],205                    parts.bucket_keys[0 .. num_buckets * 2],206                )) continue;207                populateEntries(208                    config,209                    parts.hashes,210                    parts.choices,211                    parts.counts[0..num_buckets],212                    entries[0..num_buckets],213                );214                var num_primary: u32 = 0;215                for (parts.choices) |choice| num_primary += @intFromBool(choice == 0);216                return .{217                    .config = config,218                    .entries = entries[0..num_buckets],219                    .config_idx = config_idx,220                    .attempt_idx = attempt_idx,221                    .num_primary = num_primary,222                    .num_secondary = @as(u32, @intCast(keys.len)) - num_primary,223                };224            }225        }226        return error.BuildFailed;227    }228229    fn sections(self: Self, scratch: []u32) ScratchSections {230        const hashes = scratch[0..self.key_count];231        const bucket_keys_begin = self.key_count;232        const bucket_keys_end = bucket_keys_begin + self.max_bucket_count * 2;233        const bucket_keys = scratch[bucket_keys_begin..bucket_keys_end];234        const bytes = std.mem.sliceAsBytes(scratch[bucket_keys_end..self.scratch_len]);235        const choices = bytes[0..self.key_count];236        const counts = bytes[self.key_count .. self.key_count + self.max_bucket_count];237        return .{238            .hashes = hashes,239            .bucket_keys = bucket_keys,240            .choices = choices,241            .counts = counts,242        };243    }244};245246pub fn cuckoo2x2ScratchLen(key_count: usize) Cuckoo2x2BuildError!usize {247    return (try Cuckoo2x2Plan.inspect(key_count)).scratch_len;248}249250pub fn cuckoo2x2EntriesLen(key_count: usize) Cuckoo2x2BuildError!usize {251    return (try Cuckoo2x2Plan.inspect(key_count)).entries_len;252}253254pub fn buildCuckoo2x2(255    scratch: []u32,256    entries: []u32,257    keys: []const u32,258) Cuckoo2x2BuildError!Cuckoo2x2Data {259    const plan = try Cuckoo2x2Plan.inspect(keys.len);260    return plan.build(scratch, entries, keys);261}262263pub fn makeCuckoo2x2(264    scratch: []u32,265    entries: []u32,266    keys: []const u32,267) Cuckoo2x2BuildError!Cuckoo2x2 {268    return Cuckoo2x2.init(try buildCuckoo2x2(scratch, entries, keys));269}270271const ScratchSections = struct {272    hashes: []u32,273    bucket_keys: []u32,274    choices: []u8,275    counts: []u8,276};277278fn bucket(config: Cuckoo2x2Config, hashes: []const u32, key_idx: usize, which: u8) usize {279    const value = hashes[key_idx];280    const primary = value & config.bucket_mask;281    if (which == 0) return primary;282    return primary ^ ((value >> 18) + 1);283}284285fn cuckooAssign(286    config: Cuckoo2x2Config,287    hashes: []const u32,288    choices: []u8,289    counts: []u8,290    bucket_keys: []u32,291) bool {292    std.debug.assert(hashes.len == choices.len);293    std.debug.assert(counts.len == config.numBuckets());294    std.debug.assert(bucket_keys.len == counts.len * 2);295    @memset(bucket_keys, std.math.maxInt(u32));296    for (0..hashes.len) |key_idx| {297        var displaced = key_idx;298        const primary = bucket(config, hashes, displaced, 0);299        const secondary = bucket(config, hashes, displaced, 1);300        if (counts[primary] <= counts[secondary] and counts[primary] < 2) {301            choices[displaced] = 0;302            bucket_keys[primary * 2 + counts[primary]] = @intCast(displaced);303            counts[primary] += 1;304            continue;305        }306        if (counts[secondary] < 2) {307            choices[displaced] = 1;308            bucket_keys[secondary * 2 + counts[secondary]] = @intCast(displaced);309            counts[secondary] += 1;310            continue;311        }312        if (counts[primary] < 2) {313            choices[displaced] = 0;314            bucket_keys[primary * 2 + counts[primary]] = @intCast(displaced);315            counts[primary] += 1;316            continue;317        }318319        var which: u8 = 0;320        var placed = false;321        for (0..max_displacements) |_| {322            const destination = bucket(config, hashes, displaced, which);323            if (counts[destination] < 2) {324                choices[displaced] = which;325                bucket_keys[destination * 2 + counts[destination]] = @intCast(displaced);326                counts[destination] += 1;327                placed = true;328                break;329            }330            const evicted = bucket_keys[destination * 2];331            bucket_keys[destination * 2] = bucket_keys[destination * 2 + 1];332            bucket_keys[destination * 2 + 1] = @intCast(displaced);333            choices[displaced] = which;334            displaced = evicted;335            which = 1 - choices[evicted];336        }337        if (!placed) return false;338    }339    return true;340}341342fn populateEntries(343    config: Cuckoo2x2Config,344    hashes: []const u32,345    choices: []const u8,346    counts: []u8,347    entries: []u32,348) void {349    std.debug.assert(hashes.len == choices.len);350    std.debug.assert(counts.len == entries.len);351    @memset(counts, 0);352    for (hashes, choices) |value, choice| {353        const primary = value & config.bucket_mask;354        const fingerprint = value >> 18;355        const destination = if (choice == 0) primary else primary ^ (fingerprint + 1);356        const tagged = fingerprint | if (choice == 0) @as(u32, 0x4000) else 0x8000;357        const slot = counts[destination];358        std.debug.assert(slot < 2);359        entries[destination] |= tagged << @intCast(slot * 16);360        counts[destination] += 1;361    }362    for (entries, counts) |*entry, count| {363        if (count == 1) entry.* |= entry.* << 16;364    }365}366367fn requireTag(comptime D: type) void {368    if (comptime D.Lane != u32) @compileError("Cuckoo2x2 requires u32 lanes");369}370371test "Highway Cuckoo2x2 scalar and vector membership agree" {372    const key_count = 1000;373    var keys: [key_count]u32 = undefined;374    for (&keys, 0..) |*key, index| key.* = @as(u32, @intCast(index)) *% 37 +% 1;375    const plan = try Cuckoo2x2Plan.inspect(keys.len);376    const allocator = std.testing.allocator;377    const scratch = try allocator.alloc(u32, plan.scratch_len);378    defer allocator.free(scratch);379    const entries = try allocator.alloc(u32, plan.entries_len);380    defer allocator.free(entries);381    const set = Cuckoo2x2.init(try plan.build(scratch, entries, &keys));382    for (keys) |key| try std.testing.expect(set.contains(key));383    const D = tag.FixedTag(u32, 8);384    var index: usize = 0;385    while (index + D.lane_count <= keys.len) : (index += D.lane_count) {386        const input: D.Vector = keys[index..][0..D.lane_count].*;387        try std.testing.expect(@reduce(.And, set.containsVec(D, input)));388    }389    for (0..key_count) |absent_idx| {390        const absent = @as(u32, @intCast(absent_idx)) *% 37 +% 2;391        try std.testing.expect(!set.contains(absent));392    }393}394395test "Highway Cuckoo2x2 single-worker builder oracle matches exactly" {396    const key_count = 1024;397    var keys: [key_count]u32 = undefined;398    for (&keys, 0..) |*key, index| key.* = @as(u32, @intCast(index)) *% 37 +% 1;399    const plan = try Cuckoo2x2Plan.inspect(keys.len);400    const allocator = std.testing.allocator;401    const scratch = try allocator.alloc(u32, plan.scratch_len);402    defer allocator.free(scratch);403    const entries = try allocator.alloc(u32, plan.entries_len);404    defer allocator.free(entries);405    const data = try plan.build(scratch, entries, &keys);406    try std.testing.expectEqual(@as(usize, 262_144), data.numBuckets());407    try std.testing.expectEqual(@as(u32, 0x7aa5_332d), data.config.hash_key);408    try std.testing.expectEqual(@as(usize, 0), data.config_idx);409    try std.testing.expectEqual(@as(usize, 0), data.attempt_idx);410    try std.testing.expectEqual(@as(u32, 1023), data.num_primary);411    try std.testing.expectEqual(@as(u32, 1), data.num_secondary);412    const expected_indices = [16]usize{413        90,   153,  321,  423,  429,  636,  747,  1234,414        1648, 2094, 2308, 2564, 2586, 2598, 3034, 4066,415    };416    const expected_entries = [16]u32{417        0x6dd4_6dd4,418        0x5894_5894,419        0x686c_686c,420        0x5272_5272,421        0x454f_454f,422        0x6a0e_6a0e,423        0x4a8e_4a8e,424        0x745b_745b,425        0x4499_4499,426        0x503c_503c,427        0x4800_4800,428        0x57f3_57f3,429        0x6dd1_6dd1,430        0x7c93_7c93,431        0x544e_544e,432        0x5634_5634,433    };434    for (expected_indices, expected_entries) |index, expected| {435        try std.testing.expectEqual(expected, data.entries[index]);436    }437    var digest: u64 = 1_469_598_103_934_665_603;438    var nonzero: usize = 0;439    for (data.entries) |entry| {440        digest = (digest ^ entry) *% 1_099_511_628_211;441        nonzero += @intFromBool(entry != 0);442    }443    try std.testing.expectEqual(@as(u64, 0x947b_8ca0_29af_1476), digest);444    try std.testing.expectEqual(key_count, nonzero);445}446447test "Highway Cuckoo2x2 builds and queries upstream sizes" {448    const allocator = std.testing.allocator;449    for ([_]usize{ 6000, 60_000 }) |key_count| {450        const keys = try allocator.alloc(u32, key_count);451        defer allocator.free(keys);452        hash_mod.fillRandomDistinct(u32, 0, keys);453        const plan = try Cuckoo2x2Plan.inspect(keys.len);454        const scratch = try allocator.alloc(u32, plan.scratch_len);455        defer allocator.free(scratch);456        const entries = try allocator.alloc(u32, plan.entries_len);457        defer allocator.free(entries);458        const data = try plan.build(scratch, entries, keys);459        const set = Cuckoo2x2.init(data);460        try std.testing.expectEqual(key_count, data.num_primary + data.num_secondary);461        try std.testing.expect(data.numBuckets() >= minimum_bucket_count);462        for (keys) |key| try std.testing.expect(set.contains(key));463    }464}465466test "Highway Cuckoo2x2 plans enforce capacities and distinct keys" {467    const keys = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8 };468    const plan = try Cuckoo2x2Plan.inspect(keys.len);469    const allocator = std.testing.allocator;470    const scratch = try allocator.alloc(u32, plan.scratch_len);471    defer allocator.free(scratch);472    const entries = try allocator.alloc(u32, plan.entries_len);473    defer allocator.free(entries);474    try std.testing.expectError(475        error.ScratchTooSmall,476        plan.build(scratch[0 .. scratch.len - 1], entries, &keys),477    );478    try std.testing.expectError(479        error.EntriesTooSmall,480        plan.build(scratch, entries[0 .. entries.len - 1], &keys),481    );482    try std.testing.expectError(error.PlanMismatch, plan.build(scratch, entries, keys[0..7]));483    const duplicate = [_]u32{ 1, 2, 3, 3, 5, 6, 7, 8 };484    try std.testing.expectError(error.DuplicateKey, plan.build(scratch, entries, &duplicate));485}486487test "Highway Cuckoo2x2 represents the empty set" {488    const plan = try Cuckoo2x2Plan.inspect(0);489    const allocator = std.testing.allocator;490    const scratch = try allocator.alloc(u32, plan.scratch_len);491    defer allocator.free(scratch);492    const entries = try allocator.alloc(u32, plan.entries_len);493    defer allocator.free(entries);494    const data = try plan.build(scratch, entries, &.{});495    const set = Cuckoo2x2.init(data);496    try std.testing.expect(!set.contains(0));497    try std.testing.expect(!set.contains(std.math.maxInt(u32)));498    try std.testing.expectEqual(@as(u32, 0), data.num_primary);499    try std.testing.expectEqual(@as(u32, 0), data.num_secondary);500}

1pub const cuckoo2x2 = @import("cuckoo2x2.zig");