new file mode 100644

@@ -0,0 +1 @@

+/testdata/rapid/**

@@ -6,6 +6,7 @@ import (

 	"encoding/json"

 	"errors"

 	"fmt"

+	"math/bits"

 )

 // block sequence constants

@@ -176,6 +177,52 @@ func (s *sequence) fromByteArray(data []byte) error {

 	return nil

 }

+// OnesCount calculates the number of selected bits in the range [start, end].

+func (h *Bitmap) OnesCount(start, end uint64) (uint64, error) {

+	if end < start || end >= h.bits {

+		return 0, fmt.Errorf("invalid bit range [%d, %d]", start, end)

+	}

+

+	// Account for the starting ordinal being partway into a block: count

+	// the sequence element's block once with a bitmask applied, then count

+	// the remaining repeats of the sequence element without the mask.

+	current, _, precBlocks, _ := findSequence(h.head, start/8)

+	var (

+		blocksToCount = end/blockLen - start/blockLen + 1

+		curblocks     = min(current.count-precBlocks, blocksToCount)

+		mask          = blockMAX >> (start % blockLen)

+		runlen        = uint64(1)

+		count         = uint64(0)

+	)

+	for blocksToCount > 0 {

+		if blocksToCount == 1 {

+			// We're counting the last block.

+			// (Which could be the same as the first block.)

+			// Mask off the bits beyond the end ordinal.

+			mask &= blockMAX << (blockLen - end%blockLen - 1)

+		}

+		count += uint64(bits.OnesCount32(current.block&mask)) * runlen

+		mask = blockMAX

+		blocksToCount -= runlen

+		curblocks -= runlen

+		if curblocks == 0 {

+			current = current.next

+			if current == nil {

+				break

+			}

+			curblocks = min(current.count, blocksToCount)

+		}

+		// If the block containing the end ordinal is a repeat of the

+		// current block, split the counting across two loop iterations.

+		// The final repeat of the block needs to be counted with a

+		// bitmask applied so we do not count bits beyond the end

+		// ordinal.

+		runlen = max(1, min(curblocks, blocksToCount-1))

+	}

+

+	return count, nil

+}

+

 // SetAnyInRange sets the first unset bit in the range [start, end] and returns

 // the ordinal of the set bit.

 //

@@ -5,6 +5,10 @@ import (

 	"math/rand"

 	"testing"

 	"time"

+

+	"gotest.tools/v3/assert"

+	is "gotest.tools/v3/assert/cmp"

+	"pgregory.net/rapid"

 )

 func TestSequenceGetAvailableBit(t *testing.T) {

@@ -1231,3 +1235,80 @@ func TestMarshalJSON(t *testing.T) {

 		})

 	}

 }

+

+func TestOnesCount(t *testing.T) {

+	rapid.Check(t, func(t *rapid.T) {

+		bm := New(rapid.Uint64Range(10, 1000).Draw(t, "capacity"))

+		ordinals := make([]uint64, bm.Bits())

+		for i := range ordinals {

+			ordinals[i] = uint64(i)

+		}

+		nBitsToSet := rapid.IntRange(0, int(bm.Bits()-1)).Draw(t, "nBitsToSet")

+		selectedOrdinals := rapid.Permutation(ordinals).Draw(t, "selectedOrdinals")[:nBitsToSet]

+		for _, i := range selectedOrdinals {

+			assert.NilError(t, bm.Set(i))

+		}

+		t.Logf("%v", bm)

+

+		got, err := bm.OnesCount(0, bm.Bits()-1)

+		assert.Check(t, err, "OnesCount of all ordinals should succeed")

+		assert.Check(t, is.Equal(got, bm.Bits()-bm.Unselected()), "OnesCount of all ordinals should equal Bits-Unselected")

+

+		idxgen := rapid.Uint64Range(0, bm.Bits()-1)

+		for range 1000 {

+			start, end := idxgen.Draw(t, "start"), idxgen.Draw(t, "end")

+			if start > end {

+				start, end = end, start

+			}

+			var expected uint64

+			for _, i := range selectedOrdinals {

+				if i >= start && i <= end {

+					expected++

+				}

+			}

+			got, err := bm.OnesCount(start, end)

+			assert.NilError(t, err)

+			assert.Check(t, is.Equal(got, expected))

+		}

+	})

+

+	bm := New(8*blockLen + 1)

+	for _, r := range []struct{ from, to uint64 }{

+		// 0b0111111... repeated 3x

+		{1, blockLen - 1},

+		{blockLen + 1, 2*blockLen - 1},

+		{2*blockLen + 1, 3*blockLen - 1},

+

+		// 0b0001111...

+		{3*blockLen + 3, 3*blockLen + 6},

+

+		// 0b111111....., 0b1000000....

+		{5 * blockLen, 6 * blockLen},

+

+		{8 * blockLen, 8 * blockLen},

+	} {

+		for i := r.from; i <= r.to; i++ {

+			assert.NilError(t, bm.Set(i))

+		}

+	}

+	t.Logf("%v", bm)

+

+	expectOnesCount := func(from, to uint64, expected uint64) {

+		t.Helper()

+		got, err := bm.OnesCount(from, to)

+		assert.NilError(t, err)

+		assert.Check(t, is.Equal(got, expected))

+	}

+

+	expectOnesCount(0, bm.Bits()-1, bm.Bits()-bm.Unselected())

+	expectOnesCount(0, 0, 0)

+	expectOnesCount(0, 1, 1)

+	expectOnesCount(1, 1, 1)

+	expectOnesCount(3, 12, 12-3+1)

+	expectOnesCount(blockLen-1, 2*blockLen+1, blockLen+1)

+	expectOnesCount(5, 3*blockLen-3, blockLen-1-4+blockLen-1+blockLen-1-2)

+	expectOnesCount(3*blockLen, 3*blockLen+4, 2)

+	expectOnesCount(3*blockLen, 5*blockLen-1, 4)

+	expectOnesCount(6*blockLen, 8*blockLen-1, 1)

+	expectOnesCount(8*blockLen, 8*blockLen, 1)

+}