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)
+	}
+}