Merge pull request #1 from kortschak/b

Use cznic/b B+tree implementation in place of GoLLRB for memstore

graph/memstore/b/keys_test.go

@@ -0,0 +1,396 @@
+// Copyright 2014 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package b
+
+import (
+	"math"
+	"runtime/debug"
+	"testing"
+
+	"github.com/cznic/mathutil"
+)
+
+func rng() *mathutil.FC32 {
+	x, err := mathutil.NewFC32(math.MinInt32/4, math.MaxInt32/4, false)
+	if err != nil {
+		panic(err)
+	}
+
+	return x
+}
+
+func cmp(a, b int64) int {
+	return int(a - b)
+}
+
+func BenchmarkSetSeq1e3(b *testing.B) {
+	benchmarkSetSeq(b, 1e3)
+}
+
+func BenchmarkSetSeq1e4(b *testing.B) {
+	benchmarkSetSeq(b, 1e4)
+}
+
+func BenchmarkSetSeq1e5(b *testing.B) {
+	benchmarkSetSeq(b, 1e5)
+}
+
+func BenchmarkSetSeq1e6(b *testing.B) {
+	benchmarkSetSeq(b, 1e6)
+}
+
+func benchmarkSetSeq(b *testing.B, n int) {
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		r := TreeNew(cmp)
+		debug.FreeOSMemory()
+		b.StartTimer()
+		for j := int64(0); j < int64(n); j++ {
+			r.Set(j, struct{}{})
+		}
+		b.StopTimer()
+		r.Close()
+	}
+	b.StopTimer()
+}
+
+func BenchmarkGetSeq1e3(b *testing.B) {
+	benchmarkGetSeq(b, 1e3)
+}
+
+func BenchmarkGetSeq1e4(b *testing.B) {
+	benchmarkGetSeq(b, 1e4)
+}
+
+func BenchmarkGetSeq1e5(b *testing.B) {
+	benchmarkGetSeq(b, 1e5)
+}
+
+func BenchmarkGetSeq1e6(b *testing.B) {
+	benchmarkGetSeq(b, 1e6)
+}
+
+func benchmarkGetSeq(b *testing.B, n int) {
+	r := TreeNew(cmp)
+	for i := int64(0); i < int64(n); i++ {
+		r.Set(i, struct{}{})
+	}
+	debug.FreeOSMemory()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for j := int64(0); j < int64(n); j++ {
+			r.Get(j)
+		}
+	}
+	b.StopTimer()
+	r.Close()
+}
+
+func BenchmarkSetRnd1e3(b *testing.B) {
+	benchmarkSetRnd(b, 1e3)
+}
+
+func BenchmarkSetRnd1e4(b *testing.B) {
+	benchmarkSetRnd(b, 1e4)
+}
+
+func BenchmarkSetRnd1e5(b *testing.B) {
+	benchmarkSetRnd(b, 1e5)
+}
+
+func BenchmarkSetRnd1e6(b *testing.B) {
+	benchmarkSetRnd(b, 1e6)
+}
+
+func benchmarkSetRnd(b *testing.B, n int) {
+	rng := rng()
+	a := make([]int, n)
+	for i := range a {
+		a[i] = rng.Next()
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		r := TreeNew(cmp)
+		debug.FreeOSMemory()
+		b.StartTimer()
+		for _, v := range a {
+			r.Set(int64(v), struct{}{})
+		}
+		b.StopTimer()
+		r.Close()
+	}
+	b.StopTimer()
+}
+
+func BenchmarkGetRnd1e3(b *testing.B) {
+	benchmarkGetRnd(b, 1e3)
+}
+
+func BenchmarkGetRnd1e4(b *testing.B) {
+	benchmarkGetRnd(b, 1e4)
+}
+
+func BenchmarkGetRnd1e5(b *testing.B) {
+	benchmarkGetRnd(b, 1e5)
+}
+
+func BenchmarkGetRnd1e6(b *testing.B) {
+	benchmarkGetRnd(b, 1e6)
+}
+
+func benchmarkGetRnd(b *testing.B, n int) {
+	r := TreeNew(cmp)
+	rng := rng()
+	a := make([]int64, n)
+	for i := range a {
+		a[i] = int64(rng.Next())
+	}
+	for _, v := range a {
+		r.Set(v, struct{}{})
+	}
+	debug.FreeOSMemory()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for _, v := range a {
+			r.Get(v)
+		}
+	}
+	b.StopTimer()
+	r.Close()
+}
+
+func BenchmarkDelSeq1e3(b *testing.B) {
+	benchmarkDelSeq(b, 1e3)
+}
+
+func BenchmarkDelSeq1e4(b *testing.B) {
+	benchmarkDelSeq(b, 1e4)
+}
+
+func BenchmarkDelSeq1e5(b *testing.B) {
+	benchmarkDelSeq(b, 1e5)
+}
+
+func BenchmarkDelSeq1e6(b *testing.B) {
+	benchmarkDelSeq(b, 1e6)
+}
+
+func benchmarkDelSeq(b *testing.B, n int) {
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		r := TreeNew(cmp)
+		for j := int64(0); j < int64(n); j++ {
+			r.Set(j, struct{}{})
+		}
+		debug.FreeOSMemory()
+		b.StartTimer()
+		for j := int64(0); j < int64(n); j++ {
+			r.Delete(j)
+		}
+	}
+	b.StopTimer()
+}
+
+func BenchmarkDelRnd1e3(b *testing.B) {
+	benchmarkDelRnd(b, 1e3)
+}
+
+func BenchmarkDelRnd1e4(b *testing.B) {
+	benchmarkDelRnd(b, 1e4)
+}
+
+func BenchmarkDelRnd1e5(b *testing.B) {
+	benchmarkDelRnd(b, 1e5)
+}
+
+func BenchmarkDelRnd1e6(b *testing.B) {
+	benchmarkDelRnd(b, 1e6)
+}
+
+func benchmarkDelRnd(b *testing.B, n int) {
+	rng := rng()
+	a := make([]int64, n)
+	for i := range a {
+		a[i] = int64(rng.Next())
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		r := TreeNew(cmp)
+		for _, v := range a {
+			r.Set(v, struct{}{})
+		}
+		debug.FreeOSMemory()
+		b.StartTimer()
+		for _, v := range a {
+			r.Delete(v)
+		}
+		b.StopTimer()
+		r.Close()
+	}
+	b.StopTimer()
+}
+
+func BenchmarkSeekSeq1e3(b *testing.B) {
+	benchmarkSeekSeq(b, 1e3)
+}
+
+func BenchmarkSeekSeq1e4(b *testing.B) {
+	benchmarkSeekSeq(b, 1e4)
+}
+
+func BenchmarkSeekSeq1e5(b *testing.B) {
+	benchmarkSeekSeq(b, 1e5)
+}
+
+func BenchmarkSeekSeq1e6(b *testing.B) {
+	benchmarkSeekSeq(b, 1e6)
+}
+
+func benchmarkSeekSeq(b *testing.B, n int) {
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		t := TreeNew(cmp)
+		for j := int64(0); j < int64(n); j++ {
+			t.Set(j, struct{}{})
+		}
+		debug.FreeOSMemory()
+		b.StartTimer()
+		for j := int64(0); j < int64(n); j++ {
+			e, _ := t.Seek(j)
+			e.Close()
+		}
+		b.StopTimer()
+		t.Close()
+	}
+	b.StopTimer()
+}
+
+func BenchmarkSeekRnd1e3(b *testing.B) {
+	benchmarkSeekRnd(b, 1e3)
+}
+
+func BenchmarkSeekRnd1e4(b *testing.B) {
+	benchmarkSeekRnd(b, 1e4)
+}
+
+func BenchmarkSeekRnd1e5(b *testing.B) {
+	benchmarkSeekRnd(b, 1e5)
+}
+
+func BenchmarkSeekRnd1e6(b *testing.B) {
+	benchmarkSeekRnd(b, 1e6)
+}
+
+func benchmarkSeekRnd(b *testing.B, n int) {
+	r := TreeNew(cmp)
+	rng := rng()
+	a := make([]int64, n)
+	for i := range a {
+		a[i] = int64(rng.Next())
+	}
+	for _, v := range a {
+		r.Set(v, struct{}{})
+	}
+	debug.FreeOSMemory()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		for _, v := range a {
+			e, _ := r.Seek(v)
+			e.Close()
+		}
+	}
+	b.StopTimer()
+	r.Close()
+}
+
+func BenchmarkNext1e3(b *testing.B) {
+	benchmarkNext(b, 1e3)
+}
+
+func BenchmarkNext1e4(b *testing.B) {
+	benchmarkNext(b, 1e4)
+}
+
+func BenchmarkNext1e5(b *testing.B) {
+	benchmarkNext(b, 1e5)
+}
+
+func BenchmarkNext1e6(b *testing.B) {
+	benchmarkNext(b, 1e6)
+}
+
+func benchmarkNext(b *testing.B, n int) {
+	t := TreeNew(cmp)
+	for i := int64(0); i < int64(n); i++ {
+		t.Set(i, struct{}{})
+	}
+	debug.FreeOSMemory()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		en, err := t.SeekFirst()
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		m := 0
+		for {
+			if _, _, err = en.Next(); err != nil {
+				break
+			}
+			m++
+		}
+		if m != n {
+			b.Fatal(m)
+		}
+	}
+	b.StopTimer()
+	t.Close()
+}
+
+func BenchmarkPrev1e3(b *testing.B) {
+	benchmarkPrev(b, 1e3)
+}
+
+func BenchmarkPrev1e4(b *testing.B) {
+	benchmarkPrev(b, 1e4)
+}
+
+func BenchmarkPrev1e5(b *testing.B) {
+	benchmarkPrev(b, 1e5)
+}
+
+func BenchmarkPrev1e6(b *testing.B) {
+	benchmarkPrev(b, 1e6)
+}
+
+func benchmarkPrev(b *testing.B, n int) {
+	t := TreeNew(cmp)
+	for i := int64(0); i < int64(n); i++ {
+		t.Set(i, struct{}{})
+	}
+	debug.FreeOSMemory()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		en, err := t.SeekLast()
+		if err != nil {
+			b.Fatal(err)
+		}
+
+		m := 0
+		for {
+			if _, _, err = en.Prev(); err != nil {
+				break
+			}
+			m++
+		}
+		if m != n {
+			b.Fatal(m)
+		}
+	}
+}

