diff --git a/src/kern/cpu.zig b/src/kern/cpu.zig index 73b1d88..bf49b74 100644 --- a/src/kern/cpu.zig +++ b/src/kern/cpu.zig @@ -68,3 +68,4 @@ pub fn ExportedCpuLocal(comptime T: type, comptime init: T, comptime name: []con /// Bitmask of CPUs. pub const CpuMask = rtl.BitMap(config.ncpus); +pub const AtomicCpuMask = rtl.AtomicBitMap(config.ncpus); diff --git a/src/rtl/bitmap.zig b/src/rtl/bitmap.zig index a599afb..f6f9fa3 100644 --- a/src/rtl/bitmap.zig +++ b/src/rtl/bitmap.zig @@ -5,6 +5,10 @@ fn get_closest_type(comptime N: usize) type { return @Int(.unsigned, bits); } +fn bit_mask(comptime Word: type, bit_index: usize) Word { + return @as(Word, 1) << @truncate(bit_index); +} + /// Static bitmap that uses the most efficient underlying word size for N < 64 pub fn BitMap(comptime N: usize) type { return struct { @@ -28,7 +32,7 @@ pub fn BitMap(comptime N: usize) type { 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)); + self.storage[word_index] |= bit_mask(Word, bit_index); } /// Clear a given bit. @@ -36,7 +40,7 @@ pub fn BitMap(comptime N: usize) type { 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)); + self.storage[word_index] &= ~bit_mask(Word, bit_index); } /// Get a given bit. @@ -44,7 +48,7 @@ pub fn BitMap(comptime N: usize) type { 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; + return self.storage[word_index] & bit_mask(Word, bit_index) != 0; } /// Check whether all bits in the bitmap are equal to a given value. @@ -58,6 +62,63 @@ pub fn BitMap(comptime N: usize) type { }; } +/// Atomic static bitmap with the same API as `BitMap`, but each operation accepts a memory ordering. +pub fn AtomicBitMap(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); + const arr_len = (N + bits_per_word - 1) / bits_per_word; + + storage: [arr_len]std.atomic.Value(Word), + + /// Initialize a bitmap, setting all bits to `val`. + pub fn init(val: bool) Self { + var storage: [arr_len]std.atomic.Value(Word) = undefined; + const word: Word = if (val) std.math.maxInt(Self.Word) else 0; + inline for (&storage) |*elem| { + elem.* = std.atomic.Value(Word).init(word); + } + return .{ .storage = storage }; + } + + /// Set a given bit. + pub fn set(self: *Self, bit: usize, comptime ordering: std.builtin.AtomicOrder) 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].fetchOr(bit_mask(Word, bit_index), ordering); + } + + /// Clear a given bit. + pub fn clear(self: *Self, bit: usize, comptime ordering: std.builtin.AtomicOrder) 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].fetchAnd(~bit_mask(Word, bit_index), ordering); + } + + /// Get a given bit. + pub fn get(self: *const Self, bit: usize, comptime ordering: std.builtin.AtomicOrder) 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].load(ordering) & bit_mask(Word, 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, comptime ordering: std.builtin.AtomicOrder) bool { + const v: Self.Word = if (value) std.math.maxInt(Self.Word) else 0; + for (self.storage) |word| { + if (word.load(ordering) != 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); @@ -82,3 +143,19 @@ test BitMap { std.debug.assert(bitmap.get(0) == false); } + +test AtomicBitMap { + var bitmap = AtomicBitMap(10).init(false); + + std.debug.assert(@TypeOf(bitmap).Word == u16); + std.debug.assert(bitmap.get(0, .seq_cst) == false); + bitmap.set(0, .seq_cst); + + std.debug.assert(bitmap.get(0, .seq_cst) == true); + std.debug.assert(bitmap.is_all(false, .seq_cst) == false); + + bitmap.clear(0, .seq_cst); + + std.debug.assert(bitmap.get(0, .seq_cst) == false); + std.debug.assert(bitmap.is_all(false, .seq_cst) == true); +} diff --git a/src/rtl/root.zig b/src/rtl/root.zig index eeb5343..6048d97 100644 --- a/src/rtl/root.zig +++ b/src/rtl/root.zig @@ -12,6 +12,7 @@ pub const BST = bst.BST; pub const RBTree = @import("rbtree.zig").RBTree; pub const AVLTree = @import("avl.zig").AVLTree; pub const BitMap = @import("bitmap.zig").BitMap; +pub const AtomicBitMap = @import("bitmap.zig").AtomicBitMap; test { std.testing.refAllDecls(@This());