diff --git a/test/behavior.zig b/test/behavior.zig index 3223131d0815..b91d41c74de5 100644 --- a/test/behavior.zig +++ b/test/behavior.zig @@ -71,4 +71,5 @@ comptime { _ = @import("cases/void.zig"); _ = @import("cases/while.zig"); _ = @import("cases/widening.zig"); + _ = @import("cases/bit_shifting.zig"); } diff --git a/test/cases/bit_shifting.zig b/test/cases/bit_shifting.zig new file mode 100644 index 000000000000..601793d5ed38 --- /dev/null +++ b/test/cases/bit_shifting.zig @@ -0,0 +1,86 @@ +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 { + // 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: 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); + } +}