graph/memstore/iterator.go

@@ -19,49 +19,36 @@ import (
 	"math"
 	"strings"
 
-	"github.com/petar/GoLLRB/llrb"
-
 	"github.com/google/cayley/graph"
 	"github.com/google/cayley/graph/iterator"
+	"github.com/google/cayley/graph/memstore/b"
 )
 
 type Iterator struct {
-	uid       uint64
-	ts        *TripleStore
-	tags      graph.Tagger
-	tree      *llrb.LLRB
-	data      string
-	isRunning bool
-	iterLast  Int64
-	result    graph.Value
+	uid    uint64
+	ts     *TripleStore
+	tags   graph.Tagger
+	tree   *b.Tree
+	iter   *b.Enumerator
+	data   string
+	result graph.Value
 }
 
-type Int64 int64
-
-func (i Int64) Less(than llrb.Item) bool {
-	return i < than.(Int64)
+func cmp(a, b int64) int {
+	return int(a - b)
 }
 
-func IterateOne(tree *llrb.LLRB, last Int64) Int64 {
-	var next Int64
-	tree.AscendGreaterOrEqual(last, func(i llrb.Item) bool {
-		if i.(Int64) == last {
-			return true
-		} else {
-			next = i.(Int64)
-			return false
-		}
-	})
-	return next
-}
-
-func NewLlrbIterator(tree *llrb.LLRB, data string, ts *TripleStore) *Iterator {
+func NewIterator(tree *b.Tree, data string, ts *TripleStore) *Iterator {
+	iter, err := tree.SeekFirst()
+	if err != nil {
+		iter = nil
+	}
 	return &Iterator{
-		uid:      iterator.NextUID(),
-		ts:       ts,
-		tree:     tree,
-		iterLast: Int64(-1),
-		data:     data,
+		uid:  iterator.NextUID(),
+		ts:   ts,
+		tree: tree,
+		iter: iter,
+		data: data,
 	}
 }
 
@@ -70,7 +57,11 @@ func (it *Iterator) UID() uint64 {
 }
 
 func (it *Iterator) Reset() {
-	it.iterLast = Int64(-1)
+	var err error
+	it.iter, err = it.tree.SeekFirst()
+	if err != nil {
+		it.iter = nil
+	}
 }
 
 func (it *Iterator) Tagger() *graph.Tagger {
@@ -88,8 +79,30 @@ func (it *Iterator) TagResults(dst map[string]graph.Value) {
 }
 
 func (it *Iterator) Clone() graph.Iterator {
-	m := NewLlrbIterator(it.tree, it.data, it.ts)
+	var iter *b.Enumerator
+	if it.result != nil {
+		var ok bool
+		iter, ok = it.tree.Seek(it.result.(int64))
+		if !ok {
+			panic("value unexpectedly missing")
+		}
+	} else {
+		var err error
+		iter, err = it.tree.SeekFirst()
+		if err != nil {
+			iter = nil
+		}
+	}
+
+	m := &Iterator{
+		uid:  iterator.NextUID(),
+		ts:   it.ts,
+		tree: it.tree,
+		iter: iter,
+		data: it.data,
+	}
 	m.tags.CopyFrom(it)
+
 	return m
 }
 
@@ -101,14 +114,18 @@ func (it *Iterator) checkValid(index int64) bool {
 
 func (it *Iterator) Next() bool {
 	graph.NextLogIn(it)
-	if it.tree.Max() == nil || it.iterLast == it.tree.Max().(Int64) {
+
+	if it.iter == nil {
 		return graph.NextLogOut(it, nil, false)
 	}
-	it.iterLast = IterateOne(it.tree, it.iterLast)
-	if !it.checkValid(int64(it.iterLast)) {
+	result, _, err := it.iter.Next()
+	if err != nil {
+		return graph.NextLogOut(it, nil, false)
+	}
+	if !it.checkValid(result) {
 		return it.Next()
 	}
-	it.result = int64(it.iterLast)
+	it.result = result
 	return graph.NextLogOut(it, it.result, true)
 }
 
@@ -135,7 +152,7 @@ func (it *Iterator) Size() (int64, bool) {
 
 func (it *Iterator) Contains(v graph.Value) bool {
 	graph.ContainsLogIn(it, v)
-	if it.tree.Has(Int64(v.(int64))) && it.checkValid(v.(int64)) {
+	if _, ok := it.tree.Get(v.(int64)); ok {
 		it.result = v
 		return graph.ContainsLogOut(it, v, true)
 	}
@@ -150,7 +167,7 @@ func (it *Iterator) DebugString(indent int) string {
 var memType graph.Type
 
 func init() {
-	memType = graph.RegisterIterator("llrb")
+	memType = graph.RegisterIterator("b+tree")
 }
 
 func Type() graph.Type { return memType }

graph/memstore/triplestore.go

@@ -18,11 +18,11 @@ import (
 	"fmt"
 
 	"github.com/barakmich/glog"
+
 	"github.com/google/cayley/graph"
 	"github.com/google/cayley/graph/iterator"
+	"github.com/google/cayley/graph/memstore/b"
 	"github.com/google/cayley/quad"
-
-	"github.com/petar/GoLLRB/llrb"
 )
 
 func init() {
@@ -32,47 +32,36 @@ func init() {
 }
 
 type QuadDirectionIndex struct {
-	subject   map[int64]*llrb.LLRB
-	predicate map[int64]*llrb.LLRB
-	object    map[int64]*llrb.LLRB
-	label     map[int64]*llrb.LLRB
+	index [4]map[int64]*b.Tree
 }
 
-func NewQuadDirectionIndex() *QuadDirectionIndex {
-	var qdi QuadDirectionIndex
-	qdi.subject = make(map[int64]*llrb.LLRB)
-	qdi.predicate = make(map[int64]*llrb.LLRB)
-	qdi.object = make(map[int64]*llrb.LLRB)
-	qdi.label = make(map[int64]*llrb.LLRB)
-	return &qdi
+func NewQuadDirectionIndex() QuadDirectionIndex {
+	return QuadDirectionIndex{[...]map[int64]*b.Tree{
+		quad.Subject - 1:   make(map[int64]*b.Tree),
+		quad.Predicate - 1: make(map[int64]*b.Tree),
+		quad.Object - 1:    make(map[int64]*b.Tree),
+		quad.Label - 1:     make(map[int64]*b.Tree),
+	}}
 }
 
-func (qdi *QuadDirectionIndex) GetForDir(d quad.Direction) map[int64]*llrb.LLRB {
-	switch d {
-	case quad.Subject:
-		return qdi.subject
-	case quad.Object:
-		return qdi.object
-	case quad.Predicate:
-		return qdi.predicate
-	case quad.Label:
-		return qdi.label
+func (qdi QuadDirectionIndex) Tree(d quad.Direction, id int64) *b.Tree {
+	if d < quad.Subject || d > quad.Label {
+		panic("illegal direction")
 	}
-	panic("illegal direction")
-}
-
-func (qdi *QuadDirectionIndex) GetOrCreate(d quad.Direction, id int64) *llrb.LLRB {
-	directionIndex := qdi.GetForDir(d)
-	if _, ok := directionIndex[id]; !ok {
-		directionIndex[id] = llrb.New()
+	tree, ok := qdi.index[d-1][id]
+	if !ok {
+		tree = b.TreeNew(cmp)
+		qdi.index[d-1][id] = tree
 	}
-	return directionIndex[id]
+	return tree
 }
 
-func (qdi *QuadDirectionIndex) Get(d quad.Direction, id int64) (*llrb.LLRB, bool) {
-	directionIndex := qdi.GetForDir(d)
-	tree, exists := directionIndex[id]
-	return tree, exists
+func (qdi QuadDirectionIndex) Get(d quad.Direction, id int64) (*b.Tree, bool) {
+	if d < quad.Subject || d > quad.Label {
+		panic("illegal direction")
+	}
+	tree, ok := qdi.index[d-1][id]
+	return tree, ok
 }
 
 type LogEntry struct {
@@ -88,21 +77,21 @@ type TripleStore struct {
 	log           []LogEntry
 	size          int64
 	index         QuadDirectionIndex
-	// vip_index map[string]map[int64]map[string]map[int64]*llrb.Tree
+	// vip_index map[string]map[int64]map[string]map[int64]*b.Tree
 }
 
 func newTripleStore() *TripleStore {
-	var ts TripleStore
-	ts.idMap = make(map[string]int64)
-	ts.revIdMap = make(map[int64]string)
-	ts.log = make([]LogEntry, 1, 200)
+	return &TripleStore{
+		idMap:    make(map[string]int64),
+		revIdMap: make(map[int64]string),
 
-	// Sentinel null entry so indices start at 1
-	ts.log[0] = LogEntry{}
-	ts.index = *NewQuadDirectionIndex()
-	ts.idCounter = 1
-	ts.quadIdCounter = 1
-	return &ts
+		// Sentinel null entry so indices start at 1
+		log: make([]LogEntry, 1, 200),
+
+		index:         NewQuadDirectionIndex(),
+		idCounter:     1,
+		quadIdCounter: 1,
+	}
 }
 
 func (ts *TripleStore) ApplyDeltas(deltas []graph.Delta) error {
@@ -120,46 +109,46 @@ func (ts *TripleStore) ApplyDeltas(deltas []graph.Delta) error {
 	return nil
 }
 
-func (ts *TripleStore) quadExists(t quad.Quad) (bool, int64) {
-	smallest := -1
-	var smallest_tree *llrb.LLRB
+const maxInt = int(^uint(0) >> 1)
+
+func (ts *TripleStore) indexOf(t quad.Quad) (int64, bool) {
+	min := maxInt
+	var tree *b.Tree
 	for d := quad.Subject; d <= quad.Label; d++ {
 		sid := t.Get(d)
 		if d == quad.Label && sid == "" {
 			continue
 		}
 		id, ok := ts.idMap[sid]
-		// If we've never heard about a node, it most not exist
+		// If we've never heard about a node, it must not exist
 		if !ok {
-			return false, 0
+			return 0, false
 		}
-		index, exists := ts.index.Get(d, id)
-		if !exists {
+		index, ok := ts.index.Get(d, id)
+		if !ok {
 			// If it's never been indexed in this direction, it can't exist.
-			return false, 0
+			return 0, false
 		}
-		if smallest == -1 || index.Len() < smallest {
-			smallest = index.Len()
-			smallest_tree = index
+		if l := index.Len(); l < min {
+			min, tree = l, index
 		}
 	}
-	it := NewLlrbIterator(smallest_tree, "", ts)
+	it := NewIterator(tree, "", ts)
 
 	for it.Next() {
 		val := it.Result()
 		if t == ts.log[val.(int64)].Quad {
-			return true, val.(int64)
+			return val.(int64), true
 		}
 	}
-	return false, 0
+	return 0, false
 }
 
 func (ts *TripleStore) AddDelta(d graph.Delta) error {
-	if exists, _ := ts.quadExists(d.Quad); exists {
+	if _, exists := ts.indexOf(d.Quad); exists {
 		return graph.ErrQuadExists
 	}
-	var quadID int64
-	quadID = ts.quadIdCounter
+	qid := ts.quadIdCounter
 	ts.log = append(ts.log, LogEntry{Delta: d})
 	ts.size++
 	ts.quadIdCounter++
@@ -181,8 +170,8 @@ func (ts *TripleStore) AddDelta(d graph.Delta) error {
 			continue
 		}
 		id := ts.idMap[d.Quad.Get(dir)]
-		tree := ts.index.GetOrCreate(dir, id)
-		tree.ReplaceOrInsert(Int64(quadID))
+		tree := ts.index.Tree(dir, id)
+		tree.Set(qid, struct{}{})
 	}
 
 	// TODO(barakmich): Add VIP indexing
@@ -190,15 +179,12 @@ func (ts *TripleStore) AddDelta(d graph.Delta) error {
 }
 
 func (ts *TripleStore) RemoveDelta(d graph.Delta) error {
-	var prevQuadID int64
-	var exists bool
-	prevQuadID = 0
-	if exists, prevQuadID = ts.quadExists(d.Quad); !exists {
+	prevQuadID, exists := ts.indexOf(d.Quad)
+	if !exists {
 		return graph.ErrQuadNotExist
 	}
 
-	var quadID int64
-	quadID = ts.quadIdCounter
+	quadID := ts.quadIdCounter
 	ts.log = append(ts.log, LogEntry{Delta: d})
 	ts.log[prevQuadID].DeletedBy = quadID
 	ts.size--
@@ -214,7 +200,7 @@ func (ts *TripleStore) TripleIterator(d quad.Direction, value graph.Value) graph
 	index, ok := ts.index.Get(d, value.(int64))
 	data := fmt.Sprintf("dir:%s val:%d", d, value.(int64))
 	if ok {
-		return NewLlrbIterator(index, data, ts)
+		return NewIterator(index, data, ts)
 	}
 	return &iterator.Null{}
 }
@@ -260,4 +246,5 @@ func (ts *TripleStore) TripleDirection(val graph.Value, d quad.Direction) graph.
 func (ts *TripleStore) NodesAllIterator() graph.Iterator {
 	return NewMemstoreNodesAllIterator(ts)
 }
+
 func (ts *TripleStore) Close() {}