From 888040c26c16ddd6b67323f73f90783680038d91 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Mon, 15 Aug 2016 16:03:25 -0600
Subject: [PATCH] add CountRange()

This commit adds `roaring.Bitmap.CountRange()` to return the number
of bits in a subrange of a bitmap. This significantly improves
server start time and is needed for upcoming zero copy bitmap
optimizations.
---
 fragment.go        |  3 +-
 roaring/roaring.go | 80 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/fragment.go b/fragment.go
index 6d32bca9d..7387852ca 100644
--- a/fragment.go
+++ b/fragment.go
@@ -234,7 +234,8 @@ func (f *Fragment) openCache() error {
 	// Read in all bitmaps by ID.
 	// This will cause them to be added to the cache.
 	for _, bitmapID := range pb.GetBitmapIDs() {
-		f.bitmap(bitmapID)
+		n := f.storage.CountRange(bitmapID*SliceWidth, (bitmapID+1)*SliceWidth)
+		f.cache.Add(bitmapID, n)
 	}
 
 	return nil
diff --git a/roaring/roaring.go b/roaring/roaring.go
index f9857d13b..7ea4e9475 100644
--- a/roaring/roaring.go
+++ b/roaring/roaring.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"hash/fnv"
 	"io"
+	"sort"
 	"unsafe"
 )
 
@@ -134,6 +135,41 @@ func (b *Bitmap) Count() (n uint64) {
 	return n
 }
 
+// CountRange returns the number of bits set between [start, end).
+func (b *Bitmap) CountRange(start, end uint64) (n uint64) {
+	i := search64(b.keys, highbits(start))
+	j := search64(b.keys, highbits(end))
+
+	// If range is entirely in one container then just count that range.
+	if i > 0 && i == j {
+		return uint64(b.containers[i].countRange(lowbits(start), lowbits(end)))
+	}
+
+	// Count first partial container.
+	if i < 0 {
+		i = -i
+	} else {
+		n += uint64(b.containers[i].countRange(lowbits(start), (bitmapN*64)+1))
+	}
+
+	// Count last container.
+	if j < 0 {
+		j = -j
+		if j > len(b.containers) {
+			j = len(b.containers)
+		}
+	} else {
+		n += uint64(b.containers[j].countRange(0, lowbits(end)))
+	}
+
+	// Count containers in between.
+	for x := i + 1; x < j; x++ {
+		n += uint64(b.containers[x].n)
+	}
+
+	return n
+}
+
 // Slice returns a slice of all integers in the bitmap.
 func (b *Bitmap) Slice() []uint64 {
 	var a []uint64
@@ -658,6 +694,50 @@ func (c *container) unmap() {
 	c.mapped = false
 }
 
+// countRange counts the number of bits set between [start, end).
+func (c *container) countRange(start, end uint32) (n int) {
+	if c.isArray() {
+		return c.arrayCountRange(start, end)
+	}
+	return c.bitmapCountRange(start, end)
+}
+
+func (c *container) arrayCountRange(start, end uint32) (n int) {
+	i := sort.Search(len(c.array), func(i int) bool { return c.array[i] >= start })
+	for ; i < len(c.array); i++ {
+		v := c.array[i]
+		if v >= end {
+			break
+		}
+		n++
+	}
+	return n
+}
+
+func (c *container) bitmapCountRange(start, end uint32) int {
+	var n uint64
+	i, j := start/64, end/64
+
+	// Count partial starting word.
+	if off := start % 64; off != 0 {
+		n += popcount(c.bitmap[i] << off)
+	}
+
+	// Count words in between.
+	for ; i < j; i++ {
+		n += popcount(c.bitmap[i])
+	}
+
+	// Count partial ending word.
+	if int(j) < len(c.bitmap) {
+		if off := end % 64; off != 0 {
+			n += popcount(c.bitmap[j] >> off)
+		}
+	}
+
+	return int(n)
+}
+
 // add adds a value to the container.
 func (c *container) add(v uint32) bool {
 	if c.isArray() {