vendor/github.com/hashicorp/go-immutable-radix/node.go

   1 package iradix
   2
   3 import (
   4         "bytes"
   5         "sort"
   6 )
   7
   8 // WalkFn is used when walking the tree. Takes a
   9 // key and value, returning if iteration should
  10 // be terminated.
  11 type WalkFn func(k []byte, v interface{}) bool
  12
  13 // leafNode is used to represent a value
  14 type leafNode struct {
  15         key []byte
  16         val interface{}
  17 }
  18
  19 // edge is used to represent an edge node
  20 type edge struct {
  21         label byte
  22         node  *Node
  23 }
  24
  25 // Node is an immutable node in the radix tree
  26 type Node struct {
  27         // leaf is used to store possible leaf
  28         leaf *leafNode
  29
  30         // prefix is the common prefix we ignore
  31         prefix []byte
  32
  33         // Edges should be stored in-order for iteration.
  34         // We avoid a fully materialized slice to save memory,
  35         // since in most cases we expect to be sparse
  36         edges edges
  37 }
  38
  39 func (n *Node) isLeaf() bool {
  40         return n.leaf != nil
  41 }
  42
  43 func (n *Node) addEdge(e edge) {
  44         num := len(n.edges)
  45         idx := sort.Search(num, func(i int) bool {
  46                 return n.edges[i].label >= e.label
  47         })
  48         n.edges = append(n.edges, e)
  49         if idx != num {
  50                 copy(n.edges[idx+1:], n.edges[idx:num])
  51                 n.edges[idx] = e
  52         }
  53 }
  54
  55 func (n *Node) replaceEdge(e edge) {
  56         num := len(n.edges)
  57         idx := sort.Search(num, func(i int) bool {
  58                 return n.edges[i].label >= e.label
  59         })
  60         if idx < num && n.edges[idx].label == e.label {
  61                 n.edges[idx].node = e.node
  62                 return
  63         }
  64         panic("replacing missing edge")
  65 }
  66
  67 func (n *Node) getEdge(label byte) (int, *Node) {
  68         num := len(n.edges)
  69         idx := sort.Search(num, func(i int) bool {
  70                 return n.edges[i].label >= label
  71         })
  72         if idx < num && n.edges[idx].label == label {
  73                 return idx, n.edges[idx].node
  74         }
  75         return -1, nil
  76 }
  77
  78 func (n *Node) delEdge(label byte) {
  79         num := len(n.edges)
  80         idx := sort.Search(num, func(i int) bool {
  81                 return n.edges[i].label >= label
  82         })
  83         if idx < num && n.edges[idx].label == label {
  84                 copy(n.edges[idx:], n.edges[idx+1:])
  85                 n.edges[len(n.edges)-1] = edge{}
  86                 n.edges = n.edges[:len(n.edges)-1]
  87         }
  88 }
  89
  90 func (n *Node) mergeChild() {
  91         e := n.edges[0]
  92         child := e.node
  93         n.prefix = concat(n.prefix, child.prefix)
  94         if child.leaf != nil {
  95                 n.leaf = new(leafNode)
  96                 *n.leaf = *child.leaf
  97         } else {
  98                 n.leaf = nil
  99         }
 100         if len(child.edges) != 0 {
 101                 n.edges = make([]edge, len(child.edges))
 102                 copy(n.edges, child.edges)
 103         } else {
 104                 n.edges = nil
 105         }
 106 }
 107
 108 func (n *Node) Get(k []byte) (interface{}, bool) {
 109         search := k
 110         for {
 111                 // Check for key exhaution
 112                 if len(search) == 0 {
 113                         if n.isLeaf() {
 114                                 return n.leaf.val, true
 115                         }
 116                         break
 117                 }
 118
 119                 // Look for an edge
 120                 _, n = n.getEdge(search[0])
 121                 if n == nil {
 122                         break
 123                 }
 124
 125                 // Consume the search prefix
 126                 if bytes.HasPrefix(search, n.prefix) {
 127                         search = search[len(n.prefix):]
 128                 } else {
 129                         break
 130                 }
 131         }
 132         return nil, false
 133 }
 134
 135 // LongestPrefix is like Get, but instead of an
 136 // exact match, it will return the longest prefix match.
 137 func (n *Node) LongestPrefix(k []byte) ([]byte, interface{}, bool) {
 138         var last *leafNode
 139         search := k
 140         for {
 141                 // Look for a leaf node
 142                 if n.isLeaf() {
 143                         last = n.leaf
 144                 }
 145
 146                 // Check for key exhaution
 147                 if len(search) == 0 {
 148                         break
 149                 }
 150
 151                 // Look for an edge
 152                 _, n = n.getEdge(search[0])
 153                 if n == nil {
 154                         break
 155                 }
 156
 157                 // Consume the search prefix
 158                 if bytes.HasPrefix(search, n.prefix) {
 159                         search = search[len(n.prefix):]
 160                 } else {
 161                         break
 162                 }
 163         }
 164         if last != nil {
 165                 return last.key, last.val, true
 166         }
 167         return nil, nil, false
 168 }
 169
 170 // Minimum is used to return the minimum value in the tree
 171 func (n *Node) Minimum() ([]byte, interface{}, bool) {
 172         for {
 173                 if n.isLeaf() {
 174                         return n.leaf.key, n.leaf.val, true
 175                 }
 176                 if len(n.edges) > 0 {
 177                         n = n.edges[0].node
 178                 } else {
 179                         break
 180                 }
 181         }
 182         return nil, nil, false
 183 }
 184
 185 // Maximum is used to return the maximum value in the tree
 186 func (n *Node) Maximum() ([]byte, interface{}, bool) {
 187         for {
 188                 if num := len(n.edges); num > 0 {
 189                         n = n.edges[num-1].node
 190                         continue
 191                 }
 192                 if n.isLeaf() {
 193                         return n.leaf.key, n.leaf.val, true
 194                 } else {
 195                         break
 196                 }
 197         }
 198         return nil, nil, false
 199 }
 200
 201 // Iterator is used to return an iterator at
 202 // the given node to walk the tree
 203 func (n *Node) Iterator() *Iterator {
 204         return &Iterator{node: n}
 205 }
 206
 207 // Walk is used to walk the tree
 208 func (n *Node) Walk(fn WalkFn) {
 209         recursiveWalk(n, fn)
 210 }
 211
 212 // WalkPrefix is used to walk the tree under a prefix
 213 func (n *Node) WalkPrefix(prefix []byte, fn WalkFn) {
 214         search := prefix
 215         for {
 216                 // Check for key exhaution
 217                 if len(search) == 0 {
 218                         recursiveWalk(n, fn)
 219                         return
 220                 }
 221
 222                 // Look for an edge
 223                 _, n = n.getEdge(search[0])
 224                 if n == nil {
 225                         break
 226                 }
 227
 228                 // Consume the search prefix
 229                 if bytes.HasPrefix(search, n.prefix) {
 230                         search = search[len(n.prefix):]
 231
 232                 } else if bytes.HasPrefix(n.prefix, search) {
 233                         // Child may be under our search prefix
 234                         recursiveWalk(n, fn)
 235                         return
 236                 } else {
 237                         break
 238                 }
 239         }
 240 }
 241
 242 // WalkPath is used to walk the tree, but only visiting nodes
 243 // from the root down to a given leaf. Where WalkPrefix walks
 244 // all the entries *under* the given prefix, this walks the
 245 // entries *above* the given prefix.
 246 func (n *Node) WalkPath(path []byte, fn WalkFn) {
 247         search := path
 248         for {
 249                 // Visit the leaf values if any
 250                 if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
 251                         return
 252                 }
 253
 254                 // Check for key exhaution
 255                 if len(search) == 0 {
 256                         return
 257                 }
 258
 259                 // Look for an edge
 260                 _, n = n.getEdge(search[0])
 261                 if n == nil {
 262                         return
 263                 }
 264
 265                 // Consume the search prefix
 266                 if bytes.HasPrefix(search, n.prefix) {
 267                         search = search[len(n.prefix):]
 268                 } else {
 269                         break
 270                 }
 271         }
 272 }
 273
 274 // recursiveWalk is used to do a pre-order walk of a node
 275 // recursively. Returns true if the walk should be aborted
 276 func recursiveWalk(n *Node, fn WalkFn) bool {
 277         // Visit the leaf values if any
 278         if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
 279                 return true
 280         }
 281
 282         // Recurse on the children
 283         for _, e := range n.edges {
 284                 if recursiveWalk(e.node, fn) {
 285                         return true
 286                 }
 287         }
 288         return false
 289 }