diff --git a/common/ibbs_tree.go b/common/ibbs_tree.go new file mode 100644 index 000000000..6ef2fd831 --- /dev/null +++ b/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) + } +} diff --git a/common/ibbs_tree_test.go b/common/ibbs_tree_test.go new file mode 100644 index 000000000..74748bf77 --- /dev/null +++ b/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) + } +}