89 lines
2.8 KiB
Zig
89 lines
2.8 KiB
Zig
const std = @import("std");
|
|
const assert = std.debug.assert;
|
|
|
|
fn ShardedTable(comptime Key: type, comptime mask_bit_count: comptime_int, comptime V: type) type {
|
|
assert(Key == @IntType(false, Key.bit_count));
|
|
assert(Key.bit_count >= mask_bit_count);
|
|
const ShardKey = @IntType(false, mask_bit_count);
|
|
const shift_amount = Key.bit_count - ShardKey.bit_count;
|
|
return struct {
|
|
const Self = @This();
|
|
shards: [1 << ShardKey.bit_count]?*Node,
|
|
|
|
pub fn create() Self {
|
|
return Self{ .shards = []?*Node{null} ** (1 << ShardKey.bit_count) };
|
|
}
|
|
|
|
fn getShardKey(key: Key) ShardKey {
|
|
// https://github.com/ziglang/zig/issues/1544
|
|
// this special case is needed because you can't u32 >> 32.
|
|
if (ShardKey == u0) return 0;
|
|
|
|
// this can be u1 >> u0
|
|
const shard_key = key >> shift_amount;
|
|
|
|
// TODO: https://github.com/ziglang/zig/issues/1544
|
|
// This cast could be implicit if we teach the compiler that
|
|
// u32 >> 30 -> u2
|
|
return @intCast(ShardKey, shard_key);
|
|
}
|
|
|
|
pub fn put(self: *Self, node: *Node) void {
|
|
const shard_key = Self.getShardKey(node.key);
|
|
node.next = self.shards[shard_key];
|
|
self.shards[shard_key] = node;
|
|
}
|
|
|
|
pub fn get(self: *Self, key: Key) ?*Node {
|
|
const shard_key = Self.getShardKey(key);
|
|
var maybe_node = self.shards[shard_key];
|
|
while (maybe_node) |node| : (maybe_node = node.next) {
|
|
if (node.key == key) return node;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
pub const Node = struct {
|
|
key: Key,
|
|
value: V,
|
|
next: ?*Node,
|
|
|
|
pub fn init(self: *Node, key: Key, value: V) void {
|
|
self.key = key;
|
|
self.value = value;
|
|
self.next = null;
|
|
}
|
|
};
|
|
};
|
|
}
|
|
|
|
test "sharded table" {
|
|
// realistic 16-way sharding
|
|
testShardedTable(u32, 4, 8);
|
|
|
|
testShardedTable(u5, 0, 32); // ShardKey == u0
|
|
testShardedTable(u5, 2, 32);
|
|
testShardedTable(u5, 5, 32);
|
|
|
|
testShardedTable(u1, 0, 2);
|
|
testShardedTable(u1, 1, 2); // this does u1 >> u0
|
|
|
|
testShardedTable(u0, 0, 1);
|
|
}
|
|
fn testShardedTable(comptime Key: type, comptime mask_bit_count: comptime_int, comptime node_count: comptime_int) void {
|
|
const Table = ShardedTable(Key, mask_bit_count, void);
|
|
|
|
var table = Table.create();
|
|
var node_buffer: [node_count]Table.Node = undefined;
|
|
for (node_buffer) |*node, i| {
|
|
const key = @intCast(Key, i);
|
|
assert(table.get(key) == null);
|
|
node.init(key, {});
|
|
table.put(node);
|
|
}
|
|
|
|
for (node_buffer) |*node, i| {
|
|
assert(table.get(@intCast(Key, i)) == node);
|
|
}
|
|
}
|