🥇 bitset: add First and FlipFirst

Author: database64128

bitset/bitset.go

@@ -34,6 +34,30 @@ func (s BitSet) Count() (count uint) {
 	return
 }
 
+// First returns the index of the first set bit.
+func (s BitSet) First() (index uint, found bool) {
+	for i, block := range s.blocks {
+		if block == 0 {
+			continue
+		}
+		return uint(i*blockBits + bits.TrailingZeros(block)), true
+	}
+	return 0, false
+}
+
+// FlipFirst unsets the first set bit and returns its index.
+func (s BitSet) FlipFirst() (index uint, found bool) {
+	for i, block := range s.blocks {
+		if block == 0 {
+			continue
+		}
+		bitIndex := bits.TrailingZeros(block)
+		s.blocks[i] ^= 1 << bitIndex
+		return uint(i*blockBits + bitIndex), true
+	}
+	return 0, false
+}
+
 func (s BitSet) checkIndex(index uint) {
 	if index >= s.capacity {
 		panic(fmt.Sprintf("bitset: index out of range [%d] with capacity %d", index, s.capacity))

bitset/bitset_test.go

@@ -23,47 +23,53 @@ func testBitSetIndexOutOfRange(t *testing.T, s BitSet) {
 func assertEmptyBitSet(t *testing.T, s BitSet) {
 	t.Helper()
 	if count := s.Count(); count != 0 {
-		t.Errorf("expected count to be 0, got %d", count)
+		t.Errorf("s.Count() = %d, want 0", count)
+	}
+	if first, found := s.First(); first != 0 || found {
+		t.Errorf("s.First() = (%d, %t), want (0, false)", first, found)
+	}
+	if index, found := s.FlipFirst(); index != 0 || found {
+		t.Errorf("s.FlipFirst() = (%d, %t), want (0, false)", index, found)
 	}
 	for i := range s.Capacity() {
 		if s.IsSet(i) {
-			t.Errorf("bit %d is unexpectedly set", i)
+			t.Errorf("s.IsSet(%d) = true, want false", i)
 		}
 	}
 }
 
 func assertOddBitSet(t *testing.T, s BitSet) {
 	t.Helper()
 	if count, expectedCount := s.Count(), s.Capacity()/2; count != expectedCount {
-		t.Errorf("expected count to be %d, got %d", expectedCount, count)
+		t.Errorf("s.Count() = %d, want %d", count, expectedCount)
 	}
 	for i := range s.Capacity() {
-		if i%2 == 0 == s.IsSet(i) {
-			t.Errorf("unexpected bit %d", i)
+		if got, want := s.IsSet(i), i%2 == 1; got != want {
+			t.Errorf("s.IsSet(%d) = %t, want %t", i, got, want)
 		}
 	}
 }
 
 func assertEvenBitSet(t *testing.T, s BitSet) {
 	t.Helper()
 	if count, expectedCount := s.Count(), (s.Capacity()+1)/2; count != expectedCount {
-		t.Errorf("expected count to be %d, got %d", expectedCount, count)
+		t.Errorf("s.Count() = %d, want %d", count, expectedCount)
 	}
 	for i := range s.Capacity() {
-		if i%2 == 1 == s.IsSet(i) {
-			t.Errorf("unexpected bit %d", i)
+		if got, want := s.IsSet(i), i%2 == 0; got != want {
+			t.Errorf("s.IsSet(%d) = %t, want %t", i, got, want)
 		}
 	}
 }
 
 func assertFullBitSet(t *testing.T, s BitSet) {
 	t.Helper()
 	if count := s.Count(); count != s.Capacity() {
-		t.Errorf("expected count to be %d, got %d", s.Capacity(), count)
+		t.Errorf("s.Count() = %d, want %d", count, s.Capacity())
 	}
 	for i := range s.Capacity() {
 		if !s.IsSet(i) {
-			t.Errorf("bit %d is unexpectedly unset", i)
+			t.Errorf("s.IsSet(%d) = false, want true", i)
 		}
 	}
 }
@@ -106,6 +112,32 @@ func fillBitSet(t *testing.T, s BitSet) {
 	}
 }
 
+func testBitSetFirst(t *testing.T, s BitSet) {
+	if first, found := s.First(); first != 0 || found {
+		t.Errorf("s.First() = (%d, %t), want (0, false)", first, found)
+	}
+	if index, found := s.FlipFirst(); index != 0 || found {
+		t.Errorf("s.FlipFirst() = (%d, %t), want (0, false)", index, found)
+	}
+
+	for i := s.Capacity() - 1; i < s.Capacity(); i-- {
+		s.Set(i)
+		if first, found := s.First(); first != i || !found {
+			t.Errorf("s.First() = (%d, %t), want (%d, true)", first, found, i)
+		}
+	}
+
+	assertFullBitSet(t, s)
+
+	for i := range s.Capacity() {
+		if index, found := s.FlipFirst(); index != i || !found {
+			t.Errorf("s.FlipFirst() = (%d, %t), want (%d, true)", index, found, i)
+		}
+	}
+
+	assertEmptyBitSet(t, s)
+}
+
 func testBitSetSetAll(t *testing.T, s BitSet) {
 	clearBitSet(t, s)
 	s.SetAll()
@@ -194,6 +226,9 @@ func TestBitSet(t *testing.T) {
 			t.Run("IndexOutOfRange", func(t *testing.T) {
 				testBitSetIndexOutOfRange(t, s)
 			})
+			t.Run("First", func(t *testing.T) {
+				testBitSetFirst(t, s)
+			})
 			t.Run("SetAll", func(t *testing.T) {
 				testBitSetSetAll(t, s)
 			})