From 3bba5f79b666b1874b50ff64a0bff168bbe354b1 Mon Sep 17 00:00:00 2001
From: SimFG <bang.fu@zilliz.com>
Date: Thu, 21 Jul 2022 20:54:32 +0800
Subject: [PATCH] Extract `wordsIndex` to calculate the index of words in a
 `uint64`, and use it to replace `i % wordSize`. In addition, change the
 `InsertAt` function to reduce one bit operation.

---
 bitset.go | 48 ++++++++++++++++++++++++++----------------------
 1 file changed, 26 insertions(+), 22 deletions(-)

diff --git a/bitset.go b/bitset.go
index b423743..49fe4bc 100644
--- a/bitset.go
+++ b/bitset.go
@@ -116,6 +116,11 @@ func wordsNeeded(i uint) int {
 	return int((i + (wordSize - 1)) >> log2WordSize)
 }
 
+// wordsIndex calculates the index of words in a `uint64`
+func wordsIndex(i uint) uint {
+	return i & (wordSize - 1)
+}
+
 // New creates a new BitSet with a hint that length bits will be required
 func New(length uint) (bset *BitSet) {
 	defer func() {
@@ -171,7 +176,7 @@ func (b *BitSet) Test(i uint) bool {
 	if i >= b.length {
 		return false
 	}
-	return b.set[i>>log2WordSize]&(1<<(i&(wordSize-1))) != 0
+	return b.set[i>>log2WordSize]&(1<<wordsIndex(i)) != 0
 }
 
 // Set bit i to 1, the capacity of the bitset is automatically
@@ -182,7 +187,7 @@ func (b *BitSet) Test(i uint) bool {
 // for providing sensible parameters in line with their memory capacity.
 func (b *BitSet) Set(i uint) *BitSet {
 	b.extendSetMaybe(i)
-	b.set[i>>log2WordSize] |= 1 << (i & (wordSize - 1))
+	b.set[i>>log2WordSize] |= 1 << wordsIndex(i)
 	return b
 }
 
@@ -191,7 +196,7 @@ func (b *BitSet) Clear(i uint) *BitSet {
 	if i >= b.length {
 		return b
 	}
-	b.set[i>>log2WordSize] &^= 1 << (i & (wordSize - 1))
+	b.set[i>>log2WordSize] &^= 1 << wordsIndex(i)
 	return b
 }
 
@@ -216,7 +221,7 @@ func (b *BitSet) Flip(i uint) *BitSet {
 	if i >= b.length {
 		return b.Set(i)
 	}
-	b.set[i>>log2WordSize] ^= 1 << (i & (wordSize - 1))
+	b.set[i>>log2WordSize] ^= 1 << wordsIndex(i)
 	return b
 }
 
@@ -233,13 +238,13 @@ func (b *BitSet) FlipRange(start, end uint) *BitSet {
 	b.extendSetMaybe(end - 1)
 	var startWord uint = start >> log2WordSize
 	var endWord uint = end >> log2WordSize
-	b.set[startWord] ^= ^(^uint64(0) << (start & (wordSize - 1)))
+	b.set[startWord] ^= ^(^uint64(0) << wordsIndex(start))
 	for i := startWord; i < endWord; i++ {
 		b.set[i] = ^b.set[i]
 	}
-	if end & (wordSize - 1) != 0 {
-		b.set[endWord] ^= ^uint64(0) >> (-end & (wordSize - 1))
-    }
+	if end&(wordSize-1) != 0 {
+		b.set[endWord] ^= ^uint64(0) >> wordsIndex(-end)
+	}
 	return b
 }
 
@@ -268,7 +273,7 @@ func (b *BitSet) Shrink(lastbitindex uint) *BitSet {
 	b.set = shrunk
 	b.length = length
 	if length < 64 {
-		b.set[idx-1] &= (allBits >> (uint64(64) - uint64(length&(wordSize-1))))
+		b.set[idx-1] &= allBits >> uint64(64-wordsIndex(length))
 	}
 	return b
 }
@@ -298,7 +303,7 @@ func (b *BitSet) Compact() *BitSet {
 // this method could be extremely slow and in some cases might cause the entire BitSet
 // to be recopied.
 func (b *BitSet) InsertAt(idx uint) *BitSet {
-	insertAtElement := (idx >> log2WordSize)
+	insertAtElement := idx >> log2WordSize
 
 	// if length of set is a multiple of wordSize we need to allocate more space first
 	if b.isLenExactMultiple() {
@@ -317,13 +322,13 @@ func (b *BitSet) InsertAt(idx uint) *BitSet {
 
 	// generate a mask to extract the data that we need to shift left
 	// within the element where we insert a bit
-	dataMask := ^(uint64(1)<<uint64(idx&(wordSize-1)) - 1)
+	dataMask := uint64(1)<<uint64(wordsIndex(idx)) - 1
 
 	// extract that data that we'll shift
-	data := b.set[i] & dataMask
+	data := b.set[i] & (^dataMask)
 
 	// set the positions of the data mask to 0 in the element where we insert
-	b.set[i] &= ^dataMask
+	b.set[i] &= dataMask
 
 	// shift data mask to the left and insert its data to the slice element
 	b.set[i] |= data << 1
@@ -371,7 +376,7 @@ func (b *BitSet) DeleteAt(i uint) *BitSet {
 
 	// generate a mask for the data that needs to be shifted right
 	// within that slice element that gets modified
-	dataMask := ^((uint64(1) << (i & (wordSize - 1))) - 1)
+	dataMask := ^((uint64(1) << wordsIndex(i)) - 1)
 
 	// extract the data that we'll shift right from the slice element
 	data := b.set[deleteAtElement] & dataMask
@@ -409,7 +414,7 @@ func (b *BitSet) NextSet(i uint) (uint, bool) {
 		return 0, false
 	}
 	w := b.set[x]
-	w = w >> (i & (wordSize - 1))
+	w = w >> wordsIndex(i)
 	if w != 0 {
 		return i + trailingZeroes64(w), true
 	}
@@ -453,7 +458,7 @@ func (b *BitSet) NextSetMany(i uint, buffer []uint) (uint, []uint) {
 	if x >= len(b.set) || capacity == 0 {
 		return 0, myanswer[:0]
 	}
-	skip := i & (wordSize - 1)
+	skip := wordsIndex(i)
 	word := b.set[x] >> skip
 	myanswer = myanswer[:capacity]
 	size := int(0)
@@ -496,8 +501,8 @@ func (b *BitSet) NextClear(i uint) (uint, bool) {
 		return 0, false
 	}
 	w := b.set[x]
-	w = w >> (i & (wordSize - 1))
-	wA := allBits >> (i & (wordSize - 1))
+	w = w >> wordsIndex(i)
+	wA := allBits >> wordsIndex(i)
 	index := i + trailingZeroes64(^w)
 	if w != wA && index < b.length {
 		return index, true
@@ -810,17 +815,17 @@ func (b *BitSet) InPlaceSymmetricDifference(compare *BitSet) {
 
 // Is the length an exact multiple of word sizes?
 func (b *BitSet) isLenExactMultiple() bool {
-	return b.length%wordSize == 0
+	return wordsIndex(b.length) == 0
 }
 
 // Clean last word by setting unused bits to 0
 func (b *BitSet) cleanLastWord() {
 	if !b.isLenExactMultiple() {
-		b.set[len(b.set)-1] &= allBits >> (wordSize - b.length%wordSize)
+		b.set[len(b.set)-1] &= allBits >> (wordSize - wordsIndex(b.length))
 	}
 }
 
-// Complement computes the (local) complement of a biset (up to length bits)
+// Complement computes the (local) complement of a bitset (up to length bits)
 func (b *BitSet) Complement() (result *BitSet) {
 	panicIfNull(b)
 	result = New(b.length)
@@ -848,7 +853,6 @@ func (b *BitSet) None() bool {
 				return false
 			}
 		}
-		return true
 	}
 	return true
 }