From d60153209945d969171d86afec22b3ab9f7bad58 Mon Sep 17 00:00:00 2001
From: Daniel Lemire <daniel@lemire.me>
Date: Mon, 25 Dec 2023 10:43:29 -0500
Subject: [PATCH] fix

---
 bitset.go      | 40 +++++++++++++++++++++++++++++++++++++---
 bitset_test.go | 37 +++++++++++++++++++++++++++++++++++++
 select.go      | 45 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 119 insertions(+), 3 deletions(-)
 create mode 100644 select.go

diff --git a/bitset.go b/bitset.go
index 8fb9e9f..9aec0f3 100644
--- a/bitset.go
+++ b/bitset.go
@@ -94,17 +94,17 @@ func (b *BitSet) SetBitsetFrom(buf []uint64) {
 	b.set = buf
 }
 
-// From is a constructor used to create a BitSet from an array of integers
+// From is a constructor used to create a BitSet from an array of words
 func From(buf []uint64) *BitSet {
 	return FromWithLength(uint(len(buf))*64, buf)
 }
 
-// FromWithLength constructs from an array of integers and length.
+// FromWithLength constructs from an array of words and length.
 func FromWithLength(len uint, set []uint64) *BitSet {
 	return &BitSet{len, set}
 }
 
-// Bytes returns the bitset as array of integers
+// Bytes returns the bitset as array of words
 func (b *BitSet) Bytes() []uint64 {
 	return b.set
 }
@@ -1135,3 +1135,37 @@ func (b *BitSet) UnmarshalJSON(data []byte) error {
 	_, err = b.ReadFrom(bytes.NewReader(buf))
 	return err
 }
+
+// Rank returns the nunber of set bits up to and including the index
+// that are set in the bitset.
+// See https://en.wikipedia.org/wiki/Ranking#Ranking_in_statistics
+func (b *BitSet) Rank(index uint) uint {
+	if index >= b.length {
+		return b.Count()
+	}
+	leftover := (index + 1) & 63
+	answer := uint(popcntSlice(b.set[:(index+1)>>6]))
+	if leftover != 0 {
+		answer += uint(popcount(b.set[(index+1)>>6] << (64 - leftover)))
+	}
+	return answer
+}
+
+// Select returns the index of the jth set bit, where j is the argument.
+// The caller is responsible to ensure that 0 <= j < Count(): when j is
+// out of range, the function returns the length of the bitset (b.length).
+//
+// Note that this function differs in convention from the Rank function which
+// returns 1 when ranking the smallest value. We follow the conventional
+// textbook definition of Select and Rank.
+func (b *BitSet) Select(index uint) uint {
+	leftover := index
+	for idx, word := range b.set {
+		w := uint(popcount(word))
+		if w > leftover {
+			return uint(idx)*64 + select64(word, leftover)
+		}
+		leftover -= w
+	}
+	return b.length
+}
diff --git a/bitset_test.go b/bitset_test.go
index 3dae32e..8f50df5 100644
--- a/bitset_test.go
+++ b/bitset_test.go
@@ -1497,6 +1497,43 @@ func TestClearAll(t *testing.T) {
 	}
 }
 
+func TestRankSelect(t *testing.T) {
+	u := []uint{2, 3, 5, 7, 11, 700, 1500}
+	b := BitSet{}
+	for _, v := range u {
+		b.Set(v)
+	}
+
+	if b.Rank(5) != 3 {
+		t.Error("Unexpected rank")
+		return
+	}
+	if b.Rank(6) != 3 {
+		t.Error("Unexpected rank")
+		return
+	}
+	if b.Rank(1500) != 7 {
+		t.Error("Unexpected rank")
+		return
+	}
+	if b.Select(0) != 2 {
+		t.Error("Unexpected select")
+		return
+	}
+	if b.Select(1) != 3 {
+		t.Error("Unexpected select")
+		return
+	}
+	if b.Select(2) != 5 {
+		t.Error("Unexpected select")
+		return
+	}
+
+	if b.Select(5) != 700 {
+		t.Error("Unexpected select")
+		return
+	}
+}
 func TestFlip(t *testing.T) {
 	b := new(BitSet)
 	c := b.Flip(11)
diff --git a/select.go b/select.go
new file mode 100644
index 0000000..f15e74a
--- /dev/null
+++ b/select.go
@@ -0,0 +1,45 @@
+package bitset
+
+func select64(w uint64, j uint) uint {
+	seen := 0
+	// Divide 64bit
+	part := w & 0xFFFFFFFF
+	n := uint(popcount(part))
+	if n <= j {
+		part = w >> 32
+		seen += 32
+		j -= n
+	}
+	ww := part
+
+	// Divide 32bit
+	part = ww & 0xFFFF
+
+	n = uint(popcount(part))
+	if n <= j {
+		part = ww >> 16
+		seen += 16
+		j -= n
+	}
+	ww = part
+
+	// Divide 16bit
+	part = ww & 0xFF
+	n = uint(popcount(part))
+	if n <= j {
+		part = ww >> 8
+		seen += 8
+		j -= n
+	}
+	ww = part
+
+	// Lookup in final byte
+	counter := 0
+	for ; counter < 8; counter++ {
+		j -= uint((ww >> counter) & 1)
+		if j+1 == 0 {
+			break
+		}
+	}
+	return uint(seen + counter)
+}