+++ /dev/null
-// Copyright (c) 2015-2016 The btcsuite developers
-// Use of this source code is governed by an ISC
-// license that can be found in the LICENSE file.
-
-package treap
-
-import (
- "bytes"
- "math/rand"
-)
-
-// Mutable represents a treap data structure which is used to hold ordered
-// key/value pairs using a combination of binary search tree and heap semantics.
-// It is a self-organizing and randomized data structure that doesn't require
-// complex operations to maintain balance. Search, insert, and delete
-// operations are all O(log n).
-type Mutable struct {
- root *treapNode
- count int
-
- // totalSize is the best estimate of the total size of of all data in
- // the treap including the keys, values, and node sizes.
- totalSize uint64
-}
-
-// Len returns the number of items stored in the treap.
-func (t *Mutable) Len() int {
- return t.count
-}
-
-// Size returns a best estimate of the total number of bytes the treap is
-// consuming including all of the fields used to represent the nodes as well as
-// the size of the keys and values. Shared values are not detected, so the
-// returned size assumes each value is pointing to different memory.
-func (t *Mutable) Size() uint64 {
- return t.totalSize
-}
-
-// get returns the treap node that contains the passed key and its parent. When
-// the found node is the root of the tree, the parent will be nil. When the key
-// does not exist, both the node and the parent will be nil.
-func (t *Mutable) get(key []byte) (*treapNode, *treapNode) {
- var parent *treapNode
- for node := t.root; node != nil; {
- // Traverse left or right depending on the result of the
- // comparison.
- compareResult := bytes.Compare(key, node.key)
- if compareResult < 0 {
- parent = node
- node = node.left
- continue
- }
- if compareResult > 0 {
- parent = node
- node = node.right
- continue
- }
-
- // The key exists.
- return node, parent
- }
-
- // A nil node was reached which means the key does not exist.
- return nil, nil
-}
-
-// Has returns whether or not the passed key exists.
-func (t *Mutable) Has(key []byte) bool {
- if node, _ := t.get(key); node != nil {
- return true
- }
- return false
-}
-
-// Get returns the value for the passed key. The function will return nil when
-// the key does not exist.
-func (t *Mutable) Get(key []byte) []byte {
- if node, _ := t.get(key); node != nil {
- return node.value
- }
- return nil
-}
-
-// relinkGrandparent relinks the node into the treap after it has been rotated
-// by changing the passed grandparent's left or right pointer, depending on
-// where the old parent was, to point at the passed node. Otherwise, when there
-// is no grandparent, it means the node is now the root of the tree, so update
-// it accordingly.
-func (t *Mutable) relinkGrandparent(node, parent, grandparent *treapNode) {
- // The node is now the root of the tree when there is no grandparent.
- if grandparent == nil {
- t.root = node
- return
- }
-
- // Relink the grandparent's left or right pointer based on which side
- // the old parent was.
- if grandparent.left == parent {
- grandparent.left = node
- } else {
- grandparent.right = node
- }
-}
-
-// Put inserts the passed key/value pair.
-func (t *Mutable) Put(key, value []byte) {
- // Use an empty byte slice for the value when none was provided. This
- // ultimately allows key existence to be determined from the value since
- // an empty byte slice is distinguishable from nil.
- if value == nil {
- value = emptySlice
- }
-
- // The node is the root of the tree if there isn't already one.
- if t.root == nil {
- node := newTreapNode(key, value, rand.Int())
- t.count = 1
- t.totalSize = nodeSize(node)
- t.root = node
- return
- }
-
- // Find the binary tree insertion point and construct a list of parents
- // while doing so. When the key matches an entry already in the treap,
- // just update its value and return.
- var parents parentStack
- var compareResult int
- for node := t.root; node != nil; {
- parents.Push(node)
- compareResult = bytes.Compare(key, node.key)
- if compareResult < 0 {
- node = node.left
- continue
- }
- if compareResult > 0 {
- node = node.right
- continue
- }
-
- // The key already exists, so update its value.
- t.totalSize -= uint64(len(node.value))
- t.totalSize += uint64(len(value))
- node.value = value
- return
- }
-
- // Link the new node into the binary tree in the correct position.
- node := newTreapNode(key, value, rand.Int())
- t.count++
- t.totalSize += nodeSize(node)
- parent := parents.At(0)
- if compareResult < 0 {
- parent.left = node
- } else {
- parent.right = node
- }
-
- // Perform any rotations needed to maintain the min-heap.
- for parents.Len() > 0 {
- // There is nothing left to do when the node's priority is
- // greater than or equal to its parent's priority.
- parent = parents.Pop()
- if node.priority >= parent.priority {
- break
- }
-
- // Perform a right rotation if the node is on the left side or
- // a left rotation if the node is on the right side.
- if parent.left == node {
- node.right, parent.left = parent, node.right
- } else {
- node.left, parent.right = parent, node.left
- }
- t.relinkGrandparent(node, parent, parents.At(0))
- }
-}
-
-// Delete removes the passed key if it exists.
-func (t *Mutable) Delete(key []byte) {
- // Find the node for the key along with its parent. There is nothing to
- // do if the key does not exist.
- node, parent := t.get(key)
- if node == nil {
- return
- }
-
- // When the only node in the tree is the root node and it is the one
- // being deleted, there is nothing else to do besides removing it.
- if parent == nil && node.left == nil && node.right == nil {
- t.root = nil
- t.count = 0
- t.totalSize = 0
- return
- }
-
- // Perform rotations to move the node to delete to a leaf position while
- // maintaining the min-heap.
- var isLeft bool
- var child *treapNode
- for node.left != nil || node.right != nil {
- // Choose the child with the higher priority.
- if node.left == nil {
- child = node.right
- isLeft = false
- } else if node.right == nil {
- child = node.left
- isLeft = true
- } else if node.left.priority >= node.right.priority {
- child = node.left
- isLeft = true
- } else {
- child = node.right
- isLeft = false
- }
-
- // Rotate left or right depending on which side the child node
- // is on. This has the effect of moving the node to delete
- // towards the bottom of the tree while maintaining the
- // min-heap.
- if isLeft {
- child.right, node.left = node, child.right
- } else {
- child.left, node.right = node, child.left
- }
- t.relinkGrandparent(child, node, parent)
-
- // The parent for the node to delete is now what was previously
- // its child.
- parent = child
- }
-
- // Delete the node, which is now a leaf node, by disconnecting it from
- // its parent.
- if parent.right == node {
- parent.right = nil
- } else {
- parent.left = nil
- }
- t.count--
- t.totalSize -= nodeSize(node)
-}
-
-// ForEach invokes the passed function with every key/value pair in the treap
-// in ascending order.
-func (t *Mutable) ForEach(fn func(k, v []byte) bool) {
- // Add the root node and all children to the left of it to the list of
- // nodes to traverse and loop until they, and all of their child nodes,
- // have been traversed.
- var parents parentStack
- for node := t.root; node != nil; node = node.left {
- parents.Push(node)
- }
- for parents.Len() > 0 {
- node := parents.Pop()
- if !fn(node.key, node.value) {
- return
- }
-
- // Extend the nodes to traverse by all children to the left of
- // the current node's right child.
- for node := node.right; node != nil; node = node.left {
- parents.Push(node)
- }
- }
-}
-
-// Reset efficiently removes all items in the treap.
-func (t *Mutable) Reset() {
- t.count = 0
- t.totalSize = 0
- t.root = nil
-}
-
-// NewMutable returns a new empty mutable treap ready for use. See the
-// documentation for the Mutable structure for more details.
-func NewMutable() *Mutable {
- return &Mutable{}
-}