From 20a38f3a490ef6d9c24f1f1bef3e952b3b61880f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Gr=C3=A9goire?= Date: Sun, 21 Jun 2026 21:42:19 -0400 Subject: [PATCH] rtl: add count and iter to bitmap --- src/rtl/bitmap.zig | 174 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 174 insertions(+) diff --git a/src/rtl/bitmap.zig b/src/rtl/bitmap.zig index f6f9fa3..b763e17 100644 --- a/src/rtl/bitmap.zig +++ b/src/rtl/bitmap.zig @@ -19,6 +19,16 @@ pub fn BitMap(comptime N: usize) type { storage: [arr_len]Word, + fn valid_word_mask(word_index: usize) Word { + std.debug.assert(word_index < arr_len); + + if (word_index != arr_len - 1) return std.math.maxInt(Word); + + const rem = N % bits_per_word; + if (rem == 0) return std.math.maxInt(Word); + return (@as(Word, 1) << @truncate(rem)) - 1; + } + /// Initialize a bitmap, setting all bits to `val`. pub fn init(val: bool) Self { return .{ @@ -51,6 +61,42 @@ pub fn BitMap(comptime N: usize) type { return self.storage[word_index] & bit_mask(Word, bit_index) != 0; } + /// Count the number of set bits in the bitmap. + pub fn count(self: *const Self) usize { + var total: usize = 0; + for (self.storage, 0..) |word, i| { + total += @intCast(@popCount(word & valid_word_mask(i))); + } + return total; + } + + pub const Iterator = struct { + bitmap: *const Self, + word_index: usize = 0, + word: Word = 0, + + /// Return the index of the next set bit, or null when exhausted. + pub fn next(self: *Iterator) ?usize { + while (true) { + if (self.word != 0) { + const bit_index: usize = @intCast(@ctz(self.word)); + self.word &= self.word - 1; + return ((self.word_index - 1) * bits_per_word) + bit_index; + } + + if (self.word_index >= arr_len) return null; + + self.word = self.bitmap.storage[self.word_index] & valid_word_mask(self.word_index); + self.word_index += 1; + } + } + }; + + /// Iterate over the indexes of set bits in ascending order. + pub fn iter(self: *const Self) Iterator { + return .{ .bitmap = self }; + } + /// 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; @@ -72,6 +118,16 @@ pub fn AtomicBitMap(comptime N: usize) type { storage: [arr_len]std.atomic.Value(Word), + fn valid_word_mask(word_index: usize) Word { + std.debug.assert(word_index < arr_len); + + if (word_index != arr_len - 1) return std.math.maxInt(Word); + + const rem = N % bits_per_word; + if (rem == 0) return std.math.maxInt(Word); + return (@as(Word, 1) << @truncate(rem)) - 1; + } + /// Initialize a bitmap, setting all bits to `val`. pub fn init(val: bool) Self { var storage: [arr_len]std.atomic.Value(Word) = undefined; @@ -108,6 +164,44 @@ pub fn AtomicBitMap(comptime N: usize) type { return self.storage[word_index].load(ordering) & bit_mask(Word, bit_index) != 0; } + /// Count the number of set bits in the bitmap. + pub fn count(self: *const Self, comptime ordering: std.builtin.AtomicOrder) usize { + var total: usize = 0; + for (self.storage, 0..) |word, i| { + total += @intCast(@popCount(word.load(ordering) & valid_word_mask(i))); + } + return total; + } + + pub fn Iterator(comptime ordering: std.builtin.AtomicOrder) type { + return struct { + bitmap: *const Self, + word_index: usize = 0, + word: Word = 0, + + /// Return the index of the next set bit, or null when exhausted. + pub fn next(self: *@This()) ?usize { + while (true) { + if (self.word != 0) { + const bit_index: usize = @intCast(@ctz(self.word)); + self.word &= self.word - 1; + return ((self.word_index - 1) * bits_per_word) + bit_index; + } + + if (self.word_index >= arr_len) return null; + + self.word = self.bitmap.storage[self.word_index].load(ordering) & valid_word_mask(self.word_index); + self.word_index += 1; + } + } + }; + } + + /// Iterate over the indexes of set bits in ascending order. + pub fn iter(self: *const Self, comptime ordering: std.builtin.AtomicOrder) Iterator(ordering) { + return .{ .bitmap = self }; + } + /// 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; @@ -144,6 +238,46 @@ test BitMap { std.debug.assert(bitmap.get(0) == false); } +test "BitMap count and iter" { + var bitmap = BitMap(70).init(false); + + std.debug.assert(bitmap.count() == 0); + { + var it = bitmap.iter(); + std.debug.assert(it.next() == null); + } + + bitmap.set(0); + bitmap.set(3); + bitmap.set(63); + bitmap.set(64); + bitmap.set(69); + + std.debug.assert(bitmap.count() == 5); + + var it = bitmap.iter(); + std.debug.assert(it.next().? == 0); + std.debug.assert(it.next().? == 3); + std.debug.assert(it.next().? == 63); + std.debug.assert(it.next().? == 64); + std.debug.assert(it.next().? == 69); + std.debug.assert(it.next() == null); + + bitmap.clear(63); + + std.debug.assert(bitmap.count() == 4); + + var all = BitMap(70).init(true); + std.debug.assert(all.count() == 70); + + var all_it = all.iter(); + var expected: usize = 0; + while (all_it.next()) |bit| : (expected += 1) { + std.debug.assert(bit == expected); + } + std.debug.assert(expected == 70); +} + test AtomicBitMap { var bitmap = AtomicBitMap(10).init(false); @@ -159,3 +293,43 @@ test AtomicBitMap { std.debug.assert(bitmap.get(0, .seq_cst) == false); std.debug.assert(bitmap.is_all(false, .seq_cst) == true); } + +test "AtomicBitMap count and iter" { + var bitmap = AtomicBitMap(70).init(false); + + std.debug.assert(bitmap.count(.seq_cst) == 0); + { + var it = bitmap.iter(.seq_cst); + std.debug.assert(it.next() == null); + } + + bitmap.set(0, .seq_cst); + bitmap.set(3, .seq_cst); + bitmap.set(63, .seq_cst); + bitmap.set(64, .seq_cst); + bitmap.set(69, .seq_cst); + + std.debug.assert(bitmap.count(.seq_cst) == 5); + + var it = bitmap.iter(.seq_cst); + std.debug.assert(it.next().? == 0); + std.debug.assert(it.next().? == 3); + std.debug.assert(it.next().? == 63); + std.debug.assert(it.next().? == 64); + std.debug.assert(it.next().? == 69); + std.debug.assert(it.next() == null); + + bitmap.clear(63, .seq_cst); + + std.debug.assert(bitmap.count(.seq_cst) == 4); + + var all = AtomicBitMap(70).init(true); + std.debug.assert(all.count(.seq_cst) == 70); + + var all_it = all.iter(.seq_cst); + var expected: usize = 0; + while (all_it.next()) |bit| : (expected += 1) { + std.debug.assert(bit == expected); + } + std.debug.assert(expected == 70); +}