diff --git a/rbf/rbf.go b/rbf/rbf.go index f13a5a7fc..365907121 100644 --- a/rbf/rbf.go +++ b/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)) } diff --git a/rbf/tx.go b/rbf/tx.go index 705edfc5b..a952104ec 100644 --- a/rbf/tx.go +++ b/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 } } diff --git a/tx_internal_test.go b/tx_internal_test.go index a5da9cccd..3ff4ddd63 100644 --- a/tx_internal_test.go +++ b/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 } }