added ibbs_tree

common/ibbs_tree.go

@@ -0,0 +1,266 @@
+package common
+
+import ()
+
+// This immutable balanced binary tree happens to be an
+// immutable AVL+ tree, adapted from tendermint/merkle.
+// Unlike that one, this is in-memory, non-merkleized,
+// and nodes can be nil to signify an empty tree.
+type IBBSTree struct {
+	key    uint64
+	value  interface{}
+	size   uint64
+	height uint8
+	left   *IBBSTree
+	right  *IBBSTree
+}
+
+// Creates an empty tree.
+func NewIBBSTree() *IBBSTree {
+	return nil
+}
+
+// Creates a single tree node from key and value.
+func NewIBBSTreeNode(key uint64, value interface{}) *IBBSTree {
+	return &IBBSTree{
+		key:   key,
+		value: value,
+		size:  1,
+	}
+}
+
+func (self *IBBSTree) Copy() *IBBSTree {
+	if self == nil {
+		return nil
+	}
+	return &IBBSTree{
+		key:    self.key,
+		value:  self.value,
+		size:   self.size,
+		height: self.height,
+		left:   self.left,
+		right:  self.right,
+	}
+}
+
+func (self *IBBSTree) Size() uint64 {
+	if self == nil {
+		return 0
+	}
+	return self.size
+}
+
+func (self *IBBSTree) Height() uint8 {
+	if self == nil {
+		return 0
+	}
+	return self.height
+}
+
+func (self *IBBSTree) Has(key uint64) (has bool) {
+	if self == nil {
+		return false
+	}
+	if self.key == key {
+		return true
+	}
+	if self.height == 0 {
+		return false
+	} else {
+		if key < self.key {
+			return self.left.Has(key)
+		} else {
+			return self.right.Has(key)
+		}
+	}
+}
+
+func (self *IBBSTree) Get(key uint64) (value interface{}) {
+	if self == nil {
+		return nil
+	}
+	if self.height == 0 {
+		if self.key == key {
+			return self.value
+		} else {
+			return nil
+		}
+	} else {
+		if key < self.key {
+			return self.left.Get(key)
+		} else {
+			return self.right.Get(key)
+		}
+	}
+}
+
+func (self *IBBSTree) Set(key uint64, value interface{}) (_ *IBBSTree, updated bool) {
+	if self == nil {
+		return NewIBBSTreeNode(key, value), false
+	}
+	if self.height == 0 {
+		if key < self.key {
+			return &IBBSTree{
+				key:    self.key,
+				height: 1,
+				size:   2,
+				left:   NewIBBSTreeNode(key, value),
+				right:  self,
+			}, false
+		} else if self.key == key {
+			return NewIBBSTreeNode(key, value), true
+		} else {
+			return &IBBSTree{
+				key:    key,
+				height: 1,
+				size:   2,
+				left:   self,
+				right:  NewIBBSTreeNode(key, value),
+			}, false
+		}
+	} else {
+		self = self.Copy()
+		if key < self.key {
+			self.left, updated = self.left.Set(key, value)
+		} else {
+			self.right, updated = self.right.Set(key, value)
+		}
+		if updated {
+			return self, updated
+		} else {
+			self.calcHeightAndSize()
+			return self.balance(), updated
+		}
+	}
+}
+
+func (self *IBBSTree) Remove(key uint64) (newSelf *IBBSTree, value interface{}, removed bool) {
+	newSelf, _, _, value, removed = self.remove(key)
+	return
+}
+
+// newKey: new leftmost leaf key for tree after successfully removing 'key' if changed.
+func (self *IBBSTree) remove(key uint64) (newSelf *IBBSTree, hasNewKey bool, newKey uint64, value interface{}, removed bool) {
+	if self == nil {
+		return nil, false, 0, nil, false
+	}
+	if self.height == 0 {
+		if self.key == key {
+			return nil, false, 0, self.value, true
+		} else {
+			return self, false, 0, nil, false
+		}
+	} else {
+		if key < self.key {
+			var newLeft *IBBSTree
+			newLeft, hasNewKey, newKey, value, removed = self.left.remove(key)
+			if !removed {
+				return self, false, 0, value, false
+			} else if newLeft == nil { // left node held value, was removed
+				return self.right, true, self.key, value, true
+			}
+			self = self.Copy()
+			self.left = newLeft
+		} else {
+			var newRight *IBBSTree
+			newRight, hasNewKey, newKey, value, removed = self.right.remove(key)
+			if !removed {
+				return self, false, 0, value, false
+			} else if newRight == nil { // right node held value, was removed
+				return self.left, false, 0, value, true
+			}
+			self = self.Copy()
+			self.right = newRight
+			if hasNewKey {
+				self.key = newKey
+				hasNewKey = false
+				newKey = 0
+			}
+		}
+		self.calcHeightAndSize()
+		return self.balance(), hasNewKey, newKey, value, true
+	}
+}
+
+func (self *IBBSTree) rotateRight() *IBBSTree {
+	self = self.Copy()
+	sl := self.left.Copy()
+	slr := sl.right
+
+	sl.right = self
+	self.left = slr
+
+	self.calcHeightAndSize()
+	sl.calcHeightAndSize()
+
+	return sl
+}
+
+func (self *IBBSTree) rotateLeft() *IBBSTree {
+	self = self.Copy()
+	sr := self.right.Copy()
+	srl := sr.left
+
+	sr.left = self
+	self.right = srl
+
+	self.calcHeightAndSize()
+	sr.calcHeightAndSize()
+
+	return sr
+}
+
+func (self *IBBSTree) calcHeightAndSize() {
+	self.height = MaxUint8(self.left.Height(), self.right.Height()) + 1
+	self.size = self.left.Size() + self.right.Size()
+}
+
+func (self *IBBSTree) calcBalance() int {
+	return int(self.left.Height()) - int(self.right.Height())
+}
+
+func (self *IBBSTree) balance() (newSelf *IBBSTree) {
+	balance := self.calcBalance()
+	if balance > 1 {
+		if self.left.calcBalance() >= 0 {
+			// Left Left Case
+			return self.rotateRight()
+		} else {
+			// Left Right Case
+			self = self.Copy()
+			self.left = self.left.rotateLeft()
+			//self.calcHeightAndSize()
+			return self.rotateRight()
+		}
+	}
+	if balance < -1 {
+		if self.right.calcBalance() <= 0 {
+			// Right Right Case
+			return self.rotateLeft()
+		} else {
+			// Right Left Case
+			self = self.Copy()
+			self.right = self.right.rotateRight()
+			//self.calcHeightAndSize()
+			return self.rotateLeft()
+		}
+	}
+	// Nothing changed
+	return self
+}
+
+// Iteration stops when stop is returned.
+func (self *IBBSTree) Iterate(cb func(uint64, interface{}) bool) bool {
+	if self == nil {
+		return false
+	}
+	if self.height == 0 {
+		return cb(self.key, self.value)
+	} else {
+		stop := self.left.Iterate(cb)
+		if stop {
+			return stop
+		}
+		return self.right.Iterate(cb)
+	}
+}

