Merge pull request #1543 from ziglang/bit_shifting

somewhat realistic usecase test for shifting strange integer sizes

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");
 }

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);
+    }
+}