update dtrie to use new bitmap

Author: Theodus

trie/dtrie/dtrie.go

@@ -32,13 +32,37 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 // Jurgen J. Vinju
 package dtrie
 
+import "fmt"
+
 // Dtrie is a persistent hash trie that dynamically expands or shrinks
 // to provide efficient memory allocation.
 type Dtrie struct {
 	root   *node
 	hasher func(v interface{}) uint32
 }
 
+type entry struct {
+	hash  uint32
+	key   interface{}
+	value interface{}
+}
+
+func (e *entry) KeyHash() uint32 {
+	return e.hash
+}
+
+func (e *entry) Key() interface{} {
+	return e.key
+}
+
+func (e *entry) Value() interface{} {
+	return e.value
+}
+
+func (e *entry) String() string {
+	return fmt.Sprintf("%v:%v", e.key, e.value)
+}
+
 // New creates an empty DTrie with the given hashing function.
 // If nil is passed in, the default hashing function will be used.
 func New(hasher func(v interface{}) uint32) *Dtrie {
@@ -67,8 +91,8 @@ func (d *Dtrie) Get(key interface{}) Entry {
 
 // Insert adds an entry to the Dtrie, replacing the existing value if
 // the key already exists and returns the resulting Dtrie.
-func (d *Dtrie) Insert(entry Entry) *Dtrie {
-	root := insert(d.root, entry)
+func (d *Dtrie) Insert(key, value interface{}) *Dtrie {
+	root := insert(d.root, &entry{d.hasher(key), key, value})
 	return &Dtrie{root, d.hasher}
 }
 

trie/dtrie/dtrie_test.go

@@ -27,7 +27,6 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 package dtrie
 
 import (
-	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -40,28 +39,6 @@ func TestDefaultHasher(t *testing.T) {
 	assert.NotEqual(t, defaultHasher("foo"), defaultHasher("bar"))
 }
 
-type testEntry struct {
-	hash  uint32
-	key   int
-	value int
-}
-
-func (e *testEntry) KeyHash() uint32 {
-	return e.hash
-}
-
-func (e *testEntry) Key() interface{} {
-	return e.key
-}
-
-func (e *testEntry) Value() interface{} {
-	return e.value
-}
-
-func (e *testEntry) String() string {
-	return fmt.Sprint(e.value)
-}
-
 func collisionHash(key interface{}) uint32 {
 	return uint32(0xffffffff) // for testing collisions
 }
@@ -74,7 +51,7 @@ func TestInsert(t *testing.T) {
 func insertTest(t *testing.T, hashfunc func(interface{}) uint32, count int) *node {
 	n := emptyNode(0, 32)
 	for i := 0; i < count; i++ {
-		n = insert(n, &testEntry{hashfunc(i), i, i})
+		n = insert(n, &entry{hashfunc(i), i, i})
 	}
 	return n
 }
@@ -117,7 +94,7 @@ func TestUpdate(t *testing.T) {
 func updateTest(t *testing.T, hashfunc func(interface{}) uint32, count int) {
 	n := insertTest(t, hashfunc, count)
 	for i := 0; i < count; i++ {
-		n = insert(n, &testEntry{hashfunc(i), i, -i})
+		n = insert(n, &entry{hashfunc(i), i, -i})
 	}
 }
 
@@ -161,7 +138,7 @@ func BenchmarkInsert(b *testing.B) {
 	n := emptyNode(0, 32)
 	b.ResetTimer()
 	for i := b.N; i > 0; i-- {
-		n = insert(n, &testEntry{defaultHasher(i), i, i})
+		n = insert(n, &entry{defaultHasher(i), i, i})
 	}
 }
 
@@ -188,6 +165,6 @@ func BenchmarkUpdate(b *testing.B) {
 	n := insertTest(nil, defaultHasher, b.N)
 	b.ResetTimer()
 	for i := b.N; i > 0; i-- {
-		n = insert(n, &testEntry{defaultHasher(i), i, -i})
+		n = insert(n, &entry{defaultHasher(i), i, -i})
 	}
 }

trie/dtrie/node.go

@@ -30,13 +30,14 @@ import (
 	"fmt"
 	"sync"
 
-	"github.com/Workiva/go-datastructures/bitmap"
+	//TODO switch back to "github.com/Workiva/go-datastructures/bitarray"
+	"github.com/theodus/go-datastructures/bitarray"
 )
 
 type node struct {
 	entries []Entry
-	nodeMap bitmap.Bitmap32
-	dataMap bitmap.Bitmap32
+	nodeMap bitarray.Bitmap32
+	dataMap bitarray.Bitmap32
 	level   uint32 // level starts at 0
 }
 
@@ -74,12 +75,14 @@ func emptyNode(level uint32, capacity int) *node {
 func insert(n *node, entry Entry) *node {
 	index := mask(entry.KeyHash(), n.level)
 	newNode := n
-	if newNode.level == 6 { // handle hash collisions on 6th level
-		if newNode.entries[index] == nil {
-			newNode.entries[index] = entry
-			newNode.dataMap = newNode.dataMap.SetBit(uint(index))
-			return newNode
-		}
+	// insert directly
+	if newNode.entries[index] == nil {
+		newNode.entries[index] = entry
+		newNode.dataMap = newNode.dataMap.SetBit(uint(index))
+		return newNode
+	}
+	// handle hash collisions on 6th level
+	if newNode.level == 6 {
 		if newNode.dataMap.HasBit(uint(index)) {
 			if newNode.entries[index].Key() == entry.Key() {
 				newNode.entries[index] = entry
@@ -96,16 +99,12 @@ func insert(n *node, entry Entry) *node {
 		cNode.entries = append(cNode.entries, entry)
 		return newNode
 	}
-	// insert directly
-	if !newNode.dataMap.HasBit(uint(index)) && !newNode.nodeMap.HasBit(uint(index)) {
-		newNode.entries[index] = entry
-		newNode.dataMap = newNode.dataMap.SetBit(uint(index))
-		return newNode
-	}
-	if newNode.nodeMap.HasBit(uint(index)) { // insert into sub-node
+	// insert into sub-node
+	if newNode.nodeMap.HasBit(uint(index)) {
 		newNode.entries[index] = insert(newNode.entries[index].(*node), entry)
 		return newNode
 	}
+	// replace existing entry
 	if newNode.entries[index].Key() == entry.Key() {
 		newNode.entries[index] = entry
 		return newNode