| ... | ... | @@ -0,0 +1,86 @@ |
| 1 | const std = @import("std"); |
| 2 | const assert = std.debug.assert; |
| 3 | |
| 4 | fn ShardedTable(comptime Key: type, comptime mask_bit_count: comptime_int, comptime V: type) type { |
| 5 | assert(Key == @IntType(false, Key.bit_count)); |
| 6 | assert(Key.bit_count >= mask_bit_count); |
| 7 | const ShardKey = @IntType(false, mask_bit_count); |
| 8 | const shift_amount = Key.bit_count - ShardKey.bit_count; |
| 9 | return struct { |
| 10 | const Self = @This(); |
| 11 | shards: [1 << ShardKey.bit_count]?*Node, |
| 12 | |
| 13 | pub fn create() Self { |
| 14 | return Self{ .shards = []?*Node{null} ** (1 << ShardKey.bit_count) }; |
| 15 | } |
| 16 | |
| 17 | fn getShardKey(key: Key) ShardKey { |
| 18 | // this special case is needed because you can't u32 >> 32. |
| 19 | if (ShardKey == u0) return 0; |
| 20 | |
| 21 | // this can be u1 >> u0 |
| 22 | const shard_key = key >> shift_amount; |
| 23 | |
| 24 | // TODO: this cast could be implicit if we teach the compiler that |
| 25 | // u32 >> 30 -> u2 |
| 26 | return @intCast(ShardKey, shard_key); |
| 27 | } |
| 28 | |
| 29 | pub fn put(self: *Self, node: *Node) void { |
| 30 | const shard_key = Self.getShardKey(node.key); |
| 31 | node.next = self.shards[shard_key]; |
| 32 | self.shards[shard_key] = node; |
| 33 | } |
| 34 | |
| 35 | pub fn get(self: *Self, key: Key) ?*Node { |
| 36 | const shard_key = Self.getShardKey(key); |
| 37 | var maybe_node = self.shards[shard_key]; |
| 38 | while (maybe_node) |node| : (maybe_node = node.next) { |
| 39 | if (node.key == key) return node; |
| 40 | } |
| 41 | return null; |
| 42 | } |
| 43 | |
| 44 | pub const Node = struct { |
| 45 | key: Key, |
| 46 | value: V, |
| 47 | next: ?*Node, |
| 48 | |
| 49 | pub fn init(self: *Node, key: Key, value: V) void { |
| 50 | self.key = key; |
| 51 | self.value = value; |
| 52 | self.next = null; |
| 53 | } |
| 54 | }; |
| 55 | }; |
| 56 | } |
| 57 | |
| 58 | test "sharded table" { |
| 59 | // realistic 16-way sharding |
| 60 | testShardedTable(u32, 4, 8); |
| 61 | |
| 62 | testShardedTable(u5, 0, 32); // ShardKey == u0 |
| 63 | testShardedTable(u5, 2, 32); |
| 64 | testShardedTable(u5, 5, 32); |
| 65 | |
| 66 | testShardedTable(u1, 0, 2); |
| 67 | testShardedTable(u1, 1, 2); // this does u1 >> u0 |
| 68 | |
| 69 | testShardedTable(u0, 0, 1); |
| 70 | } |
| 71 | fn testShardedTable(comptime Key: type, comptime mask_bit_count: comptime_int, comptime node_count: comptime_int) void { |
| 72 | const Table = ShardedTable(Key, mask_bit_count, void); |
| 73 | |
| 74 | var table = Table.create(); |
| 75 | var node_buffer: [node_count]Table.Node = undefined; |
| 76 | for (node_buffer) |*node, i| { |
| 77 | const key = @intCast(Key, i); |
| 78 | assert(table.get(key) == null); |
| 79 | node.init(key, {}); |
| 80 | table.put(node); |
| 81 | } |
| 82 | |
| 83 | for (node_buffer) |*node, i| { |
| 84 | assert(table.get(@intCast(Key, i)) == node); |
| 85 | } |
| 86 | } |