vendor/github.com/armon/go-radix/radix.go

   1 package radix
   2
   3 import (
   4         "sort"
   5         "strings"
   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(s string, v interface{}) bool
  12
  13 // leafNode is used to represent a value
  14 type leafNode struct {
  15         key string
  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 type node struct {
  26         // leaf is used to store possible leaf
  27         leaf *leafNode
  28
  29         // prefix is the common prefix we ignore
  30         prefix string
  31
  32         // Edges should be stored in-order for iteration.
  33         // We avoid a fully materialized slice to save memory,
  34         // since in most cases we expect to be sparse
  35         edges edges
  36 }
  37
  38 func (n *node) isLeaf() bool {
  39         return n.leaf != nil
  40 }
  41
  42 func (n *node) addEdge(e edge) {
  43         n.edges = append(n.edges, e)
  44         n.edges.Sort()
  45 }
  46
  47 func (n *node) replaceEdge(e edge) {
  48         num := len(n.edges)
  49         idx := sort.Search(num, func(i int) bool {
  50                 return n.edges[i].label >= e.label
  51         })
  52         if idx < num && n.edges[idx].label == e.label {
  53                 n.edges[idx].node = e.node
  54                 return
  55         }
  56         panic("replacing missing edge")
  57 }
  58
  59 func (n *node) getEdge(label byte) *node {
  60         num := len(n.edges)
  61         idx := sort.Search(num, func(i int) bool {
  62                 return n.edges[i].label >= label
  63         })
  64         if idx < num && n.edges[idx].label == label {
  65                 return n.edges[idx].node
  66         }
  67         return nil
  68 }
  69
  70 func (n *node) delEdge(label byte) {
  71         num := len(n.edges)
  72         idx := sort.Search(num, func(i int) bool {
  73                 return n.edges[i].label >= label
  74         })
  75         if idx < num && n.edges[idx].label == label {
  76                 copy(n.edges[idx:], n.edges[idx+1:])
  77                 n.edges[len(n.edges)-1] = edge{}
  78                 n.edges = n.edges[:len(n.edges)-1]
  79         }
  80 }
  81
  82 type edges []edge
  83
  84 func (e edges) Len() int {
  85         return len(e)
  86 }
  87
  88 func (e edges) Less(i, j int) bool {
  89         return e[i].label < e[j].label
  90 }
  91
  92 func (e edges) Swap(i, j int) {
  93         e[i], e[j] = e[j], e[i]
  94 }
  95
  96 func (e edges) Sort() {
  97         sort.Sort(e)
  98 }
  99
 100 // Tree implements a radix tree. This can be treated as a
 101 // Dictionary abstract data type. The main advantage over
 102 // a standard hash map is prefix-based lookups and
 103 // ordered iteration,
 104 type Tree struct {
 105         root *node
 106         size int
 107 }
 108
 109 // New returns an empty Tree
 110 func New() *Tree {
 111         return NewFromMap(nil)
 112 }
 113
 114 // NewFromMap returns a new tree containing the keys
 115 // from an existing map
 116 func NewFromMap(m map[string]interface{}) *Tree {
 117         t := &Tree{root: &node{}}
 118         for k, v := range m {
 119                 t.Insert(k, v)
 120         }
 121         return t
 122 }
 123
 124 // Len is used to return the number of elements in the tree
 125 func (t *Tree) Len() int {
 126         return t.size
 127 }
 128
 129 // longestPrefix finds the length of the shared prefix
 130 // of two strings
 131 func longestPrefix(k1, k2 string) int {
 132         max := len(k1)
 133         if l := len(k2); l < max {
 134                 max = l
 135         }
 136         var i int
 137         for i = 0; i < max; i++ {
 138                 if k1[i] != k2[i] {
 139                         break
 140                 }
 141         }
 142         return i
 143 }
 144
 145 // Insert is used to add a newentry or update
 146 // an existing entry. Returns if updated.
 147 func (t *Tree) Insert(s string, v interface{}) (interface{}, bool) {
 148         var parent *node
 149         n := t.root
 150         search := s
 151         for {
 152                 // Handle key exhaution
 153                 if len(search) == 0 {
 154                         if n.isLeaf() {
 155                                 old := n.leaf.val
 156                                 n.leaf.val = v
 157                                 return old, true
 158                         }
 159
 160                         n.leaf = &leafNode{
 161                                 key: s,
 162                                 val: v,
 163                         }
 164                         t.size++
 165                         return nil, false
 166                 }
 167
 168                 // Look for the edge
 169                 parent = n
 170                 n = n.getEdge(search[0])
 171
 172                 // No edge, create one
 173                 if n == nil {
 174                         e := edge{
 175                                 label: search[0],
 176                                 node: &node{
 177                                         leaf: &leafNode{
 178                                                 key: s,
 179                                                 val: v,
 180                                         },
 181                                         prefix: search,
 182                                 },
 183                         }
 184                         parent.addEdge(e)
 185                         t.size++
 186                         return nil, false
 187                 }
 188
 189                 // Determine longest prefix of the search key on match
 190                 commonPrefix := longestPrefix(search, n.prefix)
 191                 if commonPrefix == len(n.prefix) {
 192                         search = search[commonPrefix:]
 193                         continue
 194                 }
 195
 196                 // Split the node
 197                 t.size++
 198                 child := &node{
 199                         prefix: search[:commonPrefix],
 200                 }
 201                 parent.replaceEdge(edge{
 202                         label: search[0],
 203                         node:  child,
 204                 })
 205
 206                 // Restore the existing node
 207                 child.addEdge(edge{
 208                         label: n.prefix[commonPrefix],
 209                         node:  n,
 210                 })
 211                 n.prefix = n.prefix[commonPrefix:]
 212
 213                 // Create a new leaf node
 214                 leaf := &leafNode{
 215                         key: s,
 216                         val: v,
 217                 }
 218
 219                 // If the new key is a subset, add to to this node
 220                 search = search[commonPrefix:]
 221                 if len(search) == 0 {
 222                         child.leaf = leaf
 223                         return nil, false
 224                 }
 225
 226                 // Create a new edge for the node
 227                 child.addEdge(edge{
 228                         label: search[0],
 229                         node: &node{
 230                                 leaf:   leaf,
 231                                 prefix: search,
 232                         },
 233                 })
 234                 return nil, false
 235         }
 236 }
 237
 238 // Delete is used to delete a key, returning the previous
 239 // value and if it was deleted
 240 func (t *Tree) Delete(s string) (interface{}, bool) {
 241         var parent *node
 242         var label byte
 243         n := t.root
 244         search := s
 245         for {
 246                 // Check for key exhaution
 247                 if len(search) == 0 {
 248                         if !n.isLeaf() {
 249                                 break
 250                         }
 251                         goto DELETE
 252                 }
 253
 254                 // Look for an edge
 255                 parent = n
 256                 label = search[0]
 257                 n = n.getEdge(label)
 258                 if n == nil {
 259                         break
 260                 }
 261
 262                 // Consume the search prefix
 263                 if strings.HasPrefix(search, n.prefix) {
 264                         search = search[len(n.prefix):]
 265                 } else {
 266                         break
 267                 }
 268         }
 269         return nil, false
 270
 271 DELETE:
 272         // Delete the leaf
 273         leaf := n.leaf
 274         n.leaf = nil
 275         t.size--
 276
 277         // Check if we should delete this node from the parent
 278         if parent != nil && len(n.edges) == 0 {
 279                 parent.delEdge(label)
 280         }
 281
 282         // Check if we should merge this node
 283         if n != t.root && len(n.edges) == 1 {
 284                 n.mergeChild()
 285         }
 286
 287         // Check if we should merge the parent's other child
 288         if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() {
 289                 parent.mergeChild()
 290         }
 291
 292         return leaf.val, true
 293 }
 294
 295 func (n *node) mergeChild() {
 296         e := n.edges[0]
 297         child := e.node
 298         n.prefix = n.prefix + child.prefix
 299         n.leaf = child.leaf
 300         n.edges = child.edges
 301 }
 302
 303 // Get is used to lookup a specific key, returning
 304 // the value and if it was found
 305 func (t *Tree) Get(s string) (interface{}, bool) {
 306         n := t.root
 307         search := s
 308         for {
 309                 // Check for key exhaution
 310                 if len(search) == 0 {
 311                         if n.isLeaf() {
 312                                 return n.leaf.val, true
 313                         }
 314                         break
 315                 }
 316
 317                 // Look for an edge
 318                 n = n.getEdge(search[0])
 319                 if n == nil {
 320                         break
 321                 }
 322
 323                 // Consume the search prefix
 324                 if strings.HasPrefix(search, n.prefix) {
 325                         search = search[len(n.prefix):]
 326                 } else {
 327                         break
 328                 }
 329         }
 330         return nil, false
 331 }
 332
 333 // LongestPrefix is like Get, but instead of an
 334 // exact match, it will return the longest prefix match.
 335 func (t *Tree) LongestPrefix(s string) (string, interface{}, bool) {
 336         var last *leafNode
 337         n := t.root
 338         search := s
 339         for {
 340                 // Look for a leaf node
 341                 if n.isLeaf() {
 342                         last = n.leaf
 343                 }
 344
 345                 // Check for key exhaution
 346                 if len(search) == 0 {
 347                         break
 348                 }
 349
 350                 // Look for an edge
 351                 n = n.getEdge(search[0])
 352                 if n == nil {
 353                         break
 354                 }
 355
 356                 // Consume the search prefix
 357                 if strings.HasPrefix(search, n.prefix) {
 358                         search = search[len(n.prefix):]
 359                 } else {
 360                         break
 361                 }
 362         }
 363         if last != nil {
 364                 return last.key, last.val, true
 365         }
 366         return "", nil, false
 367 }
 368
 369 // Minimum is used to return the minimum value in the tree
 370 func (t *Tree) Minimum() (string, interface{}, bool) {
 371         n := t.root
 372         for {
 373                 if n.isLeaf() {
 374                         return n.leaf.key, n.leaf.val, true
 375                 }
 376                 if len(n.edges) > 0 {
 377                         n = n.edges[0].node
 378                 } else {
 379                         break
 380                 }
 381         }
 382         return "", nil, false
 383 }
 384
 385 // Maximum is used to return the maximum value in the tree
 386 func (t *Tree) Maximum() (string, interface{}, bool) {
 387         n := t.root
 388         for {
 389                 if num := len(n.edges); num > 0 {
 390                         n = n.edges[num-1].node
 391                         continue
 392                 }
 393                 if n.isLeaf() {
 394                         return n.leaf.key, n.leaf.val, true
 395                 }
 396                 break
 397         }
 398         return "", nil, false
 399 }
 400
 401 // Walk is used to walk the tree
 402 func (t *Tree) Walk(fn WalkFn) {
 403         recursiveWalk(t.root, fn)
 404 }
 405
 406 // WalkPrefix is used to walk the tree under a prefix
 407 func (t *Tree) WalkPrefix(prefix string, fn WalkFn) {
 408         n := t.root
 409         search := prefix
 410         for {
 411                 // Check for key exhaution
 412                 if len(search) == 0 {
 413                         recursiveWalk(n, fn)
 414                         return
 415                 }
 416
 417                 // Look for an edge
 418                 n = n.getEdge(search[0])
 419                 if n == nil {
 420                         break
 421                 }
 422
 423                 // Consume the search prefix
 424                 if strings.HasPrefix(search, n.prefix) {
 425                         search = search[len(n.prefix):]
 426
 427                 } else if strings.HasPrefix(n.prefix, search) {
 428                         // Child may be under our search prefix
 429                         recursiveWalk(n, fn)
 430                         return
 431                 } else {
 432                         break
 433                 }
 434         }
 435
 436 }
 437
 438 // WalkPath is used to walk the tree, but only visiting nodes
 439 // from the root down to a given leaf. Where WalkPrefix walks
 440 // all the entries *under* the given prefix, this walks the
 441 // entries *above* the given prefix.
 442 func (t *Tree) WalkPath(path string, fn WalkFn) {
 443         n := t.root
 444         search := path
 445         for {
 446                 // Visit the leaf values if any
 447                 if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
 448                         return
 449                 }
 450
 451                 // Check for key exhaution
 452                 if len(search) == 0 {
 453                         return
 454                 }
 455
 456                 // Look for an edge
 457                 n = n.getEdge(search[0])
 458                 if n == nil {
 459                         return
 460                 }
 461
 462                 // Consume the search prefix
 463                 if strings.HasPrefix(search, n.prefix) {
 464                         search = search[len(n.prefix):]
 465                 } else {
 466                         break
 467                 }
 468         }
 469 }
 470
 471 // recursiveWalk is used to do a pre-order walk of a node
 472 // recursively. Returns true if the walk should be aborted
 473 func recursiveWalk(n *node, fn WalkFn) bool {
 474         // Visit the leaf values if any
 475         if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
 476                 return true
 477         }
 478
 479         // Recurse on the children
 480         for _, e := range n.edges {
 481                 if recursiveWalk(e.node, fn) {
 482                         return true
 483                 }
 484         }
 485         return false
 486 }
 487
 488 // ToMap is used to walk the tree and convert it into a map
 489 func (t *Tree) ToMap() map[string]interface{} {
 490         out := make(map[string]interface{}, t.size)
 491         t.Walk(func(k string, v interface{}) bool {
 492                 out[k] = v
 493                 return false
 494         })
 495         return out
 496 }