From ce052c133c1c39e76a53b15a24da5c6ebc8daf7d Mon Sep 17 00:00:00 2001
From: Travis Turner <travis@pilosa.com>
Date: Fri, 12 Jan 2018 13:27:16 -0600
Subject: [PATCH] b+tree for Containers interface

---
 roaring/containers_btree.go                   | 125 ++++++++++++++++++
 .../{containers.go => containers_skiplist.go} |   0
 roaring/containers_test.go                    |  24 ++--
 roaring/roaring.go                            |   2 +-
 4 files changed, 138 insertions(+), 13 deletions(-)
 create mode 100644 roaring/containers_btree.go
 rename roaring/{containers.go => containers_skiplist.go} (100%)

diff --git a/roaring/containers_btree.go b/roaring/containers_btree.go
new file mode 100644
index 000000000..ce78c278a
--- /dev/null
+++ b/roaring/containers_btree.go
@@ -0,0 +1,125 @@
+package roaring
+
+import (
+	"io"
+
+	btree "github.com/pilosa/b"
+)
+
+func cmp(a, b uint64) int {
+	return int(a - b)
+}
+
+func NewBTreeContainers() *BTreeContainers {
+	return &BTreeContainers{
+		tree: btree.TreeNew(cmp),
+	}
+}
+
+type BTreeContainers struct {
+	tree *btree.Tree
+
+	lastKey       uint64
+	lastContainer *container
+}
+
+func (btc *BTreeContainers) Get(key uint64) *container {
+	var c *container
+	el, ok := btc.tree.Get(key)
+	if ok {
+		c = el.(*container)
+	}
+	return c
+}
+
+func (btc *BTreeContainers) Put(key uint64, c *container) {
+	btc.tree.Set(key, c)
+}
+
+func (btc *BTreeContainers) Remove(key uint64) {
+	btc.tree.Delete(key)
+}
+
+func (btc *BTreeContainers) GetOrCreate(key uint64) *container {
+	// Check the last* cache for same container.
+	if key == btc.lastKey && btc.lastContainer != nil {
+		return btc.lastContainer
+	}
+
+	btc.lastKey = key
+	v, ok := btc.tree.Get(key)
+	if !ok {
+		cont := newContainer()
+		btc.tree.Set(key, cont)
+		btc.lastContainer = cont
+		return cont
+	}
+
+	btc.lastContainer = v.(*container)
+	return btc.lastContainer
+}
+
+func (btc *BTreeContainers) Clone() Containers {
+	nbtc := NewBTreeContainers()
+
+	itr, err := btc.tree.SeekFirst()
+	if err == io.EOF {
+		return nbtc
+	}
+	for {
+		k, v, err := itr.Next()
+		if err == io.EOF {
+			break
+		}
+		nbtc.tree.Set(k, v.(*container).clone())
+	}
+
+	return nbtc
+}
+
+func (btc *BTreeContainers) Last() (key uint64, c *container) {
+	if btc.tree.Len() == 0 {
+		return 0, nil
+	}
+	k, v := btc.tree.Last()
+	return k, v.(*container)
+}
+
+func (btc *BTreeContainers) Size() int {
+	return btc.tree.Len()
+}
+
+func (btc *BTreeContainers) Iterator(key uint64) (citer Contiterator, found bool) {
+	e, ok := btc.tree.Seek(key)
+	if ok {
+		found = true
+	}
+
+	return &BTCIterator{
+		e: e,
+	}, found
+}
+
+type BTCIterator struct {
+	e   *btree.Enumerator
+	key interface{}
+	val interface{}
+}
+
+func (i *BTCIterator) Next() bool {
+
+	k, v, err := i.e.Next()
+	if err == io.EOF {
+		return false
+	}
+	i.key = k
+	i.val = v
+	return true
+}
+
+func (i *BTCIterator) Value() (uint64, *container) {
+	if i.val == nil {
+		return 0, nil
+	}
+	return i.key.(uint64), i.val.(*container)
+}
diff --git a/roaring/containers.go b/roaring/containers_skiplist.go
similarity index 100%
rename from roaring/containers.go
rename to roaring/containers_skiplist.go
diff --git a/roaring/containers_test.go b/roaring/containers_test.go
index 464a92fd6..6d155f397 100644
--- a/roaring/containers_test.go
+++ b/roaring/containers_test.go
@@ -5,19 +5,19 @@ import (
 )
 
 func TestContainersIterator(t *testing.T) {
-	slc := NewSkipListContainers()
-	itr, found := slc.Iterator(0)
+	btc := NewBTreeContainers()
+	itr, found := btc.Iterator(0)
 	if found {
-		t.Fatalf("shouldn't have found 0 in empty slc")
+		t.Fatalf("shouldn't have found 0 in empty btc")
 	}
 	if itr.Next() {
-		t.Fatal("Next() should be false for empty slc")
+		t.Fatal("Next() should be false for empty btc")
 	}
 
-	slc.Put(1, &container{n: 1})
-	slc.Put(2, &container{n: 2})
+	btc.Put(1, &container{n: 1})
+	btc.Put(2, &container{n: 2})
 
-	itr, found = slc.Iterator(0)
+	itr, found = btc.Iterator(0)
 	if found {
 		t.Fatalf("shouldn't have found 0")
 	}
@@ -39,11 +39,11 @@ func TestContainersIterator(t *testing.T) {
 		t.Fatalf("itr should be done, but got true")
 	}
 
-	slc.Put(3, &container{n: 3})
-	slc.Put(5, &container{n: 5})
-	slc.Put(6, &container{n: 6})
+	btc.Put(3, &container{n: 3})
+	btc.Put(5, &container{n: 5})
+	btc.Put(6, &container{n: 6})
 
-	itr, found = slc.Iterator(3)
+	itr, found = btc.Iterator(3)
 	if !itr.Next() {
 		t.Fatalf("3 should be next, but got false")
 	}
@@ -60,7 +60,7 @@ func TestContainersIterator(t *testing.T) {
 		t.Fatalf("Wrong k/v, exp: 5,5 got: %v,%v", key, val.n)
 	}
 
-	itr, found = slc.Iterator(4)
+	itr, found = btc.Iterator(4)
 	if found {
 		t.Fatalf("shouldn't have found 4")
 	}
diff --git a/roaring/roaring.go b/roaring/roaring.go
index 367983322..fb1994bfd 100644
--- a/roaring/roaring.go
+++ b/roaring/roaring.go
@@ -110,7 +110,7 @@ type Bitmap struct {
 // NewBitmap returns a Bitmap with an initial set of values.
 func NewBitmap(a ...uint64) *Bitmap {
 	b := &Bitmap{
-		conts: NewSkipListContainers(),
+		conts: NewBTreeContainers(),
 	}
 	b.Add(a...)
 	return b