vendor/github.com/btcsuite/btcd/database/internal/treap/mutable.go

   1 // Copyright (c) 2015-2016 The btcsuite developers
   2 // Use of this source code is governed by an ISC
   3 // license that can be found in the LICENSE file.
   4
   5 package treap
   6
   7 import (
   8         "bytes"
   9         "math/rand"
  10 )
  11
  12 // Mutable represents a treap data structure which is used to hold ordered
  13 // key/value pairs using a combination of binary search tree and heap semantics.
  14 // It is a self-organizing and randomized data structure that doesn't require
  15 // complex operations to maintain balance.  Search, insert, and delete
  16 // operations are all O(log n).
  17 type Mutable struct {
  18         root  *treapNode
  19         count int
  20
  21         // totalSize is the best estimate of the total size of of all data in
  22         // the treap including the keys, values, and node sizes.
  23         totalSize uint64
  24 }
  25
  26 // Len returns the number of items stored in the treap.
  27 func (t *Mutable) Len() int {
  28         return t.count
  29 }
  30
  31 // Size returns a best estimate of the total number of bytes the treap is
  32 // consuming including all of the fields used to represent the nodes as well as
  33 // the size of the keys and values.  Shared values are not detected, so the
  34 // returned size assumes each value is pointing to different memory.
  35 func (t *Mutable) Size() uint64 {
  36         return t.totalSize
  37 }
  38
  39 // get returns the treap node that contains the passed key and its parent.  When
  40 // the found node is the root of the tree, the parent will be nil.  When the key
  41 // does not exist, both the node and the parent will be nil.
  42 func (t *Mutable) get(key []byte) (*treapNode, *treapNode) {
  43         var parent *treapNode
  44         for node := t.root; node != nil; {
  45                 // Traverse left or right depending on the result of the
  46                 // comparison.
  47                 compareResult := bytes.Compare(key, node.key)
  48                 if compareResult < 0 {
  49                         parent = node
  50                         node = node.left
  51                         continue
  52                 }
  53                 if compareResult > 0 {
  54                         parent = node
  55                         node = node.right
  56                         continue
  57                 }
  58
  59                 // The key exists.
  60                 return node, parent
  61         }
  62
  63         // A nil node was reached which means the key does not exist.
  64         return nil, nil
  65 }
  66
  67 // Has returns whether or not the passed key exists.
  68 func (t *Mutable) Has(key []byte) bool {
  69         if node, _ := t.get(key); node != nil {
  70                 return true
  71         }
  72         return false
  73 }
  74
  75 // Get returns the value for the passed key.  The function will return nil when
  76 // the key does not exist.
  77 func (t *Mutable) Get(key []byte) []byte {
  78         if node, _ := t.get(key); node != nil {
  79                 return node.value
  80         }
  81         return nil
  82 }
  83
  84 // relinkGrandparent relinks the node into the treap after it has been rotated
  85 // by changing the passed grandparent's left or right pointer, depending on
  86 // where the old parent was, to point at the passed node.  Otherwise, when there
  87 // is no grandparent, it means the node is now the root of the tree, so update
  88 // it accordingly.
  89 func (t *Mutable) relinkGrandparent(node, parent, grandparent *treapNode) {
  90         // The node is now the root of the tree when there is no grandparent.
  91         if grandparent == nil {
  92                 t.root = node
  93                 return
  94         }
  95
  96         // Relink the grandparent's left or right pointer based on which side
  97         // the old parent was.
  98         if grandparent.left == parent {
  99                 grandparent.left = node
 100         } else {
 101                 grandparent.right = node
 102         }
 103 }
 104
 105 // Put inserts the passed key/value pair.
 106 func (t *Mutable) Put(key, value []byte) {
 107         // Use an empty byte slice for the value when none was provided.  This
 108         // ultimately allows key existence to be determined from the value since
 109         // an empty byte slice is distinguishable from nil.
 110         if value == nil {
 111                 value = emptySlice
 112         }
 113
 114         // The node is the root of the tree if there isn't already one.
 115         if t.root == nil {
 116                 node := newTreapNode(key, value, rand.Int())
 117                 t.count = 1
 118                 t.totalSize = nodeSize(node)
 119                 t.root = node
 120                 return
 121         }
 122
 123         // Find the binary tree insertion point and construct a list of parents
 124         // while doing so.  When the key matches an entry already in the treap,
 125         // just update its value and return.
 126         var parents parentStack
 127         var compareResult int
 128         for node := t.root; node != nil; {
 129                 parents.Push(node)
 130                 compareResult = bytes.Compare(key, node.key)
 131                 if compareResult < 0 {
 132                         node = node.left
 133                         continue
 134                 }
 135                 if compareResult > 0 {
 136                         node = node.right
 137                         continue
 138                 }
 139
 140                 // The key already exists, so update its value.
 141                 t.totalSize -= uint64(len(node.value))
 142                 t.totalSize += uint64(len(value))
 143                 node.value = value
 144                 return
 145         }
 146
 147         // Link the new node into the binary tree in the correct position.
 148         node := newTreapNode(key, value, rand.Int())
 149         t.count++
 150         t.totalSize += nodeSize(node)
 151         parent := parents.At(0)
 152         if compareResult < 0 {
 153                 parent.left = node
 154         } else {
 155                 parent.right = node
 156         }
 157
 158         // Perform any rotations needed to maintain the min-heap.
 159         for parents.Len() > 0 {
 160                 // There is nothing left to do when the node's priority is
 161                 // greater than or equal to its parent's priority.
 162                 parent = parents.Pop()
 163                 if node.priority >= parent.priority {
 164                         break
 165                 }
 166
 167                 // Perform a right rotation if the node is on the left side or
 168                 // a left rotation if the node is on the right side.
 169                 if parent.left == node {
 170                         node.right, parent.left = parent, node.right
 171                 } else {
 172                         node.left, parent.right = parent, node.left
 173                 }
 174                 t.relinkGrandparent(node, parent, parents.At(0))
 175         }
 176 }
 177
 178 // Delete removes the passed key if it exists.
 179 func (t *Mutable) Delete(key []byte) {
 180         // Find the node for the key along with its parent.  There is nothing to
 181         // do if the key does not exist.
 182         node, parent := t.get(key)
 183         if node == nil {
 184                 return
 185         }
 186
 187         // When the only node in the tree is the root node and it is the one
 188         // being deleted, there is nothing else to do besides removing it.
 189         if parent == nil && node.left == nil && node.right == nil {
 190                 t.root = nil
 191                 t.count = 0
 192                 t.totalSize = 0
 193                 return
 194         }
 195
 196         // Perform rotations to move the node to delete to a leaf position while
 197         // maintaining the min-heap.
 198         var isLeft bool
 199         var child *treapNode
 200         for node.left != nil || node.right != nil {
 201                 // Choose the child with the higher priority.
 202                 if node.left == nil {
 203                         child = node.right
 204                         isLeft = false
 205                 } else if node.right == nil {
 206                         child = node.left
 207                         isLeft = true
 208                 } else if node.left.priority >= node.right.priority {
 209                         child = node.left
 210                         isLeft = true
 211                 } else {
 212                         child = node.right
 213                         isLeft = false
 214                 }
 215
 216                 // Rotate left or right depending on which side the child node
 217                 // is on.  This has the effect of moving the node to delete
 218                 // towards the bottom of the tree while maintaining the
 219                 // min-heap.
 220                 if isLeft {
 221                         child.right, node.left = node, child.right
 222                 } else {
 223                         child.left, node.right = node, child.left
 224                 }
 225                 t.relinkGrandparent(child, node, parent)
 226
 227                 // The parent for the node to delete is now what was previously
 228                 // its child.
 229                 parent = child
 230         }
 231
 232         // Delete the node, which is now a leaf node, by disconnecting it from
 233         // its parent.
 234         if parent.right == node {
 235                 parent.right = nil
 236         } else {
 237                 parent.left = nil
 238         }
 239         t.count--
 240         t.totalSize -= nodeSize(node)
 241 }
 242
 243 // ForEach invokes the passed function with every key/value pair in the treap
 244 // in ascending order.
 245 func (t *Mutable) ForEach(fn func(k, v []byte) bool) {
 246         // Add the root node and all children to the left of it to the list of
 247         // nodes to traverse and loop until they, and all of their child nodes,
 248         // have been traversed.
 249         var parents parentStack
 250         for node := t.root; node != nil; node = node.left {
 251                 parents.Push(node)
 252         }
 253         for parents.Len() > 0 {
 254                 node := parents.Pop()
 255                 if !fn(node.key, node.value) {
 256                         return
 257                 }
 258
 259                 // Extend the nodes to traverse by all children to the left of
 260                 // the current node's right child.
 261                 for node := node.right; node != nil; node = node.left {
 262                         parents.Push(node)
 263                 }
 264         }
 265 }
 266
 267 // Reset efficiently removes all items in the treap.
 268 func (t *Mutable) Reset() {
 269         t.count = 0
 270         t.totalSize = 0
 271         t.root = nil
 272 }
 273
 274 // NewMutable returns a new empty mutable treap ready for use.  See the
 275 // documentation for the Mutable structure for more details.
 276 func NewMutable() *Mutable {
 277         return &Mutable{}
 278 }