impl. AppendTo and AsSlice

bitset.go

@@ -520,6 +520,45 @@ func (b *BitSet) DeleteAt(i uint) *BitSet {
 	return b
 }
 
+// AppendTo appends all set bits to buf and returns the (maybe extended) buf.
+// In case of allocation failure, the function will panic.
+//
+// See also [BitSet.AsSlice] and [BitSet.NextSetMany].
+func (b *BitSet) AppendTo(buf []uint) []uint {
+	for idx, word := range b.set {
+		for word != 0 {
+			buf = append(buf, uint(idx<<log2WordSize+bits.TrailingZeros64(word)))
+
+			// clear the rightmost set bit
+			word &= word - 1
+		}
+	}
+
+	return buf
+}
+
+// AsSlice returns all set bits as slice.
+// It panics if the capacity of buf is < b.Count()
+//
+// See also [BitSet.AppendTo] and [BitSet.NextSetMany].
+func (b *BitSet) AsSlice(buf []uint) []uint {
+	buf = buf[:cap(buf)] // len = cap
+
+	size := 0
+	for idx, word := range b.set {
+		for ; word != 0; size++ {
+			// panics if capacity of buf is exceeded.
+			buf[size] = uint(idx<<log2WordSize + bits.TrailingZeros64(word))
+
+			// clear the rightmost set bit
+			word &= word - 1
+		}
+	}
+
+	buf = buf[:size]
+	return buf
+}
+
 // NextSet returns the next bit set from the specified index,
 // including possibly the current index
 // along with an error code (true = valid, false = no set bit found)
@@ -570,8 +609,11 @@ func (b *BitSet) NextSet(i uint) (uint, bool) {
 //	indices := make([]uint, bitmap.Count())
 //	bitmap.NextSetMany(0, indices)
 //
-// However if bitmap.Count() is large, it might be preferable to
-// use several calls to NextSetMany, for performance reasons.
+// It is also possible to retrieve all set bits with [BitSet.AppendTo]
+// or [BitSet.AsSlice].
+//
+// However if Count() is large, it might be preferable to
+// use several calls to NextSetMany for memory reasons.
 func (b *BitSet) NextSetMany(i uint, buffer []uint) (uint, []uint) {
 	capacity := cap(buffer)
 	result := buffer[:capacity]

bitset_test.go

@@ -507,6 +507,104 @@ func TestNextSetMany(t *testing.T) {
 	}
 }
 
+func TestAppendTo(t *testing.T) {
+	testCases := []struct {
+		name string
+		set  []uint
+		buf  []uint
+	}{
+		{
+			name: "null",
+			set:  nil,
+			buf:  nil,
+		},
+		{
+			name: "one",
+			set:  []uint{42},
+			buf:  make([]uint, 0, 5),
+		},
+		{
+			name: "many",
+			set:  []uint{1, 42, 55, 258, 7211, 54666},
+			buf:  make([]uint, 0),
+		},
+	}
+
+	for _, tc := range testCases {
+		var b BitSet
+		for _, u := range tc.set {
+			b.Set(u)
+		}
+
+		tc.buf = b.AppendTo(tc.buf)
+
+		if !reflect.DeepEqual(tc.buf, tc.set) {
+			t.Errorf("AppendTo, %s: returned buf is not equal as expected:\ngot:  %v\nwant: %v",
+				tc.name, tc.buf, tc.set)
+		}
+	}
+}
+
+func TestAsSlice(t *testing.T) {
+	testCases := []struct {
+		name string
+		set  []uint
+		buf  []uint
+	}{
+		{
+			name: "null",
+			set:  nil,
+			buf:  nil,
+		},
+		{
+			name: "one",
+			set:  []uint{42},
+			buf:  make([]uint, 1),
+		},
+		{
+			name: "many",
+			set:  []uint{1, 42, 55, 258, 7211, 54666},
+			buf:  make([]uint, 6),
+		},
+	}
+
+	for _, tc := range testCases {
+		var b BitSet
+		for _, u := range tc.set {
+			b.Set(u)
+		}
+
+		tc.buf = b.AsSlice(tc.buf)
+
+		if !reflect.DeepEqual(tc.buf, tc.set) {
+			t.Errorf("AsSlice, %s: returned buf is not equal as expected:\ngot:  %v\nwant: %v",
+				tc.name, tc.buf, tc.set)
+		}
+	}
+}
+
+func TestPanicAppendTo(t *testing.T) {
+	defer func() {
+		if r := recover(); r != nil {
+			t.Error("AppendTo with empty buf should not have caused a panic")
+		}
+	}()
+	v := New(1000)
+	v.Set(1000)
+	_ = v.AppendTo(nil)
+}
+
+func TestPanicAsSlice(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Error("AsSlice with buf too small should have caused a panic")
+		}
+	}()
+	v := New(1000)
+	v.Set(1000)
+	_ = v.AsSlice(nil)
+}
+
 func TestSetTo(t *testing.T) {
 	v := New(1000)
 	v.SetTo(100, true)