Implement GetSetBits for bitArray and sparseBitArray

bitarray/bitarray.go

@@ -20,6 +20,8 @@ efficient way.  This is *NOT* a threadsafe package.
 */
 package bitarray
 
+import "math/bits"
+
 // bitArray is a struct that maintains state of a bit array.
 type bitArray struct {
 	blocks  []block
@@ -117,8 +119,66 @@ func (ba *bitArray) GetBit(k uint64) (bool, error) {
 }
 
 // GetSetBits gets the position of bits set in the array.
-func (ba *bitArray) GetSetBits(_ uint64, buffer []uint64) []uint64 {
-	return buffer[:0]
+func (ba *bitArray) GetSetBits(from uint64, buffer []uint64) []uint64 {
+	fromBlockIndex, fromOffset := getIndexAndRemainder(from)
+	return getSetBitsInBlocks(
+		fromBlockIndex,
+		fromOffset,
+		ba.blocks[fromBlockIndex:],
+		nil,
+		buffer,
+	)
+}
+
+// getSetBitsInBlocks fills a buffer with positions of set bits in the provided blocks. Optionally, indices may be
+// provided for sparse/non-consecutive blocks.
+func getSetBitsInBlocks(
+	fromBlockIndex, fromOffset uint64,
+	blocks []block,
+	indices []uint64,
+	buffer []uint64,
+) []uint64 {
+	bufferCapacity := cap(buffer)
+	results := buffer[:bufferCapacity]
+	resultSize := 0
+
+	for i, block := range blocks {
+		blockIndex := fromBlockIndex + uint64(i)
+		if indices != nil {
+			blockIndex = indices[i]
+		}
+
+		isFirstBlock := blockIndex == fromBlockIndex
+		if isFirstBlock {
+			block >>= fromOffset
+		}
+
+		for block != 0 {
+			trailing := bits.TrailingZeros64(uint64(block))
+
+			if isFirstBlock {
+				results[resultSize] = uint64(trailing) + (blockIndex << 6) + fromOffset
+			} else {
+				results[resultSize] = uint64(trailing) + (blockIndex << 6)
+			}
+			resultSize++
+
+			if resultSize == cap(results) {
+				return results[:resultSize]
+			}
+
+			// Example of this expression:
+			// 	block					01001100
+			// 	^block					10110011
+			// 	(^block) + 1 			10110100
+			// 	block & (^block) + 1	00000100
+			// 	block ^ mask 			01001000
+			mask := block & ((^block) + 1)
+			block = block ^ mask
+		}
+	}
+
+	return results[:resultSize]
 }
 
 // ClearBit will unset a bit at the given index if it is set.

bitarray/sparse_bitarray.go

@@ -128,8 +128,21 @@ func (sba *sparseBitArray) GetBit(k uint64) (bool, error) {
 }
 
 // GetSetBits gets the position of bits set in the array.
-func (sba *sparseBitArray) GetSetBits(_ uint64, buffer []uint64) []uint64 {
-	return buffer[:0]
+func (sba *sparseBitArray) GetSetBits(from uint64, buffer []uint64) []uint64 {
+	fromBlockIndex, fromOffset := getIndexAndRemainder(from)
+
+	fromBlockLocation := sba.indices.search(fromBlockIndex)
+	if int(fromBlockLocation) == len(sba.indices) {
+		return buffer[:0]
+	}
+
+	return getSetBitsInBlocks(
+		fromBlockIndex,
+		fromOffset,
+		sba.blocks[fromBlockLocation:],
+		sba.indices[fromBlockLocation:],
+		buffer,
+	)
 }
 
 // ToNums converts this sparse bitarray to a list of numbers contained