rtl: add BitMap

Adds a new BitMap type in rtl for static bitmaps, now used to back
ke.CpuMask.

src/kern/cpu.zig

@@ -67,42 +67,4 @@ pub fn ExportedCpuLocal(comptime T: type, comptime init: T, comptime name: []con
 }
 
 /// Bitmask of CPUs.
-pub const CpuMask = struct {
-    const bits_per_word = @bitSizeOf(usize);
-    const num_words = (config.ncpus + bits_per_word - 1) / bits_per_word;
-
-    bits: [num_words]usize,
-
-    /// Return an empty `CpuMask`.
-    pub fn empty() CpuMask {
-        return CpuMask{ .bits = @splat(0) };
-    }
-
-    /// Set the bit corresponding to `cpu_id`.
-    pub fn set(self: *CpuMask, cpu_id: usize) void {
-        const word_index = cpu_id / bits_per_word;
-        const bit_index = cpu_id % bits_per_word;
-        self.bits[word_index] |= (@as(usize, 1) << @intCast(bit_index));
-    }
-
-    /// Clear the bit corresponding to `cpu_id`.
-    pub fn clear(self: *CpuMask, cpu_id: usize) void {
-        const word_index = cpu_id / bits_per_word;
-        const bit_index = cpu_id % bits_per_word;
-        self.bits[word_index] &= ~(@as(usize, 1) << @intCast(bit_index));
-    }
-
-    /// Check whether the bit corresponding to `cpu_id` is set.
-    pub fn is_set(self: *const CpuMask, cpu_id: usize) bool {
-        const word_index = cpu_id / bits_per_word;
-        const bit_index = cpu_id % bits_per_word;
-        return (self.bits[word_index] & (@as(usize, 1) << @intCast(bit_index))) != 0;
-    }
-
-    pub fn is_full(self: *const CpuMask) bool {
-        for (self.bits) |word| {
-            if (word != std.math.maxInt(usize)) return false;
-        }
-        return true;
-    }
-};
+pub const CpuMask = rtl.BitMap(config.ncpus);

src/kern/sched.zig

@@ -846,7 +846,7 @@ fn find_most_loaded(exclude: ?*ke.CpuMask) ?u32 {
 
     for (0..ke.ncpus) |i| {
         if (exclude) |mask| {
-            if (mask.is_set(i)) {
+            if (mask.get(i)) {
                 continue;
             }
         }
@@ -881,7 +881,7 @@ fn find_least_loaded(exclude: ?*ke.CpuMask) ?u32 {
 
     for (0..ke.ncpus) |i| {
         if (exclude) |mask| {
-            if (mask.is_set(i)) {
+            if (mask.get(i)) {
                 continue;
             }
         }
@@ -969,7 +969,7 @@ pub fn idle(_: ?*anyopaque) noreturn {
 /// Called every second to balance work between CPUs on CPU0.
 /// Takes a thread from the most loaded CPU and puts it on the least loaded one.
 fn balance(_: ?*anyopaque) void {
-    var high_mask = ke.CpuMask.empty();
+    var high_mask = ke.CpuMask.init(false);
     var low_mask: ke.CpuMask = undefined;
 
     while (true) {
@@ -984,7 +984,7 @@ fn balance(_: ?*anyopaque) void {
         high_mask.set(high.?);
         low_mask = high_mask;
 
-        if (high_mask.is_full()) {
+        if (high_mask.is_all(true)) {
             // All CPUs are masked, nothing to steal.
             break;
         }

src/rtl/bitmap.zig

@@ -0,0 +1,82 @@
+const std = @import("std");
+
+fn get_closest_type(comptime N: usize) type {
+    const bits = std.math.ceilPowerOfTwo(usize, N) catch unreachable;
+    return std.meta.Int(.unsigned, bits);
+}
+
+/// Static bitmap that uses the most efficient underlying word size for N < 64
+pub fn BitMap(comptime N: usize) type {
+    return struct {
+        pub const Word = if (N > 32) u64 else get_closest_type(N);
+        const Self = @This();
+        const bits_per_word = @bitSizeOf(Word);
+
+        storage: [bits_per_word / N]Word,
+
+        /// Initialize a bitmap, setting all bits to `val`.
+        pub fn init(val: bool) Self {
+            return .{
+                .storage = @splat(if (val) std.math.maxInt(Self.Word) else 0),
+            };
+        }
+
+        /// Set a given bit.
+        pub fn set(self: *Self, bit: usize) void {
+            std.debug.assert(bit < N);
+            const word_index = bit / Self.bits_per_word;
+            const bit_index = bit % Self.bits_per_word;
+
+            self.storage[word_index] |= (@as(Word, 1) << @truncate(bit_index));
+        }
+
+        /// Clear a given bit.
+        pub fn clear(self: *Self, bit: usize) void {
+            std.debug.assert(bit < N);
+            const word_index = bit / Self.bits_per_word;
+            const bit_index = bit % Self.bits_per_word;
+            self.storage[word_index] &= ~(@as(Word, 1) << @truncate(bit_index));
+        }
+
+        /// Get a given bit.
+        pub fn get(self: *const Self, bit: usize) bool {
+            std.debug.assert(bit < N);
+            const word_index = bit / Self.bits_per_word;
+            const bit_index = bit % Self.bits_per_word;
+            return self.storage[word_index] & (@as(Word, 1) << @truncate(bit_index)) != 0;
+        }
+
+        /// Check whether all bits in the bitmap are equal to a given value.
+        pub fn is_all(self: *const Self, value: bool) bool {
+            const v: Self.Word = if (value) std.math.maxInt(Self.Word) else 0;
+            for (self.storage) |word| {
+                if (word != v) return false;
+            }
+            return true;
+        }
+    };
+}
+
+test "BitMap has proper underlying type" {
+    std.debug.assert(BitMap(8).Word == u8);
+    std.debug.assert(BitMap(10).Word == u16);
+    std.debug.assert(BitMap(16).Word == u16);
+    std.debug.assert(BitMap(21).Word == u32);
+    std.debug.assert(BitMap(32).Word == u32);
+    std.debug.assert(BitMap(64).Word == u64);
+    std.debug.assert(BitMap(128).Word == u64);
+}
+
+test BitMap {
+    var bitmap = BitMap(10).init(false);
+
+    std.debug.assert(@TypeOf(bitmap).Word == u16);
+    std.debug.assert(bitmap.get(0) == false);
+    bitmap.set(0);
+
+    std.debug.assert(bitmap.get(0) == true);
+
+    bitmap.clear(0);
+
+    std.debug.assert(bitmap.get(0) == false);
+}

src/rtl/root.zig

@@ -10,6 +10,7 @@ pub const TaggedPtr = @import("tagged_ptr.zig").TaggedPtr;
 pub const bst = @import("bst.zig");
 pub const BST = bst.BST;
 pub const RBTree = @import("rbtree.zig").RBTree;
+pub const BitMap = @import("bitmap.zig").BitMap;
 
 test {
     std.testing.refAllDecls(@This());