Merge pull request #1367 from seebs/countRange

RBF countRange: Handle partial counts on bitmapPtr containers

rbf/rbf.go

@@ -462,7 +462,7 @@ func (c *leafCell) lastValue(tx *Tx) uint16 {
 // We have to take int32 rather than uint16 because the interval is [start, end),
 // and otherwise we have no way to ask to count the entire container (the
 // high bit will be missed).
-func (c *leafCell) countRange(start, end int32) (n int) {
+func (c *leafCell) countRange(tx *Tx, start, end int32) (n int) {
 	// If the full range is being queried, simply use the precalculated count.
 	if start == 0 && end > math.MaxUint16 {
 		return c.BitN
@@ -475,6 +475,10 @@ func (c *leafCell) countRange(start, end int32) (n int) {
 		return int(roaring.RunCountRange(toInterval16(c.Data), start, end))
 	case ContainerTypeBitmap:
 		return int(roaring.BitmapCountRange(toArray64(c.Data), start, end))
+	case ContainerTypeBitmapPtr:
+		_, a, err := tx.leafCellBitmap(toPgno(c.Data))
+		panicOn(err)
+		return int(roaring.BitmapCountRange(a, start, end))
 	default:
 		panic(fmt.Sprintf("invalid container type: %d", c.Type))
 	}

rbf/tx.go

@@ -1318,7 +1318,6 @@ func (tx *Tx) CountRange(name string, start, end uint64) (uint64, error) {
 	} else if err != nil {
 		return 0, err
 	}
-
 	var n uint64
 	for {
 		if err := csr.Next(); err == io.EOF {
@@ -1341,7 +1340,7 @@ func (tx *Tx) CountRange(name string, start, end uint64) (uint64, error) {
 
 		// If range is entirely in one container then just count that range.
 		if skey == ekey {
-			return uint64(c.countRange(int32(lowbits(start)), ebits)), nil
+			return uint64(c.countRange(tx, int32(lowbits(start)), ebits)), nil
 		}
 		// INVAR: skey < ekey
 
@@ -1351,7 +1350,7 @@ func (tx *Tx) CountRange(name string, start, end uint64) (uint64, error) {
 			break
 		}
 		if k == skey {
-			n += uint64(c.countRange(int32(lowbits(start)), roaring.MaxContainerVal+1))
+			n += uint64(c.countRange(tx, int32(lowbits(start)), roaring.MaxContainerVal+1))
 			continue
 		}
 		if k < ekey {
@@ -1359,7 +1358,7 @@ func (tx *Tx) CountRange(name string, start, end uint64) (uint64, error) {
 			continue
 		}
 		if k == ekey && ebits > 0 {
-			n += uint64(c.countRange(0, ebits))
+			n += uint64(c.countRange(tx, 0, ebits))
 			break
 		}
 	}

tx_internal_test.go

@@ -37,20 +37,29 @@ func requireCountRangeSampleData(tb testing.TB) (*fragment, Tx) {
 		// request that each container get its own copy of the bitmap.
 		var bitmapSample [1025]uint64
 		for i := range arraySample {
-			arraySample[i] = uint16(i)
+			arraySample[i] = uint16(i * 2)
 		}
-		for i := 0; i < 4096/64; i++ {
-			bitmapSample[i] = ^uint64(0)
+		// Put corresponding bits in the bitmap...
+		for i := 0; i < 4096/32; i++ {
+			// bit 0 is 0x1, bit 2 is 0x4, so even-numbered bits
+			// are 0x5555....
+			bitmapSample[i] = 0x5555555555555555
 		}
 		bm := roaring.NewSliceBitmap()
 		for n := 0; n < 4096 && n < countRangeMaxN; n++ {
 			c := roaring.NewContainerArray(arraySample[:n])
 			bm.Put(uint64(n), c)
 		}
-		for n := 4096; n < countRangeMaxN; n++ {
+		// Start filling in the missing bits. This starts us out with
+		// bitmap containers, but then eventually converts to things
+		// that are more likely to be run containers. At the end of this,
+		// we should have exactly the first 8,192 bits set, for a single
+		// run of 8k.
+		for n := 4096; n < 8192; n++ {
 			c := roaring.NewContainerBitmapN(bitmapSample[:], int32(n))
 			bm.Put(uint64(n), c)
-			bitmapSample[n/64] |= 1 << (n % 64)
+			w := n - 4096
+			bitmapSample[w/32] |= 1 << (((n % 32) * 2) + 1)
 		}
 		var asBytes bytes.Buffer
 		n, err := bm.WriteTo(&asBytes)
@@ -90,11 +99,14 @@ func TestTx_CountRange(t *testing.T) {
 	expected := uint64(0)
 	j := uint64(0)
 	for i := uint64(0); i < countRangeMaxN; i += 7 {
+		expected += i
 		if i%4 == 3 {
 			expected -= (j * 7) + 21
 			j += 7
 		}
-		got, err := tx.CountRange("i", "f", viewStandard, 0, uint64(j)<<16, uint64(i)<<16)
+		// Every other bit gets set, for a total of i bits in container
+		// i, so they're all in the first (i*2) bits of the container.
+		got, err := tx.CountRange("i", "f", viewStandard, 0, uint64(j)<<16, (uint64(i)<<16)+(i*2))
 		if err != nil {
 			t.Fatalf("counting range: %v", err)
 		}
@@ -102,7 +114,8 @@ func TestTx_CountRange(t *testing.T) {
 			t.Fatalf("counting from container %d to %d, expected %d, got %d",
 				j, i, expected, got)
 		}
-		expected += (i * 7) + 21
+		// The -i here undoes the +i at the top of this loop.
+		expected += (i * 7) + 21 - i
 	}
 }