lib/simd/src/bitset.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const builtin = @import("builtin");
3
4 const runtime_assertions = builtin.mode == .debug or builtin.mode == .safe;
5
6 pub const BitSet64 = struct {
7 bits: u64 = 0,
8
9 pub const max_size: usize = 64;
10
11 pub fn maxSize(_: *const @This()) usize {
12 return max_size;
13 }
14
15 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 }
20
21 pub fn setNonzeroBitsFrom64(self: *@This(), bits: u64) void {
22 self.bits |= bits;
23 }
24
25 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 }
30
31 pub fn get(self: *const @This(), index: usize) bool {
32 std.debug.assert(index < max_size);
33 return self.bits & bit(index) != 0;
34 }
35
36 pub fn any(self: *const @This()) bool {
37 return self.bits != 0;
38 }
39
40 pub fn all(self: *const @This()) bool {
41 return self.bits == std.math.maxInt(u64);
42 }
43
44 pub fn first(self: *const @This()) usize {
45 std.debug.assert(self.any());
46 return firstBit(self.bits);
47 }
48
49 pub fn first0(self: *const @This()) usize {
50 std.debug.assert(!self.all());
51 return firstBit(~self.bits);
52 }
53
54 pub fn get64(self: *const @This()) u64 {
55 return self.bits;
56 }
57
58 pub fn foreach(self: *const @This(), function: anytype) void {
59 visitWord(self.bits, 0, function);
60 }
61
62 pub fn count(self: *const @This()) usize {
63 return @popCount(self.bits);
64 }
65 };
66
67 pub fn BitSet(comptime size: usize) type {
68 requireCapacity(size);
69 const word_count = wordsFor(size);
70
71 return struct {
72 words: [word_count]BitSet64 = @as([word_count]BitSet64, @splat(.{})),
73
74 const Self = @This();
75 pub const capacity: usize = size;
76 pub const max_size: usize = capacity;
77
78 pub fn maxSize(_: *const Self) usize {
79 return max_size;
80 }
81
82 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 }
87
88 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 }
93
94 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 }
98
99 pub fn any(self: *const Self) bool {
100 for (&self.words) |*word| {
101 if (word.any()) return true;
102 }
103 return false;
104 }
105
106 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 else
114 self.words[word_count - 1].count() == remainder;
115 }
116
117 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 }
124
125 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 }
136
137 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 }
142
143 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 }
151
152 pub 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);
158
159 return struct {
160 words: [word_count]AtomicWord = @as([word_count]AtomicWord, @splat(AtomicWord.init(0))),
161
162 const Self = @This();
163 pub const capacity: usize = size;
164 pub const max_size: usize = capacity;
165
166 pub fn maxSize(_: *const Self) usize {
167 return max_size;
168 }
169
170 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 }
177
178 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 }
186
187 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 }
192
193 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 }
199
200 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 else
209 @popCount(last) == remainder;
210 }
211
212 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 }
222
223 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 }
236
237 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 }
246
247 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 }
253
254 fn atomicWordIndex(index: usize) usize {
255 return index / word_bits;
256 }
257
258 fn atomicBitIndex(index: usize) usize {
259 return index % word_bits;
260 }
261
262 fn atomicBit(index: usize) Word {
263 std.debug.assert(index < word_bits);
264 return @as(Word, 1) << @intCast(index);
265 }
266
267 fn atomicFirstBit(bits: Word) usize {
268 std.debug.assert(bits != 0);
269 return @ctz(bits);
270 }
271
272 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 }
282
283 pub 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);
287
288 return struct {
289 nonzero: BitSet64 = .{},
290 words: [word_count]BitSet64 = @as([word_count]BitSet64, @splat(.{})),
291
292 const Self = @This();
293 pub const capacity: usize = size;
294 pub const max_size: usize = capacity;
295
296 pub fn maxSize(_: *const Self) usize {
297 return max_size;
298 }
299
300 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 }
307
308 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 }
314
315 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 }
322
323 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 }
327
328 pub fn any(self: *const Self) bool {
329 return self.nonzero.any();
330 }
331
332 pub fn all(self: *const Self) bool {
333 if (self.nonzero.count() != word_count) return false;
334 return self.count() == max_size;
335 }
336
337 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 }
342
343 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 }
354
355 pub fn get64(self: *const Self) u64 {
356 return self.words[0].get64();
357 }
358
359 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 }
367
368 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 }
381
382 pub const DefaultBitSet4096 = BitSet4096(4096);
383
384 fn 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 }
392
393 fn firstBit(bits: u64) usize {
394 std.debug.assert(bits != 0);
395 return @ctz(bits);
396 }
397
398 fn bit(index: usize) u64 {
399 std.debug.assert(index < 64);
400 return @as(u64, 1) << @intCast(index);
401 }
402
403 fn wordIndex(index: usize) usize {
404 return index / 64;
405 }
406
407 fn bitIndex(index: usize) usize {
408 return index % 64;
409 }
410
411 fn wordsFor(max_size: usize) usize {
412 return (max_size + 63) / 64;
413 }
414
415 fn validBits(count: usize) u64 {
416 std.debug.assert(count <= 64);
417 return if (count == 64)
418 std.math.maxInt(u64)
419 else
420 (@as(u64, 1) << @intCast(count)) - 1;
421 }
422
423 fn requireCapacity(comptime max_size: usize) void {
424 if (max_size == 0) @compileError("bit sets require a nonzero capacity");
425 }
426
427 const Collector = struct {
428 values: []usize,
429 count: usize = 0,
430
431 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 };
437
438 fn 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());
451
452 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());
464
465 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 }
472
473 fn 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 }
487
488 fn 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());
503
504 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 }
516
517 fn 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 }
537
538 const Mutate64 = struct {
539 set: *BitSet64,
540 visited: *BitSet64,
541
542 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 };
550
551 fn MutateFuture(comptime Set: type) type {
552 return struct {
553 set: *Set,
554 visited: *Set,
555
556 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 }
565
566 fn 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 }
577
578 test "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 }
595
596 test "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 }
607
608 test "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 }
622
623 test "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 }