feat(rank): Implement the Rank API (#20)

Implement Rank(uint64) API which returns the rank of a given value in the whole sorted set.

Author: ahsanbarkati

bitmap.go

@@ -464,6 +464,9 @@ func (ra *Bitmap) Select(x uint64) (uint64, error) {
 }
 
 func (ra *Bitmap) Contains(x uint64) bool {
+	if ra == nil {
+		return false
+	}
 	key := x & mask
 	offset, has := ra.keys.getValue(key)
 	if !has {
@@ -484,6 +487,9 @@ func (ra *Bitmap) Contains(x uint64) bool {
 }
 
 func (ra *Bitmap) Remove(x uint64) bool {
+	if ra == nil {
+		return false
+	}
 	key := x & mask
 	offset, has := ra.keys.getValue(key)
 	if !has {
@@ -563,6 +569,17 @@ func (ra *Bitmap) RemoveRange(lo, hi uint64) {
 	ra.Cleanup()
 }
 
+func (ra *Bitmap) Reset() {
+	// reset ra.data to size enough for one container and corresponding key.
+	// 2 u64 is needed for header and another 2 u16 for the key 0.
+	ra.data = ra.data[:16+minContainerSize]
+	ra.keys = toUint64Slice(ra.data)
+
+	offset := ra.newContainer(minContainerSize)
+	ra.keys.setAt(indexNodeStart+1, offset)
+	ra.keys.setNumKeys(1)
+}
+
 func (ra *Bitmap) GetCardinality() int {
 	if ra == nil {
 		return 0
@@ -715,6 +732,11 @@ func (ra *Bitmap) extreme(dir int) uint64 {
 }
 
 func (ra *Bitmap) And(bm *Bitmap) {
+	if bm == nil {
+		ra.Reset()
+		return
+	}
+
 	a, b := ra, bm
 	ai, an := 0, a.keys.numKeys()
 	bi, bn := 0, b.keys.numKeys()
@@ -787,6 +809,9 @@ func And(a, b *Bitmap) *Bitmap {
 }
 
 func (ra *Bitmap) AndNot(bm *Bitmap) {
+	if bm == nil {
+		return
+	}
 	a, b := ra, bm
 	var ai, bi int
 
@@ -839,6 +864,9 @@ func (ra *Bitmap) AndNot(bm *Bitmap) {
 
 // TODO: Check if we want to use lazyMode
 func (dst *Bitmap) Or(src *Bitmap) {
+	if src == nil {
+		return
+	}
 	dst.or(src, runInline)
 }
 
@@ -926,6 +954,39 @@ func Or(a, b *Bitmap) *Bitmap {
 	return res
 }
 
+func (ra *Bitmap) Rank(x uint64) int {
+	key := x & mask
+	offset, has := ra.keys.getValue(key)
+	if !has {
+		return -1
+	}
+	c := ra.getContainer(offset)
+	y := uint16(x)
+
+	// Find the rank within the container
+	var rank int
+	switch c[indexType] {
+	case typeArray:
+		rank = array(c).rank(y)
+	case typeBitmap:
+		rank = bitmap(c).rank(y)
+	}
+	if rank < 0 {
+		return -1
+	}
+
+	// Add up cardinalities of all the containers on the left of container containing x.
+	n := ra.keys.numKeys()
+	for i := 0; i < n; i++ {
+		if ra.keys.key(i) == key {
+			break
+		}
+		cont := ra.getContainer(ra.keys.val(i))
+		rank += getCardinality(cont)
+	}
+	return rank
+}
+
 func (ra *Bitmap) Cleanup() {
 	for idx := 1; idx < ra.keys.numKeys(); {
 		off := ra.keys.val(idx)

bitmap_test.go

@@ -654,5 +654,38 @@ func TestCleanup(t *testing.T) {
 	for i := 65536; i < n; i++ {
 		require.Falsef(t, a.Contains(uint64(i)), "idx: %d", i)
 	}
+}
+
+func TestRank(t *testing.T) {
+	a := NewBitmap()
+	n := int(1e6)
+	for i := uint64(0); i < uint64(n); i++ {
+		a.Set(i)
+	}
+	for i := 0; i < n; i++ {
+		require.Equal(t, i, a.Rank(uint64(i)))
+	}
+	require.Equal(t, -1, a.Rank(uint64(n)))
+
+	// Check ranks after removing an element.
+	a.Remove(100)
+	for i := 0; i < n; i++ {
+		if i < 100 {
+			require.Equal(t, i, a.Rank(uint64(i)))
+		} else if i == 100 {
+			require.Equal(t, -1, a.Rank(uint64(i)))
+		} else {
+			require.Equal(t, i-1, a.Rank(uint64(i)))
+		}
+	}
 
+	// Check ranks after removing a range of elements.
+	a.RemoveRange(0, uint64(1e4))
+	for i := 0; i < n; i++ {
+		if i < 1e4 {
+			require.Equal(t, -1, a.Rank(uint64(n)))
+		} else {
+			require.Equal(t, i-1e4, a.Rank(uint64(i)))
+		}
+	}
 }

container.go

@@ -123,6 +123,16 @@ func (c array) find(x uint16) int {
 	}
 	return N
 }
+
+func (c array) rank(x uint16) int {
+	N := getCardinality(c)
+	idx := c.find(x)
+	if idx == N {
+		return -1
+	}
+	return idx
+}
+
 func (c array) has(x uint16) bool {
 	N := getCardinality(c)
 	idx := c.find(x)
@@ -441,13 +451,32 @@ func (b bitmap) has(x uint16) bool {
 	return has > 0
 }
 
+func (b bitmap) rank(x uint16) int {
+	idx := x >> 4
+	pos := x & 0xF
+	if b[startIdx+idx]&bitmapMask[pos] == 0 {
+		return -1
+	}
+
+	var rank int
+	for i := 0; i < int(idx); i++ {
+		rank += bits.OnesCount16(b[int(startIdx)+i])
+	}
+	for p := uint16(0); p <= pos; p++ {
+		if b[startIdx+idx]&bitmapMask[p] > 0 {
+			rank++
+		}
+	}
+	return rank - 1
+}
+
 // TODO: This can perhaps be using SIMD instructions.
 func (b bitmap) andBitmap(other bitmap) []uint16 {
 	out := make([]uint16, maxContainerSize)
 	out[indexSize] = maxContainerSize
 	out[indexType] = typeBitmap
 	var num int
-	for i := 4; i < len(b); i++ {
+	for i := int(startIdx); i < len(b); i++ {
 		out[i] = b[i] & other[i]
 		num += bits.OnesCount16(out[i])
 	}