change bitmap to go-datastructures version

Author: Theodus

trie/dtrie/dtrie_test.go

@@ -33,19 +33,6 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestPopCount(t *testing.T) {
-	b := []uint32{
-		uint32(0x55555555), // 0x55555555 = 01010101 01010101 01010101 01010101
-		uint32(0x33333333), // 0x33333333 = 00110011 00110011 00110011 00110011
-		uint32(0x0F0F0F0F), // 0x0F0F0F0F = 00001111 00001111 00001111 00001111
-		uint32(0x00FF00FF), // 0x00FF00FF = 00000000 11111111 00000000 11111111
-		uint32(0x0000FFFF), // 0x0000FFFF = 00000000 00000000 11111111 11111111
-	}
-	for _, x := range b {
-		assert.Equal(t, 16, popCount(x))
-	}
-}
-
 func TestDefaultHasher(t *testing.T) {
 	assert.Equal(t,
 		defaultHasher(map[int]string{11234: "foo"}),

trie/dtrie/node.go

@@ -29,12 +29,14 @@ package dtrie
 import (
 	"fmt"
 	"sync"
+
+	"github.com/Workiva/go-datastructures/bitmap"
 )
 
 type node struct {
 	entries []Entry
-	nodeMap uint32
-	dataMap uint32
+	nodeMap bitmap.Bitmap32
+	dataMap bitmap.Bitmap32
 	level   uint32 // level starts at 0
 }
 
@@ -75,10 +77,10 @@ func insert(n *node, entry Entry) *node {
 	if newNode.level == 6 { // handle hash collisions on 6th level
 		if newNode.entries[index] == nil {
 			newNode.entries[index] = entry
-			newNode.dataMap = setBit(newNode.dataMap, index)
+			newNode.dataMap = newNode.dataMap.SetBit(uint(index))
 			return newNode
 		}
-		if hasBit(newNode.dataMap, index) {
+		if newNode.dataMap.HasBit(uint(index)) {
 			if newNode.entries[index].Key() == entry.Key() {
 				newNode.entries[index] = entry
 				return newNode
@@ -87,19 +89,20 @@ func insert(n *node, entry Entry) *node {
 			cNode.entries[0] = newNode.entries[index]
 			cNode.entries[1] = entry
 			newNode.entries[index] = cNode
-			newNode.dataMap = clearBit(newNode.dataMap, index)
+			newNode.dataMap = newNode.dataMap.ClearBit(uint(index))
 			return newNode
 		}
 		cNode := newNode.entries[index].(*collisionNode)
 		cNode.entries = append(cNode.entries, entry)
 		return newNode
 	}
-	if !hasBit(newNode.dataMap, index) && !hasBit(newNode.nodeMap, index) { // insert directly
+	// insert directly
+	if !newNode.dataMap.HasBit(uint(index)) && !newNode.nodeMap.HasBit(uint(index)) {
 		newNode.entries[index] = entry
-		newNode.dataMap = setBit(newNode.dataMap, index)
+		newNode.dataMap = newNode.dataMap.SetBit(uint(index))
 		return newNode
 	}
-	if hasBit(newNode.nodeMap, index) { // insert into sub-node
+	if newNode.nodeMap.HasBit(uint(index)) { // insert into sub-node
 		newNode.entries[index] = insert(newNode.entries[index].(*node), entry)
 		return newNode
 	}
@@ -116,19 +119,19 @@ func insert(n *node, entry Entry) *node {
 	}
 	subNode = insert(subNode, newNode.entries[index])
 	subNode = insert(subNode, entry)
-	newNode.dataMap = clearBit(newNode.dataMap, index)
-	newNode.nodeMap = setBit(newNode.nodeMap, index)
+	newNode.dataMap = newNode.dataMap.ClearBit(uint(index))
+	newNode.nodeMap = newNode.nodeMap.SetBit(uint(index))
 	newNode.entries[index] = subNode
 	return newNode
 }
 
 // returns nil if not found
 func get(n *node, keyHash uint32, key interface{}) Entry {
 	index := mask(keyHash, n.level)
-	if hasBit(n.dataMap, index) {
+	if n.dataMap.HasBit(uint(index)) {
 		return n.entries[index]
 	}
-	if hasBit(n.nodeMap, index) {
+	if n.nodeMap.HasBit(uint(index)) {
 		return get(n.entries[index].(*node), keyHash, key)
 	}
 	if n.level == 6 { // get from collisionNode
@@ -148,26 +151,26 @@ func get(n *node, keyHash uint32, key interface{}) Entry {
 func remove(n *node, keyHash uint32, key interface{}) *node {
 	index := mask(keyHash, n.level)
 	newNode := n
-	if hasBit(n.dataMap, index) {
+	if n.dataMap.HasBit(uint(index)) {
 		newNode.entries[index] = nil
-		newNode.dataMap = clearBit(newNode.dataMap, index)
+		newNode.dataMap = newNode.dataMap.ClearBit(uint(index))
 		return newNode
 	}
-	if hasBit(n.nodeMap, index) {
+	if n.nodeMap.HasBit(uint(index)) {
 		subNode := newNode.entries[index].(*node)
 		subNode = remove(subNode, keyHash, key)
 		// compress if only 1 entry exists in sub-node
-		if popCount(subNode.nodeMap) == 0 && popCount(subNode.dataMap) == 1 {
+		if subNode.nodeMap.PopCount() == 0 && subNode.dataMap.PopCount() == 1 {
 			var e Entry
 			for i := uint32(0); i < 32; i++ {
-				if hasBit(subNode.dataMap, i) {
+				if subNode.dataMap.HasBit(uint(i)) {
 					e = subNode.entries[i]
 					break
 				}
 			}
 			newNode.entries[index] = e
-			newNode.nodeMap = clearBit(newNode.nodeMap, index)
-			newNode.dataMap = setBit(newNode.dataMap, index)
+			newNode.nodeMap = newNode.nodeMap.ClearBit(uint(index))
+			newNode.dataMap = newNode.dataMap.SetBit(uint(index))
 		}
 		newNode.entries[index] = subNode
 		return newNode
@@ -183,7 +186,7 @@ func remove(n *node, keyHash uint32, key interface{}) *node {
 		// compress if only 1 entry exists in collisionNode
 		if len(cNode.entries) == 1 {
 			newNode.entries[index] = cNode.entries[0]
-			newNode.dataMap = setBit(newNode.dataMap, index)
+			newNode.dataMap = newNode.dataMap.SetBit(uint(index))
 		}
 		return newNode
 	}
@@ -208,9 +211,9 @@ func pushEntries(n *node, stop <-chan struct{}, out chan Entry) {
 		default:
 			index := uint32(i)
 			switch {
-			case hasBit(n.dataMap, index):
+			case n.dataMap.HasBit(uint(index)):
 				out <- e
-			case hasBit(n.nodeMap, index):
+			case n.nodeMap.HasBit(uint(index)):
 				wg.Add(1)
 				go func() {
 					defer wg.Done()

trie/dtrie/util.go

@@ -35,29 +35,6 @@ func mask(hash, level uint32) uint32 {
 	return (hash >> (5 * level)) & 0x01f
 }
 
-func setBit(bitmap uint32, pos uint32) uint32 {
-	return bitmap | (1 << pos)
-}
-
-func clearBit(bitmap uint32, pos uint32) uint32 {
-	return bitmap & ^(1 << pos)
-}
-
-func hasBit(bitmap uint32, pos uint32) bool {
-	return (bitmap & (1 << pos)) != 0
-}
-
-func popCount(bitmap uint32) int {
-	// bit population count, see
-	// http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel
-	bitmap -= (bitmap >> 1) & 0x55555555
-	bitmap = (bitmap>>2)&0x33333333 + bitmap&0x33333333
-	bitmap += bitmap >> 4
-	bitmap &= 0x0f0f0f0f
-	bitmap *= 0x01010101
-	return int(byte(bitmap >> 24))
-}
-
 func defaultHasher(value interface{}) uint32 {
 	switch v := value.(type) {
 	case uint8: