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) updateEdge(label byte, node *node) {
  48         num := len(n.edges)
  49         idx := sort.Search(num, func(i int) bool {
  50                 return n.edges[i].label >= label
  51         })
  52         if idx < num && n.edges[idx].label == label {
  53                 n.edges[idx].node = 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.updateEdge(search[0], child)
 202
 203                 // Restore the existing node
 204                 child.addEdge(edge{
 205                         label: n.prefix[commonPrefix],
 206                         node:  n,
 207                 })
 208                 n.prefix = n.prefix[commonPrefix:]
 209
 210                 // Create a new leaf node
 211                 leaf := &leafNode{
 212                         key: s,
 213                         val: v,
 214                 }
 215
 216                 // If the new key is a subset, add to to this node
 217                 search = search[commonPrefix:]
 218                 if len(search) == 0 {
 219                         child.leaf = leaf
 220                         return nil, false
 221                 }
 222
 223                 // Create a new edge for the node
 224                 child.addEdge(edge{
 225                         label: search[0],
 226                         node: &node{
 227                                 leaf:   leaf,
 228                                 prefix: search,
 229                         },
 230                 })
 231                 return nil, false
 232         }
 233 }
 234
 235 // Delete is used to delete a key, returning the previous
 236 // value and if it was deleted
 237 func (t *Tree) Delete(s string) (interface{}, bool) {
 238         var parent *node
 239         var label byte
 240         n := t.root
 241         search := s
 242         for {
 243                 // Check for key exhaution
 244                 if len(search) == 0 {
 245                         if !n.isLeaf() {
 246                                 break
 247                         }
 248                         goto DELETE
 249                 }
 250
 251                 // Look for an edge
 252                 parent = n
 253                 label = search[0]
 254                 n = n.getEdge(label)
 255                 if n == nil {
 256                         break
 257                 }
 258
 259                 // Consume the search prefix
 260                 if strings.HasPrefix(search, n.prefix) {
 261                         search = search[len(n.prefix):]
 262                 } else {
 263                         break
 264                 }
 265         }
 266         return nil, false
 267
 268 DELETE:
 269         // Delete the leaf
 270         leaf := n.leaf
 271         n.leaf = nil
 272         t.size--
 273
 274         // Check if we should delete this node from the parent
 275         if parent != nil && len(n.edges) == 0 {
 276                 parent.delEdge(label)
 277         }
 278
 279         // Check if we should merge this node
 280         if n != t.root && len(n.edges) == 1 {
 281                 n.mergeChild()
 282         }
 283
 284         // Check if we should merge the parent's other child
 285         if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() {
 286                 parent.mergeChild()
 287         }
 288
 289         return leaf.val, true
 290 }
 291
 292 // DeletePrefix is used to delete the subtree under a prefix
 293 // Returns how many nodes were deleted
 294 // Use this to delete large subtrees efficiently
 295 func (t *Tree) DeletePrefix(s string) int {
 296         return t.deletePrefix(nil, t.root, s)
 297 }
 298
 299 // delete does a recursive deletion
 300 func (t *Tree) deletePrefix(parent, n *node, prefix string) int {
 301         // Check for key exhaustion
 302         if len(prefix) == 0 {
 303                 // Remove the leaf node
 304                 subTreeSize := 0
 305                 //recursively walk from all edges of the node to be deleted
 306                 recursiveWalk(n, func(s string, v interface{}) bool {
 307                         subTreeSize++
 308                         return false
 309                 })
 310                 if n.isLeaf() {
 311                         n.leaf = nil
 312                 }
 313                 n.edges = nil // deletes the entire subtree
 314
 315                 // Check if we should merge the parent's other child
 316                 if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() {
 317                         parent.mergeChild()
 318                 }
 319                 t.size -= subTreeSize
 320                 return subTreeSize
 321         }
 322
 323         // Look for an edge
 324         label := prefix[0]
 325         child := n.getEdge(label)
 326         if child == nil || (!strings.HasPrefix(child.prefix, prefix) && !strings.HasPrefix(prefix, child.prefix)) {
 327                 return 0
 328         }
 329
 330         // Consume the search prefix
 331         if len(child.prefix) > len(prefix) {
 332                 prefix = prefix[len(prefix):]
 333         } else {
 334                 prefix = prefix[len(child.prefix):]
 335         }
 336         return t.deletePrefix(n, child, prefix)
 337 }
 338
 339 func (n *node) mergeChild() {
 340         e := n.edges[0]
 341         child := e.node
 342         n.prefix = n.prefix + child.prefix
 343         n.leaf = child.leaf
 344         n.edges = child.edges
 345 }
 346
 347 // Get is used to lookup a specific key, returning
 348 // the value and if it was found
 349 func (t *Tree) Get(s string) (interface{}, bool) {
 350         n := t.root
 351         search := s
 352         for {
 353                 // Check for key exhaution
 354                 if len(search) == 0 {
 355                         if n.isLeaf() {
 356                                 return n.leaf.val, true
 357                         }
 358                         break
 359                 }
 360
 361                 // Look for an edge
 362                 n = n.getEdge(search[0])
 363                 if n == nil {
 364                         break
 365                 }
 366
 367                 // Consume the search prefix
 368                 if strings.HasPrefix(search, n.prefix) {
 369                         search = search[len(n.prefix):]
 370                 } else {
 371                         break
 372                 }
 373         }
 374         return nil, false
 375 }
 376
 377 // LongestPrefix is like Get, but instead of an
 378 // exact match, it will return the longest prefix match.
 379 func (t *Tree) LongestPrefix(s string) (string, interface{}, bool) {
 380         var last *leafNode
 381         n := t.root
 382         search := s
 383         for {
 384                 // Look for a leaf node
 385                 if n.isLeaf() {
 386                         last = n.leaf
 387                 }
 388
 389                 // Check for key exhaution
 390                 if len(search) == 0 {
 391                         break
 392                 }
 393
 394                 // Look for an edge
 395                 n = n.getEdge(search[0])
 396                 if n == nil {
 397                         break
 398                 }
 399
 400                 // Consume the search prefix
 401                 if strings.HasPrefix(search, n.prefix) {
 402                         search = search[len(n.prefix):]
 403                 } else {
 404                         break
 405                 }
 406         }
 407         if last != nil {
 408                 return last.key, last.val, true
 409         }
 410         return "", nil, false
 411 }
 412
 413 // Minimum is used to return the minimum value in the tree
 414 func (t *Tree) Minimum() (string, interface{}, bool) {
 415         n := t.root
 416         for {
 417                 if n.isLeaf() {
 418                         return n.leaf.key, n.leaf.val, true
 419                 }
 420                 if len(n.edges) > 0 {
 421                         n = n.edges[0].node
 422                 } else {
 423                         break
 424                 }
 425         }
 426         return "", nil, false
 427 }
 428
 429 // Maximum is used to return the maximum value in the tree
 430 func (t *Tree) Maximum() (string, interface{}, bool) {
 431         n := t.root
 432         for {
 433                 if num := len(n.edges); num > 0 {
 434                         n = n.edges[num-1].node
 435                         continue
 436                 }
 437                 if n.isLeaf() {
 438                         return n.leaf.key, n.leaf.val, true
 439                 }
 440                 break
 441         }
 442         return "", nil, false
 443 }
 444
 445 // Walk is used to walk the tree
 446 func (t *Tree) Walk(fn WalkFn) {
 447         recursiveWalk(t.root, fn)
 448 }
 449
 450 // WalkPrefix is used to walk the tree under a prefix
 451 func (t *Tree) WalkPrefix(prefix string, fn WalkFn) {
 452         n := t.root
 453         search := prefix
 454         for {
 455                 // Check for key exhaution
 456                 if len(search) == 0 {
 457                         recursiveWalk(n, fn)
 458                         return
 459                 }
 460
 461                 // Look for an edge
 462                 n = n.getEdge(search[0])
 463                 if n == nil {
 464                         break
 465                 }
 466
 467                 // Consume the search prefix
 468                 if strings.HasPrefix(search, n.prefix) {
 469                         search = search[len(n.prefix):]
 470
 471                 } else if strings.HasPrefix(n.prefix, search) {
 472                         // Child may be under our search prefix
 473                         recursiveWalk(n, fn)
 474                         return
 475                 } else {
 476                         break
 477                 }
 478         }
 479
 480 }
 481
 482 // WalkPath is used to walk the tree, but only visiting nodes
 483 // from the root down to a given leaf. Where WalkPrefix walks
 484 // all the entries *under* the given prefix, this walks the
 485 // entries *above* the given prefix.
 486 func (t *Tree) WalkPath(path string, fn WalkFn) {
 487         n := t.root
 488         search := path
 489         for {
 490                 // Visit the leaf values if any
 491                 if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
 492                         return
 493                 }
 494
 495                 // Check for key exhaution
 496                 if len(search) == 0 {
 497                         return
 498                 }
 499
 500                 // Look for an edge
 501                 n = n.getEdge(search[0])
 502                 if n == nil {
 503                         return
 504                 }
 505
 506                 // Consume the search prefix
 507                 if strings.HasPrefix(search, n.prefix) {
 508                         search = search[len(n.prefix):]
 509                 } else {
 510                         break
 511                 }
 512         }
 513 }
 514
 515 // recursiveWalk is used to do a pre-order walk of a node
 516 // recursively. Returns true if the walk should be aborted
 517 func recursiveWalk(n *node, fn WalkFn) bool {
 518         // Visit the leaf values if any
 519         if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
 520                 return true
 521         }
 522
 523         // Recurse on the children
 524         for _, e := range n.edges {
 525                 if recursiveWalk(e.node, fn) {
 526                         return true
 527                 }
 528         }
 529         return false
 530 }
 531
 532 // ToMap is used to walk the tree and convert it into a map
 533 func (t *Tree) ToMap() map[string]interface{} {
 534         out := make(map[string]interface{}, t.size)
 535         t.Walk(func(k string, v interface{}) bool {
 536                 out[k] = v
 537                 return false
 538         })
 539         return out
 540 }