From a238afb21a75a0ed584d8c82c3164cba034d0792 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@seebs.net>
Date: Mon, 25 Jan 2021 14:50:37 -0600
Subject: [PATCH] Handle BitmapPtr cells in countRange

We need to be able to count bits in BitmapPtr containers. This only
comes up if you have a non-container-aligned range count, which we
never do in real production yet, but the API allows it so it should
work. In order to do this, we need to provide the tx to countRange
so it can grab pages as needed. Arguably, we should be able to avoid
actually creating/copying that page since we're only using it
internally, never returning it, but this is a pretty rare case
and probably not performance-critical.
---
 rbf/rbf.go | 6 +++++-
 rbf/tx.go  | 7 +++----
 2 files changed, 8 insertions(+), 5 deletions(-)

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
 		}
 	}