authorgravatar for pyrogx1133@gmail.comPyrolistical <pyrogx1133@gmail.com> 2022-12-11 09:10:54-08:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2022-12-11 19:10:54+02:00
logcd9af0f286725c19f5dc0335fb1db1ab8cf7af0d
treec6375e2d4a8751a329f4dd7fe858d298a5d2f67e
parent05890a12f532ba9d58904a14381ec174b9efe473
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

std: add EnumMultiSet

1 files changed, 340 insertions(+), 0 deletions(-)

lib/std/enums.zig+340
...@@ -304,6 +304,346 @@ pub fn EnumMap(comptime E: type, comptime V: type) type {...@@ -304,6 +304,346 @@ pub fn EnumMap(comptime E: type, comptime V: type) type {
304 return IndexedMap(EnumIndexer(E), V, mixin.EnumMapExt);304 return IndexedMap(EnumIndexer(E), V, mixin.EnumMapExt);
305}305}
306306
307/// A multiset of enum elements up to a count of usize. Backed
308/// by an EnumArray. This type does no dynamic allocation and can
309/// be copied by value.
310pub fn EnumMultiset(comptime E: type) type {
311 return BoundedEnumMultiset(E, usize);
312}
313
314/// A multiset of enum elements up to CountSize. Backed by an
315/// EnumArray. This type does no dynamic allocation and can be
316/// copied by value.
317pub fn BoundedEnumMultiset(comptime E: type, comptime CountSize: type) type {
318 return struct {
319 const Self = @This();
320
321 counts: EnumArray(E, CountSize),
322
323 /// Initializes the multiset using a struct of counts.
324 pub fn init(init_counts: EnumFieldStruct(E, CountSize, 0)) Self {
325 var self = initWithCount(0);
326 inline for (@typeInfo(E).Enum.fields) |field| {
327 const c = @field(init_counts, field.name);
328 const key = @intToEnum(E, field.value);
329 self.counts.set(key, c);
330 }
331 return self;
332 }
333
334 /// Initializes the multiset with a count of zero.
335 pub fn initEmpty() Self {
336 return initWithCount(0);
337 }
338
339 /// Initializes the multiset with all keys at the
340 /// same count.
341 pub fn initWithCount(comptime c: CountSize) Self {
342 return .{
343 .counts = EnumArray(E, CountSize).initDefault(c, .{}),
344 };
345 }
346
347 /// Returns the total number of key counts in the multiset.
348 pub fn count(self: Self) usize {
349 var sum: usize = 0;
350 for (self.counts.values) |c| {
351 sum += c;
352 }
353 return sum;
354 }
355
356 /// Checks if at least one key in multiset.
357 pub fn contains(self: Self, key: E) bool {
358 return self.counts.get(key) > 0;
359 }
360
361 /// Removes all instance of a key from multiset. Same as
362 /// setCount(key, 0).
363 pub fn removeAll(self: *Self, key: E) void {
364 return self.counts.set(key, 0);
365 }
366
367 /// Increases the key count by given amount. Caller asserts
368 /// operation will not overflow.
369 pub fn addAssertSafe(self: *Self, key: E, c: CountSize) void {
370 self.counts.getPtr(key).* += c;
371 }
372
373 /// Increases the key count by given amount.
374 pub fn add(self: *Self, key: E, c: CountSize) error{Overflow}!void {
375 self.counts.set(key, try std.math.add(CountSize, self.counts.get(key), c));
376 }
377
378 /// Decreases the key count by given amount. If amount is
379 /// greater than the number of keys in multset, then key count
380 /// will be set to zero.
381 pub fn remove(self: *Self, key: E, c: CountSize) void {
382 self.counts.getPtr(key).* -= @min(self.getCount(key), c);
383 }
384
385 /// Returns the count for a key.
386 pub fn getCount(self: Self, key: E) CountSize {
387 return self.counts.get(key);
388 }
389
390 /// Set the count for a key.
391 pub fn setCount(self: *Self, key: E, c: CountSize) void {
392 self.counts.set(key, c);
393 }
394
395 /// Increases the all key counts by given multiset. Caller
396 /// asserts operation will not overflow any key.
397 pub fn addSetAssertSafe(self: *Self, other: Self) void {
398 inline for (@typeInfo(E).Enum.fields) |field| {
399 const key = @intToEnum(E, field.value);
400 self.addAssertSafe(key, other.getCount(key));
401 }
402 }
403
404 /// Increases the all key counts by given multiset.
405 pub fn addSet(self: *Self, other: Self) error{Overflow}!void {
406 inline for (@typeInfo(E).Enum.fields) |field| {
407 const key = @intToEnum(E, field.value);
408 try self.add(key, other.getCount(key));
409 }
410 }
411
412 /// Deccreases the all key counts by given multiset. If
413 /// the given multiset has more key counts than this,
414 /// then that key will have a key count of zero.
415 pub fn removeSet(self: *Self, other: Self) void {
416 inline for (@typeInfo(E).Enum.fields) |field| {
417 const key = @intToEnum(E, field.value);
418 self.remove(key, other.getCount(key));
419 }
420 }
421
422 /// Returns true iff all key counts are the same as
423 /// given multiset.
424 pub fn eql(self: Self, other: Self) bool {
425 inline for (@typeInfo(E).Enum.fields) |field| {
426 const key = @intToEnum(E, field.value);
427 if (self.getCount(key) != other.getCount(key)) {
428 return false;
429 }
430 }
431 return true;
432 }
433
434 /// Returns a multiset with the total key count of this
435 /// multiset and the other multiset. Caller asserts
436 /// operation will not overflow any key.
437 pub fn plusAssertSafe(self: Self, other: Self) Self {
438 var result = self;
439 result.addSetAssertSafe(other);
440 return result;
441 }
442
443 /// Returns a multiset with the total key count of this
444 /// multiset and the other multiset.
445 pub fn plus(self: Self, other: Self) error{Overflow}!Self {
446 var result = self;
447 try result.addSet(other);
448 return result;
449 }
450
451 /// Returns a multiset with the key count of this
452 /// multiset minus the corresponding key count in the
453 /// other multiset. If the other multiset contains
454 /// more key count than this set, that key will have
455 /// a count of zero.
456 pub fn minus(self: Self, other: Self) Self {
457 var result = self;
458 result.removeSet(other);
459 return result;
460 }
461
462 pub const Entry = EnumArray(E, CountSize).Entry;
463 pub const Iterator = EnumArray(E, CountSize).Iterator;
464
465 /// Returns an iterator over this multiset. Keys with zero
466 /// counts are included. Modifications to the set during
467 /// iteration may or may not be observed by the iterator,
468 /// but will not invalidate it.
469 pub fn iterator(self: *Self) Iterator {
470 return self.counts.iterator();
471 }
472 };
473}
474
475test "EnumMultiset" {
476 const Ball = enum { red, green, blue };
477
478 const empty = EnumMultiset(Ball).initEmpty();
479 const r0_g1_b2 = EnumMultiset(Ball).init(.{
480 .red = 0,
481 .green = 1,
482 .blue = 2,
483 });
484 const ten_of_each = EnumMultiset(Ball).initWithCount(10);
485
486 try testing.expectEqual(empty.count(), 0);
487 try testing.expectEqual(r0_g1_b2.count(), 3);
488 try testing.expectEqual(ten_of_each.count(), 30);
489
490 try testing.expect(!empty.contains(.red));
491 try testing.expect(!empty.contains(.green));
492 try testing.expect(!empty.contains(.blue));
493
494 try testing.expect(!r0_g1_b2.contains(.red));
495 try testing.expect(r0_g1_b2.contains(.green));
496 try testing.expect(r0_g1_b2.contains(.blue));
497
498 try testing.expect(ten_of_each.contains(.red));
499 try testing.expect(ten_of_each.contains(.green));
500 try testing.expect(ten_of_each.contains(.blue));
501
502 {
503 var copy = ten_of_each;
504 copy.removeAll(.red);
505 try testing.expect(!copy.contains(.red));
506
507 // removeAll second time does nothing
508 copy.removeAll(.red);
509 try testing.expect(!copy.contains(.red));
510 }
511
512 {
513 var copy = ten_of_each;
514 copy.addAssertSafe(.red, 6);
515 try testing.expectEqual(copy.getCount(.red), 16);
516 }
517
518 {
519 var copy = ten_of_each;
520 try copy.add(.red, 6);
521 try testing.expectEqual(copy.getCount(.red), 16);
522
523 try testing.expectError(error.Overflow, copy.add(.red, std.math.maxInt(usize)));
524 }
525
526 {
527 var copy = ten_of_each;
528 copy.remove(.red, 4);
529 try testing.expectEqual(copy.getCount(.red), 6);
530
531 // subtracting more it contains does not underflow
532 copy.remove(.green, 14);
533 try testing.expectEqual(copy.getCount(.green), 0);
534 }
535
536 try testing.expectEqual(empty.getCount(.green), 0);
537 try testing.expectEqual(r0_g1_b2.getCount(.green), 1);
538 try testing.expectEqual(ten_of_each.getCount(.green), 10);
539
540 {
541 var copy = empty;
542 copy.setCount(.red, 6);
543 try testing.expectEqual(copy.getCount(.red), 6);
544 }
545
546 {
547 var copy = r0_g1_b2;
548 copy.addSetAssertSafe(ten_of_each);
549 try testing.expectEqual(copy.getCount(.red), 10);
550 try testing.expectEqual(copy.getCount(.green), 11);
551 try testing.expectEqual(copy.getCount(.blue), 12);
552 }
553
554 {
555 var copy = r0_g1_b2;
556 try copy.addSet(ten_of_each);
557 try testing.expectEqual(copy.getCount(.red), 10);
558 try testing.expectEqual(copy.getCount(.green), 11);
559 try testing.expectEqual(copy.getCount(.blue), 12);
560
561 const full = EnumMultiset(Ball).initWithCount(std.math.maxInt(usize));
562 try testing.expectError(error.Overflow, copy.addSet(full));
563 }
564
565 {
566 var copy = ten_of_each;
567 copy.removeSet(r0_g1_b2);
568 try testing.expectEqual(copy.getCount(.red), 10);
569 try testing.expectEqual(copy.getCount(.green), 9);
570 try testing.expectEqual(copy.getCount(.blue), 8);
571
572 copy.removeSet(ten_of_each);
573 try testing.expectEqual(copy.getCount(.red), 0);
574 try testing.expectEqual(copy.getCount(.green), 0);
575 try testing.expectEqual(copy.getCount(.blue), 0);
576 }
577
578 try testing.expect(empty.eql(empty));
579 try testing.expect(r0_g1_b2.eql(r0_g1_b2));
580 try testing.expect(ten_of_each.eql(ten_of_each));
581 try testing.expect(!empty.eql(r0_g1_b2));
582 try testing.expect(!r0_g1_b2.eql(ten_of_each));
583 try testing.expect(!ten_of_each.eql(empty));
584
585 {
586 const result = r0_g1_b2.plusAssertSafe(ten_of_each);
587 try testing.expectEqual(result.getCount(.red), 10);
588 try testing.expectEqual(result.getCount(.green), 11);
589 try testing.expectEqual(result.getCount(.blue), 12);
590 }
591
592 {
593 const result = try r0_g1_b2.plus(ten_of_each);
594 try testing.expectEqual(result.getCount(.red), 10);
595 try testing.expectEqual(result.getCount(.green), 11);
596 try testing.expectEqual(result.getCount(.blue), 12);
597
598 const full = EnumMultiset(Ball).initWithCount(std.math.maxInt(usize));
599 try testing.expectError(error.Overflow, result.plus(full));
600 }
601
602 {
603 const result = ten_of_each.minus(r0_g1_b2);
604 try testing.expectEqual(result.getCount(.red), 10);
605 try testing.expectEqual(result.getCount(.green), 9);
606 try testing.expectEqual(result.getCount(.blue), 8);
607 }
608
609 {
610 const result = ten_of_each.minus(r0_g1_b2).minus(ten_of_each);
611 try testing.expectEqual(result.getCount(.red), 0);
612 try testing.expectEqual(result.getCount(.green), 0);
613 try testing.expectEqual(result.getCount(.blue), 0);
614 }
615
616 {
617 var copy = empty;
618 var it = copy.iterator();
619 var entry = it.next().?;
620 try testing.expectEqual(entry.key, .red);
621 try testing.expectEqual(entry.value.*, 0);
622 entry = it.next().?;
623 try testing.expectEqual(entry.key, .green);
624 try testing.expectEqual(entry.value.*, 0);
625 entry = it.next().?;
626 try testing.expectEqual(entry.key, .blue);
627 try testing.expectEqual(entry.value.*, 0);
628 try testing.expectEqual(it.next(), null);
629 }
630
631 {
632 var copy = r0_g1_b2;
633 var it = copy.iterator();
634 var entry = it.next().?;
635 try testing.expectEqual(entry.key, .red);
636 try testing.expectEqual(entry.value.*, 0);
637 entry = it.next().?;
638 try testing.expectEqual(entry.key, .green);
639 try testing.expectEqual(entry.value.*, 1);
640 entry = it.next().?;
641 try testing.expectEqual(entry.key, .blue);
642 try testing.expectEqual(entry.value.*, 2);
643 try testing.expectEqual(it.next(), null);
644 }
645}
646
307/// An array keyed by an enum, backed by a dense array.647/// An array keyed by an enum, backed by a dense array.
308/// If the enum is not dense, a mapping will be constructed from648/// If the enum is not dense, a mapping will be constructed from
309/// enum values to dense indices. This type does no dynamic649/// enum values to dense indices. This type does no dynamic