common/ibbs_tree_test.go

@@ -0,0 +1,203 @@
+package common
+
+import (
+	"fmt"
+	"runtime"
+	"testing"
+)
+
+func init() {
+	// TODO: seed rand?
+}
+
+func (self *IBBSTree) lmd() *IBBSTree {
+	if self.height == 0 {
+		return self
+	}
+	return self.left.lmd()
+}
+
+func TestUnit(t *testing.T) {
+
+	// Convenience for a new node
+	N := func(l, r interface{}) *IBBSTree {
+		var left, right *IBBSTree
+		if _, ok := l.(*IBBSTree); ok {
+			left = l.(*IBBSTree)
+		} else {
+			left = NewIBBSTreeNode(uint64(l.(int)), nil)
+		}
+		if _, ok := r.(*IBBSTree); ok {
+			right = r.(*IBBSTree)
+		} else {
+			right = NewIBBSTreeNode(uint64(r.(int)), nil)
+		}
+
+		n := &IBBSTree{
+			key:   right.lmd().key,
+			left:  left,
+			right: right,
+		}
+		n.calcHeightAndSize()
+		return n
+	}
+
+	// Convenience for simple printing of keys & tree structure
+	var P func(*IBBSTree) string
+	P = func(n *IBBSTree) string {
+		if n.height == 0 {
+			return fmt.Sprintf("%v", n.key)
+		} else {
+			return fmt.Sprintf("(%v %v)", P(n.left), P(n.right))
+		}
+	}
+
+	expectSet := func(n *IBBSTree, i uint64, repr string) {
+		n2, updated := n.Set(i, nil)
+		// ensure node was added & structure is as expected.
+		if updated == true || P(n2) != repr {
+			t.Fatalf("Adding %v to %v:\nExpected         %v\nUnexpectedly got %v updated:%v",
+				i, P(n), repr, P(n2), updated)
+		}
+	}
+
+	expectRemove := func(n *IBBSTree, i uint64, repr string) {
+		n2, value, removed := n.Remove(i)
+		// ensure node was removed & structure is as expected.
+		if value != nil || P(n2) != repr || !removed {
+			t.Fatalf("Removing %v from %v:\nExpected         %v\nUnexpectedly got %v value:%v err:%v",
+				i, P(n), repr, P(n2), value, removed)
+		}
+	}
+
+	//////// Test Set cases:
+
+	// Case 1:
+	n1 := N(4, 20)
+
+	expectSet(n1, 8, "((4 8) 20)")
+	expectSet(n1, 25, "(4 (20 25))")
+
+	n2 := N(4, N(20, 25))
+
+	expectSet(n2, 8, "((4 8) (20 25))")
+	expectSet(n2, 30, "((4 20) (25 30))")
+
+	n3 := N(N(1, 2), 6)
+
+	expectSet(n3, 4, "((1 2) (4 6))")
+	expectSet(n3, 8, "((1 2) (6 8))")
+
+	n4 := N(N(1, 2), N(N(5, 6), N(7, 9)))
+
+	expectSet(n4, 8, "(((1 2) (5 6)) ((7 8) 9))")
+	expectSet(n4, 10, "(((1 2) (5 6)) (7 (9 10)))")
+
+	//////// Test Remove cases:
+
+	n10 := N(N(1, 2), 3)
+
+	expectRemove(n10, 2, "(1 3)")
+	expectRemove(n10, 3, "(1 2)")
+
+	n11 := N(N(N(1, 2), 3), N(4, 5))
+
+	expectRemove(n11, 4, "((1 2) (3 5))")
+	expectRemove(n11, 3, "((1 2) (4 5))")
+
+}
+
+func TestIntegration(t *testing.T) {
+
+	type record struct {
+		key   uint64
+		value string
+	}
+
+	records := make([]*record, 400)
+	var tree *IBBSTree = NewIBBSTree()
+	var val interface{}
+	var removed bool
+	var updated bool
+
+	randomRecord := func() *record {
+		return &record{RandUInt64(), RandStr(20)}
+	}
+
+	for i := range records {
+		r := randomRecord()
+		records[i] = r
+		//t.Log("New record", r)
+		//PrintIBBSTree(tree.root)
+		tree, updated = tree.Set(r.key, "")
+		if updated {
+			t.Error("should have not been updated")
+		}
+		tree, updated = tree.Set(r.key, r.value)
+		if !updated {
+			t.Error("should have been updated")
+		}
+		if tree.Size() != uint64(i+1) {
+			t.Error("size was wrong", tree.Size(), i+1)
+		}
+	}
+
+	for _, r := range records {
+		if has := tree.Has(r.key); !has {
+			t.Error("Missing key", r.key)
+		}
+		if val := tree.Get(r.key); val.(string) != r.value {
+			t.Error("wrong value")
+		}
+	}
+
+	for i, x := range records {
+		if tree, val, removed = tree.Remove(x.key); !removed {
+			t.Error("not removed")
+		} else if val.(string) != x.value {
+			t.Error("wrong value")
+		}
+		for _, r := range records[i+1:] {
+			if has := tree.Has(r.key); !has {
+				t.Error("Missing key", r.key)
+			}
+			val := tree.Get(r.key)
+			if val.(string) != r.value {
+				t.Error("wrong value")
+			}
+		}
+		if tree.Size() != uint64(len(records)-(i+1)) {
+			t.Error("size was wrong", tree.Size(), (len(records) - (i + 1)))
+		}
+	}
+}
+
+func BenchmarkIBBSTree(b *testing.B) {
+	b.StopTimer()
+
+	type record struct {
+		key   uint64
+		value interface{}
+	}
+
+	randomRecord := func() *record {
+		return &record{RandUInt64(), RandUInt64()}
+	}
+
+	t := NewIBBSTree()
+	for i := 0; i < 1000000; i++ {
+		r := randomRecord()
+		t, _ = t.Set(r.key, r.value)
+	}
+
+	fmt.Println("ok, starting")
+
+	runtime.GC()
+
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		r := randomRecord()
+		t, _ = t.Set(r.key, r.value)
+		t, _, _ = t.Remove(r.key)
+	}
+}