tiny.simd.bitset

1const std = @import("std");2const builtin = @import("builtin");34const runtime_assertions = builtin.mode == .debug or builtin.mode == .safe;56pub const BitSet64 = struct {7    bits: u64 = 0,89    pub const max_size: usize = 64;1011    pub fn maxSize(_: *const @This()) usize {12        return max_size;13    }1415    pub fn set(self: *@This(), index: usize) void {16        std.debug.assert(index < max_size);17        self.bits |= bit(index);18        std.debug.assert(self.get(index));19    }2021    pub fn setNonzeroBitsFrom64(self: *@This(), bits: u64) void {22        self.bits |= bits;23    }2425    pub fn clear(self: *@This(), index: usize) void {26        std.debug.assert(index < max_size);27        self.bits &= ~bit(index);28        std.debug.assert(!self.get(index));29    }3031    pub fn get(self: *const @This(), index: usize) bool {32        std.debug.assert(index < max_size);33        return self.bits & bit(index) != 0;34    }3536    pub fn any(self: *const @This()) bool {37        return self.bits != 0;38    }3940    pub fn all(self: *const @This()) bool {41        return self.bits == std.math.maxInt(u64);42    }4344    pub fn first(self: *const @This()) usize {45        std.debug.assert(self.any());46        return firstBit(self.bits);47    }4849    pub fn first0(self: *const @This()) usize {50        std.debug.assert(!self.all());51        return firstBit(~self.bits);52    }5354    pub fn get64(self: *const @This()) u64 {55        return self.bits;56    }5758    pub fn foreach(self: *const @This(), function: anytype) void {59        visitWord(self.bits, 0, function);60    }6162    pub fn count(self: *const @This()) usize {63        return @popCount(self.bits);64    }65};6667pub fn BitSet(comptime size: usize) type {68    requireCapacity(size);69    const word_count = wordsFor(size);7071    return struct {72        words: [word_count]BitSet64 = @as([word_count]BitSet64, @splat(.{})),7374        const Self = @This();75        pub const capacity: usize = size;76        pub const max_size: usize = capacity;7778        pub fn maxSize(_: *const Self) usize {79            return max_size;80        }8182        pub fn set(self: *Self, index: usize) void {83            std.debug.assert(index < max_size);84            self.words[wordIndex(index)].set(bitIndex(index));85            std.debug.assert(self.get(index));86        }8788        pub fn clear(self: *Self, index: usize) void {89            std.debug.assert(index < max_size);90            self.words[wordIndex(index)].clear(bitIndex(index));91            std.debug.assert(!self.get(index));92        }9394        pub fn get(self: *const Self, index: usize) bool {95            std.debug.assert(index < max_size);96            return self.words[wordIndex(index)].get(bitIndex(index));97        }9899        pub fn any(self: *const Self) bool {100            for (&self.words) |*word| {101                if (word.any()) return true;102            }103            return false;104        }105106        pub fn all(self: *const Self) bool {107            for (self.words[0 .. word_count - 1]) |word| {108                if (!word.all()) return false;109            }110            const remainder = max_size % 64;111            return if (remainder == 0)112                self.words[word_count - 1].all()113            else114                self.words[word_count - 1].count() == remainder;115        }116117        pub fn first(self: *const Self) usize {118            std.debug.assert(self.any());119            for (&self.words, 0..) |*word, word_index| {120                if (word.any()) return word_index * 64 + word.first();121            }122            unreachable;123        }124125        pub fn first0(self: *const Self) usize {126            std.debug.assert(!self.all());127            for (&self.words, 0..) |*word, word_index| {128                if (!word.all()) {129                    const index = word_index * 64 + word.first0();130                    std.debug.assert(index < max_size);131                    return index;132                }133            }134            unreachable;135        }136137        pub fn foreach(self: *const Self, function: anytype) void {138            for (&self.words, 0..) |*word, word_index| {139                visitWord(word.get64(), word_index * 64, function);140            }141        }142143        pub fn count(self: *const Self) usize {144            var total: usize = 0;145            for (&self.words) |*word| total += word.count();146            std.debug.assert(total <= max_size);147            return total;148        }149    };150}151152pub fn AtomicBitSet(comptime size: usize) type {153    requireCapacity(size);154    const Word = usize;155    const word_bits: usize = @bitSizeOf(Word);156    const word_count = std.math.divCeil(usize, size, word_bits) catch unreachable;157    const AtomicWord = std.atomic.Value(Word);158159    return struct {160        words: [word_count]AtomicWord = @as([word_count]AtomicWord, @splat(AtomicWord.init(0))),161162        const Self = @This();163        pub const capacity: usize = size;164        pub const max_size: usize = capacity;165166        pub fn maxSize(_: *const Self) usize {167            return max_size;168        }169170        pub fn set(self: *Self, index: usize) void {171            std.debug.assert(index < max_size);172            _ = self.words[atomicWordIndex(index)].fetchOr(173                atomicBit(atomicBitIndex(index)),174                .release,175            );176        }177178        pub fn clear(self: *Self, index: usize) void {179            std.debug.assert(index < max_size);180            _ = self.words[atomicWordIndex(index)].fetchAnd(181                ~atomicBit(atomicBitIndex(index)),182                .release,183            );184            if (runtime_assertions) std.debug.assert(!self.get(index));185        }186187        pub fn get(self: *const Self, index: usize) bool {188            std.debug.assert(index < max_size);189            return self.words[atomicWordIndex(index)].load(.acquire) &190                atomicBit(atomicBitIndex(index)) != 0;191        }192193        pub fn any(self: *const Self) bool {194            for (&self.words) |*word| {195                if (word.load(.acquire) != 0) return true;196            }197            return false;198        }199200        pub fn all(self: *const Self) bool {201            for (self.words[0 .. word_count - 1]) |*word| {202                if (word.load(.acquire) != std.math.maxInt(Word)) return false;203            }204            const last = self.words[word_count - 1].load(.acquire);205            const remainder = max_size % word_bits;206            return if (remainder == 0)207                last == std.math.maxInt(Word)208            else209                @popCount(last) == remainder;210        }211212        pub fn first(self: *const Self) usize {213            if (runtime_assertions) std.debug.assert(self.any());214            for (&self.words, 0..) |*word, word_index| {215                const bits = word.load(.acquire);216                if (bits != 0) {217                    return word_index * word_bits + atomicFirstBit(bits);218                }219            }220            unreachable;221        }222223        pub fn first0(self: *const Self) usize {224            if (runtime_assertions) std.debug.assert(!self.all());225            for (&self.words, 0..) |*word, word_index| {226                const inverted = ~word.load(.acquire);227                if (inverted != 0) {228                    const index = word_index * word_bits +229                        atomicFirstBit(inverted);230                    std.debug.assert(index < max_size);231                    return index;232                }233            }234            unreachable;235        }236237        pub fn foreach(self: *const Self, function: anytype) void {238            for (&self.words, 0..) |*word, word_index| {239                visitAtomicWord(240                    word.load(.acquire),241                    word_index * word_bits,242                    function,243                );244            }245        }246247        pub fn count(self: *const Self) usize {248            var total: usize = 0;249            for (&self.words) |*word| total += @popCount(word.load(.acquire));250            std.debug.assert(total <= max_size);251            return total;252        }253254        fn atomicWordIndex(index: usize) usize {255            return index / word_bits;256        }257258        fn atomicBitIndex(index: usize) usize {259            return index % word_bits;260        }261262        fn atomicBit(index: usize) Word {263            std.debug.assert(index < word_bits);264            return @as(Word, 1) << @intCast(index);265        }266267        fn atomicFirstBit(bits: Word) usize {268            std.debug.assert(bits != 0);269            return @ctz(bits);270        }271272        fn visitAtomicWord(bits: Word, base: usize, function: anytype) void {273            var remaining = bits;274            while (remaining != 0) {275                const index = atomicFirstBit(remaining);276                remaining &= remaining - 1;277                function.call(base + index);278            }279        }280    };281}282283pub fn BitSet4096(comptime size: usize) type {284    requireCapacity(size);285    if (size > 4096) @compileError("BitSet4096 supports at most 4096 bits");286    const word_count = wordsFor(size);287288    return struct {289        nonzero: BitSet64 = .{},290        words: [word_count]BitSet64 = @as([word_count]BitSet64, @splat(.{})),291292        const Self = @This();293        pub const capacity: usize = size;294        pub const max_size: usize = capacity;295296        pub fn maxSize(_: *const Self) usize {297            return max_size;298        }299300        pub fn set(self: *Self, index: usize) void {301            std.debug.assert(index < max_size);302            const word_index = wordIndex(index);303            self.words[word_index].set(bitIndex(index));304            self.nonzero.set(word_index);305            std.debug.assert(self.get(index));306        }307308        pub fn setNonzeroBitsFrom64(self: *Self, bits: u64) void {309            const valid = bits & validBits(@min(max_size, 64));310            std.debug.assert(valid == bits);311            self.words[0].setNonzeroBitsFrom64(valid);312            if (valid != 0) self.nonzero.set(0);313        }314315        pub fn clear(self: *Self, index: usize) void {316            std.debug.assert(index < max_size);317            const word_index = wordIndex(index);318            self.words[word_index].clear(bitIndex(index));319            if (!self.words[word_index].any()) self.nonzero.clear(word_index);320            std.debug.assert(!self.get(index));321        }322323        pub fn get(self: *const Self, index: usize) bool {324            std.debug.assert(index < max_size);325            return self.words[wordIndex(index)].get(bitIndex(index));326        }327328        pub fn any(self: *const Self) bool {329            return self.nonzero.any();330        }331332        pub fn all(self: *const Self) bool {333            if (self.nonzero.count() != word_count) return false;334            return self.count() == max_size;335        }336337        pub fn first(self: *const Self) usize {338            std.debug.assert(self.any());339            const word_index = self.nonzero.first();340            return word_index * 64 + self.words[word_index].first();341        }342343        pub fn first0(self: *const Self) usize {344            std.debug.assert(!self.all());345            for (&self.words, 0..) |*word, word_index| {346                if (!word.all()) {347                    const index = word_index * 64 + word.first0();348                    std.debug.assert(index < max_size);349                    return index;350                }351            }352            unreachable;353        }354355        pub fn get64(self: *const Self) u64 {356            return self.words[0].get64();357        }358359        pub fn foreach(self: *const Self, function: anytype) void {360            var remaining_words = self.nonzero.get64();361            while (remaining_words != 0) {362                const word_index = firstBit(remaining_words);363                remaining_words &= remaining_words - 1;364                visitWord(self.words[word_index].get64(), word_index * 64, function);365            }366        }367368        pub fn count(self: *const Self) usize {369            var total: usize = 0;370            var remaining_words = self.nonzero.get64();371            while (remaining_words != 0) {372                const word_index = firstBit(remaining_words);373                remaining_words &= remaining_words - 1;374                total += self.words[word_index].count();375            }376            std.debug.assert(total <= max_size);377            return total;378        }379    };380}381382pub const DefaultBitSet4096 = BitSet4096(4096);383384fn visitWord(bits: u64, base: usize, function: anytype) void {385    var remaining = bits;386    while (remaining != 0) {387        const index = firstBit(remaining);388        remaining &= remaining - 1;389        function.call(base + index);390    }391}392393fn firstBit(bits: u64) usize {394    std.debug.assert(bits != 0);395    return @ctz(bits);396}397398fn bit(index: usize) u64 {399    std.debug.assert(index < 64);400    return @as(u64, 1) << @intCast(index);401}402403fn wordIndex(index: usize) usize {404    return index / 64;405}406407fn bitIndex(index: usize) usize {408    return index % 64;409}410411fn wordsFor(max_size: usize) usize {412    return (max_size + 63) / 64;413}414415fn validBits(count: usize) u64 {416    std.debug.assert(count <= 64);417    return if (count == 64)418        std.math.maxInt(u64)419    else420        (@as(u64, 1) << @intCast(count)) - 1;421}422423fn requireCapacity(comptime max_size: usize) void {424    if (max_size == 0) @compileError("bit sets require a nonzero capacity");425}426427const Collector = struct {428    values: []usize,429    count: usize = 0,430431    pub fn call(self: *@This(), index: usize) void {432        std.debug.assert(self.count < self.values.len);433        self.values[self.count] = index;434        self.count += 1;435    }436};437438fn smoke(comptime Set: type) !void {439    var set = Set{};440    try std.testing.expectEqual(Set.max_size, set.maxSize());441    try std.testing.expect(!set.any());442    try std.testing.expect(!set.all());443    try std.testing.expect(!set.get(0));444    try std.testing.expect(!set.get(Set.max_size - 1));445    try std.testing.expectEqual(@as(usize, 0), set.first0());446    var values: [Set.max_size]usize = undefined;447    var collector = Collector{ .values = &values };448    set.foreach(&collector);449    try std.testing.expectEqual(@as(usize, 0), collector.count);450    try std.testing.expectEqual(@as(usize, 0), set.count());451452    const last = Set.max_size - 1;453    set.set(last);454    try std.testing.expect(set.get(last));455    try std.testing.expect(set.any());456    try std.testing.expect(!set.all());457    try std.testing.expectEqual(last, set.first());458    try std.testing.expectEqual(@as(usize, 0), set.first0());459    collector.count = 0;460    set.foreach(&collector);461    try std.testing.expectEqual(@as(usize, 1), collector.count);462    try std.testing.expectEqual(last, collector.values[0]);463    try std.testing.expectEqual(@as(usize, 1), set.count());464465    set.clear(last);466    set.clear(0);467    try std.testing.expect(!set.any());468    try std.testing.expect(!set.all());469    try std.testing.expectEqual(@as(usize, 0), set.first0());470    try std.testing.expectEqual(@as(usize, 0), set.count());471}472473fn verifyFull(comptime Set: type) !void {474    var set = Set{};475    for (0..Set.max_size) |index| set.set(index);476    try std.testing.expect(set.any());477    try std.testing.expect(set.all());478    try std.testing.expectEqual(@as(usize, 0), set.first());479    try std.testing.expectEqual(Set.max_size, set.count());480    const missing = Set.max_size / 2;481    set.clear(missing);482    try std.testing.expect(!set.all());483    try std.testing.expectEqual(missing, set.first0());484    set.set(missing);485    try std.testing.expect(set.all());486}487488fn verifyModel(comptime Set: type, set: *const Set, model: *const [Set.max_size]bool) !void {489    var expected_count: usize = 0;490    var expected_first: usize = Set.max_size;491    var expected_first0: usize = Set.max_size;492    for (model, 0..) |present, index| {493        try std.testing.expectEqual(present, set.get(index));494        expected_count += @intFromBool(present);495        if (present and expected_first == Set.max_size) expected_first = index;496        if (!present and expected_first0 == Set.max_size) expected_first0 = index;497    }498    try std.testing.expectEqual(expected_count, set.count());499    try std.testing.expectEqual(expected_count != 0, set.any());500    try std.testing.expectEqual(expected_count == Set.max_size, set.all());501    if (expected_first != Set.max_size) try std.testing.expectEqual(expected_first, set.first());502    if (expected_first0 != Set.max_size) try std.testing.expectEqual(expected_first0, set.first0());503504    var values: [Set.max_size]usize = undefined;505    var collector = Collector{ .values = &values };506    set.foreach(&collector);507    try std.testing.expectEqual(expected_count, collector.count);508    var seen: usize = 0;509    for (model, 0..) |present, index| {510        if (present) {511            try std.testing.expectEqual(index, collector.values[seen]);512            seen += 1;513        }514    }515}516517fn randomWalk(comptime Set: type, grow_percent: u8, seed: u64) !void {518    var prng = std.Random.DefaultPrng.init(seed ^ Set.max_size);519    const random = prng.random();520    for (0..16) |_| {521        var set = Set{};522        var model: [Set.max_size]bool = @splat(false);523        for (0..128) |_| {524            const index = random.uintLessThan(usize, Set.max_size);525            if (random.uintLessThan(u8, 100) < grow_percent) {526                set.set(index);527                model[index] = true;528            } else {529                set.clear(index);530                model[index] = false;531            }532            try std.testing.expectEqual(model[index], set.get(index));533        }534        try verifyModel(Set, &set, &model);535    }536}537538const Mutate64 = struct {539    set: *BitSet64,540    visited: *BitSet64,541542    pub fn call(self: *@This(), index: usize) void {543        self.visited.set(index);544        if (index == 1) {545            self.set.clear(2);546            self.set.set(3);547        }548    }549};550551fn MutateFuture(comptime Set: type) type {552    return struct {553        set: *Set,554        visited: *Set,555556        pub fn call(self: *@This(), index: usize) void {557            self.visited.set(index);558            if (index == 1) {559                self.set.clear(70);560                self.set.set(71);561            }562        }563    };564}565566fn verifyFutureMutation(comptime Set: type) !void {567    var set = Set{};568    set.set(1);569    set.set(70);570    var visited = Set{};571    var mutation = MutateFuture(Set){ .set = &set, .visited = &visited };572    set.foreach(&mutation);573    try std.testing.expect(visited.get(1));574    try std.testing.expect(!visited.get(70));575    try std.testing.expect(visited.get(71));576}577578test "Highway bit sets default empty and preserve full-capacity boundaries" {579    inline for (.{580        BitSet64,581        BitSet(320),582        AtomicBitSet(400),583        DefaultBitSet4096,584    }) |Set| try smoke(Set);585    inline for (.{586        BitSet64,587        BitSet(41),588        BitSet(192),589        AtomicBitSet(32),590        AtomicBitSet(192),591        BitSet4096(3000),592        DefaultBitSet4096,593    }) |Set| try verifyFull(Set);594}595596test "Highway bit sets import nonzero low words without clearing" {597    inline for (.{ BitSet64, DefaultBitSet4096 }) |Set| {598        var set = Set{};599        set.setNonzeroBitsFrom64(1);600        set.setNonzeroBitsFrom64(0x70);601        try std.testing.expectEqual(@as(u64, 0x71), set.get64());602        try std.testing.expectEqual(@as(usize, 4), set.count());603        try std.testing.expectEqual(@as(usize, 0), set.first());604        try std.testing.expectEqual(@as(usize, 1), set.first0());605    }606}607608test "Highway bit set random walks match fixed scalar models" {609    inline for (.{610        BitSet64,611        BitSet(41),612        BitSet(199),613        AtomicBitSet(32),614        AtomicBitSet(192),615        BitSet4096(3000),616        DefaultBitSet4096,617    }, 0..) |Set, index| {618        try randomWalk(Set, 40, 0x4253_4554 + index);619        try randomWalk(Set, 60, 0x5345_5442 + index);620    }621}622623test "Highway bit set foreach snapshots current words and observes future words" {624    var set64 = BitSet64{};625    set64.set(1);626    set64.set(2);627    var visited64 = BitSet64{};628    var mutation64 = Mutate64{ .set = &set64, .visited = &visited64 };629    set64.foreach(&mutation64);630    try std.testing.expect(visited64.get(1));631    try std.testing.expect(visited64.get(2));632    try std.testing.expect(!visited64.get(3));633    try verifyFutureMutation(BitSet(128));634    try verifyFutureMutation(AtomicBitSet(128));635    try verifyFutureMutation(BitSet4096(128));636}

1pub const bitset = @import("bitset.zig");