1// Copyright (c) 2014 The go-patricia AUTHORS 2// 3// Use of this source code is governed by The MIT License 4// that can be found in the LICENSE file. 5 6package patricia 7 8import ( 9 "bytes" 10 "errors" 11 "fmt" 12 "io" 13 "strings" 14) 15 16//------------------------------------------------------------------------------ 17// Trie 18//------------------------------------------------------------------------------ 19 20const ( 21 DefaultMaxPrefixPerNode = 10 22 DefaultMaxChildrenPerSparseNode = 8 23) 24 25type ( 26 Prefix []byte 27 Item interface{} 28 VisitorFunc func(prefix Prefix, item Item) error 29) 30 31// Trie is a generic patricia trie that allows fast retrieval of items by prefix. 32// and other funky stuff. 33// 34// Trie is not thread-safe. 35type Trie struct { 36 prefix Prefix 37 item Item 38 39 maxPrefixPerNode int 40 maxChildrenPerSparseNode int 41 42 children childList 43} 44 45// Public API ------------------------------------------------------------------ 46 47type Option func(*Trie) 48 49// Trie constructor. 50func NewTrie(options ...Option) *Trie { 51 trie := &Trie{} 52 53 for _, opt := range options { 54 opt(trie) 55 } 56 57 if trie.maxPrefixPerNode <= 0 { 58 trie.maxPrefixPerNode = DefaultMaxPrefixPerNode 59 } 60 if trie.maxChildrenPerSparseNode <= 0 { 61 trie.maxChildrenPerSparseNode = DefaultMaxChildrenPerSparseNode 62 } 63 64 trie.children = newSparseChildList(trie.maxChildrenPerSparseNode) 65 return trie 66} 67 68func MaxPrefixPerNode(value int) Option { 69 return func(trie *Trie) { 70 trie.maxPrefixPerNode = value 71 } 72} 73 74func MaxChildrenPerSparseNode(value int) Option { 75 return func(trie *Trie) { 76 trie.maxChildrenPerSparseNode = value 77 } 78} 79 80// Item returns the item stored in the root of this trie. 81func (trie *Trie) Item() Item { 82 return trie.item 83} 84 85// Insert inserts a new item into the trie using the given prefix. Insert does 86// not replace existing items. It returns false if an item was already in place. 87func (trie *Trie) Insert(key Prefix, item Item) (inserted bool) { 88 return trie.put(key, item, false) 89} 90 91// Set works much like Insert, but it always sets the item, possibly replacing 92// the item previously inserted. 93func (trie *Trie) Set(key Prefix, item Item) { 94 trie.put(key, item, true) 95} 96 97// Get returns the item located at key. 98// 99// This method is a bit dangerous, because Get can as well end up in an internal 100// node that is not really representing any user-defined value. So when nil is 101// a valid value being used, it is not possible to tell if the value was inserted 102// into the tree by the user or not. A possible workaround for this is not to use 103// nil interface as a valid value, even using zero value of any type is enough 104// to prevent this bad behaviour. 105func (trie *Trie) Get(key Prefix) (item Item) { 106 _, node, found, leftover := trie.findSubtree(key) 107 if !found || len(leftover) != 0 { 108 return nil 109 } 110 return node.item 111} 112 113// Match returns what Get(prefix) != nil would return. The same warning as for 114// Get applies here as well. 115func (trie *Trie) Match(prefix Prefix) (matchedExactly bool) { 116 return trie.Get(prefix) != nil 117} 118 119// MatchSubtree returns true when there is a subtree representing extensions 120// to key, that is if there are any keys in the tree which have key as prefix. 121func (trie *Trie) MatchSubtree(key Prefix) (matched bool) { 122 _, _, matched, _ = trie.findSubtree(key) 123 return 124} 125 126// Visit calls visitor on every node containing a non-nil item 127// in alphabetical order. 128// 129// If an error is returned from visitor, the function stops visiting the tree 130// and returns that error, unless it is a special error - SkipSubtree. In that 131// case Visit skips the subtree represented by the current node and continues 132// elsewhere. 133func (trie *Trie) Visit(visitor VisitorFunc) error { 134 return trie.walk(nil, visitor) 135} 136 137func (trie *Trie) size() int { 138 n := 0 139 140 trie.walk(nil, func(prefix Prefix, item Item) error { 141 n++ 142 return nil 143 }) 144 145 return n 146} 147 148func (trie *Trie) total() int { 149 return 1 + trie.children.total() 150} 151 152// VisitSubtree works much like Visit, but it only visits nodes matching prefix. 153func (trie *Trie) VisitSubtree(prefix Prefix, visitor VisitorFunc) error { 154 // Nil prefix not allowed. 155 if prefix == nil { 156 panic(ErrNilPrefix) 157 } 158 159 // Empty trie must be handled explicitly. 160 if trie.prefix == nil { 161 return nil 162 } 163 164 // Locate the relevant subtree. 165 _, root, found, leftover := trie.findSubtree(prefix) 166 if !found { 167 return nil 168 } 169 prefix = append(prefix, leftover...) 170 171 // Visit it. 172 return root.walk(prefix, visitor) 173} 174 175// VisitPrefixes visits only nodes that represent prefixes of key. 176// To say the obvious, returning SkipSubtree from visitor makes no sense here. 177func (trie *Trie) VisitPrefixes(key Prefix, visitor VisitorFunc) error { 178 // Nil key not allowed. 179 if key == nil { 180 panic(ErrNilPrefix) 181 } 182 183 // Empty trie must be handled explicitly. 184 if trie.prefix == nil { 185 return nil 186 } 187 188 // Walk the path matching key prefixes. 189 node := trie 190 prefix := key 191 offset := 0 192 for { 193 // Compute what part of prefix matches. 194 common := node.longestCommonPrefixLength(key) 195 key = key[common:] 196 offset += common 197 198 // Partial match means that there is no subtree matching prefix. 199 if common < len(node.prefix) { 200 return nil 201 } 202 203 // Call the visitor. 204 if item := node.item; item != nil { 205 if err := visitor(prefix[:offset], item); err != nil { 206 return err 207 } 208 } 209 210 if len(key) == 0 { 211 // This node represents key, we are finished. 212 return nil 213 } 214 215 // There is some key suffix left, move to the children. 216 child := node.children.next(key[0]) 217 if child == nil { 218 // There is nowhere to continue, return. 219 return nil 220 } 221 222 node = child 223 } 224} 225 226// Delete deletes the item represented by the given prefix. 227// 228// True is returned if the matching node was found and deleted. 229func (trie *Trie) Delete(key Prefix) (deleted bool) { 230 // Nil prefix not allowed. 231 if key == nil { 232 panic(ErrNilPrefix) 233 } 234 235 // Empty trie must be handled explicitly. 236 if trie.prefix == nil { 237 return false 238 } 239 240 // Find the relevant node. 241 path, found, _ := trie.findSubtreePath(key) 242 if !found { 243 return false 244 } 245 246 node := path[len(path)-1] 247 var parent *Trie 248 if len(path) != 1 { 249 parent = path[len(path)-2] 250 } 251 252 // If the item is already set to nil, there is nothing to do. 253 if node.item == nil { 254 return false 255 } 256 257 // Delete the item. 258 node.item = nil 259 260 // Initialise i before goto. 261 // Will be used later in a loop. 262 i := len(path) - 1 263 264 // In case there are some child nodes, we cannot drop the whole subtree. 265 // We can try to compact nodes, though. 266 if node.children.length() != 0 { 267 goto Compact 268 } 269 270 // In case we are at the root, just reset it and we are done. 271 if parent == nil { 272 node.reset() 273 return true 274 } 275 276 // We can drop a subtree. 277 // Find the first ancestor that has its value set or it has 2 or more child nodes. 278 // That will be the node where to drop the subtree at. 279 for ; i >= 0; i-- { 280 if current := path[i]; current.item != nil || current.children.length() >= 2 { 281 break 282 } 283 } 284 285 // Handle the case when there is no such node. 286 // In other words, we can reset the whole tree. 287 if i == -1 { 288 path[0].reset() 289 return true 290 } 291 292 // We can just remove the subtree here. 293 node = path[i] 294 if i == 0 { 295 parent = nil 296 } else { 297 parent = path[i-1] 298 } 299 // i+1 is always a valid index since i is never pointing to the last node. 300 // The loop above skips at least the last node since we are sure that the item 301 // is set to nil and it has no children, othewise we would be compacting instead. 302 node.children.remove(path[i+1].prefix[0]) 303 304Compact: 305 // The node is set to the first non-empty ancestor, 306 // so try to compact since that might be possible now. 307 if compacted := node.compact(); compacted != node { 308 if parent == nil { 309 *node = *compacted 310 } else { 311 parent.children.replace(node.prefix[0], compacted) 312 *parent = *parent.compact() 313 } 314 } 315 316 return true 317} 318 319// DeleteSubtree finds the subtree exactly matching prefix and deletes it. 320// 321// True is returned if the subtree was found and deleted. 322func (trie *Trie) DeleteSubtree(prefix Prefix) (deleted bool) { 323 // Nil prefix not allowed. 324 if prefix == nil { 325 panic(ErrNilPrefix) 326 } 327 328 // Empty trie must be handled explicitly. 329 if trie.prefix == nil { 330 return false 331 } 332 333 // Locate the relevant subtree. 334 parent, root, found, _ := trie.findSubtree(prefix) 335 if !found { 336 return false 337 } 338 339 // If we are in the root of the trie, reset the trie. 340 if parent == nil { 341 root.reset() 342 return true 343 } 344 345 // Otherwise remove the root node from its parent. 346 parent.children.remove(root.prefix[0]) 347 return true 348} 349 350// Internal helper methods ----------------------------------------------------- 351 352func (trie *Trie) empty() bool { 353 return trie.item == nil && trie.children.length() == 0 354} 355 356func (trie *Trie) reset() { 357 trie.prefix = nil 358 trie.children = newSparseChildList(trie.maxPrefixPerNode) 359} 360 361func (trie *Trie) put(key Prefix, item Item, replace bool) (inserted bool) { 362 // Nil prefix not allowed. 363 if key == nil { 364 panic(ErrNilPrefix) 365 } 366 367 var ( 368 common int 369 node *Trie = trie 370 child *Trie 371 ) 372 373 if node.prefix == nil { 374 if len(key) <= trie.maxPrefixPerNode { 375 node.prefix = key 376 goto InsertItem 377 } 378 node.prefix = key[:trie.maxPrefixPerNode] 379 key = key[trie.maxPrefixPerNode:] 380 goto AppendChild 381 } 382 383 for { 384 // Compute the longest common prefix length. 385 common = node.longestCommonPrefixLength(key) 386 key = key[common:] 387 388 // Only a part matches, split. 389 if common < len(node.prefix) { 390 goto SplitPrefix 391 } 392 393 // common == len(node.prefix) since never (common > len(node.prefix)) 394 // common == len(former key) <-> 0 == len(key) 395 // -> former key == node.prefix 396 if len(key) == 0 { 397 goto InsertItem 398 } 399 400 // Check children for matching prefix. 401 child = node.children.next(key[0]) 402 if child == nil { 403 goto AppendChild 404 } 405 node = child 406 } 407 408SplitPrefix: 409 // Split the prefix if necessary. 410 child = new(Trie) 411 *child = *node 412 *node = *NewTrie() 413 node.prefix = child.prefix[:common] 414 child.prefix = child.prefix[common:] 415 child = child.compact() 416 node.children = node.children.add(child) 417 418AppendChild: 419 // Keep appending children until whole prefix is inserted. 420 // This loop starts with empty node.prefix that needs to be filled. 421 for len(key) != 0 { 422 child := NewTrie() 423 if len(key) <= trie.maxPrefixPerNode { 424 child.prefix = key 425 node.children = node.children.add(child) 426 node = child 427 goto InsertItem 428 } else { 429 child.prefix = key[:trie.maxPrefixPerNode] 430 key = key[trie.maxPrefixPerNode:] 431 node.children = node.children.add(child) 432 node = child 433 } 434 } 435 436InsertItem: 437 // Try to insert the item if possible. 438 if replace || node.item == nil { 439 node.item = item 440 return true 441 } 442 return false 443} 444 445func (trie *Trie) compact() *Trie { 446 // Only a node with a single child can be compacted. 447 if trie.children.length() != 1 { 448 return trie 449 } 450 451 child := trie.children.head() 452 453 // If any item is set, we cannot compact since we want to retain 454 // the ability to do searching by key. This makes compaction less usable, 455 // but that simply cannot be avoided. 456 if trie.item != nil || child.item != nil { 457 return trie 458 } 459 460 // Make sure the combined prefixes fit into a single node. 461 if len(trie.prefix)+len(child.prefix) > trie.maxPrefixPerNode { 462 return trie 463 } 464 465 // Concatenate the prefixes, move the items. 466 child.prefix = append(trie.prefix, child.prefix...) 467 if trie.item != nil { 468 child.item = trie.item 469 } 470 471 return child 472} 473 474func (trie *Trie) findSubtree(prefix Prefix) (parent *Trie, root *Trie, found bool, leftover Prefix) { 475 // Find the subtree matching prefix. 476 root = trie 477 for { 478 // Compute what part of prefix matches. 479 common := root.longestCommonPrefixLength(prefix) 480 prefix = prefix[common:] 481 482 // We used up the whole prefix, subtree found. 483 if len(prefix) == 0 { 484 found = true 485 leftover = root.prefix[common:] 486 return 487 } 488 489 // Partial match means that there is no subtree matching prefix. 490 if common < len(root.prefix) { 491 leftover = root.prefix[common:] 492 return 493 } 494 495 // There is some prefix left, move to the children. 496 child := root.children.next(prefix[0]) 497 if child == nil { 498 // There is nowhere to continue, there is no subtree matching prefix. 499 return 500 } 501 502 parent = root 503 root = child 504 } 505} 506 507func (trie *Trie) findSubtreePath(prefix Prefix) (path []*Trie, found bool, leftover Prefix) { 508 // Find the subtree matching prefix. 509 root := trie 510 var subtreePath []*Trie 511 for { 512 // Append the current root to the path. 513 subtreePath = append(subtreePath, root) 514 515 // Compute what part of prefix matches. 516 common := root.longestCommonPrefixLength(prefix) 517 prefix = prefix[common:] 518 519 // We used up the whole prefix, subtree found. 520 if len(prefix) == 0 { 521 path = subtreePath 522 found = true 523 leftover = root.prefix[common:] 524 return 525 } 526 527 // Partial match means that there is no subtree matching prefix. 528 if common < len(root.prefix) { 529 leftover = root.prefix[common:] 530 return 531 } 532 533 // There is some prefix left, move to the children. 534 child := root.children.next(prefix[0]) 535 if child == nil { 536 // There is nowhere to continue, there is no subtree matching prefix. 537 return 538 } 539 540 root = child 541 } 542} 543 544func (trie *Trie) walk(actualRootPrefix Prefix, visitor VisitorFunc) error { 545 var prefix Prefix 546 // Allocate a bit more space for prefix at the beginning. 547 if actualRootPrefix == nil { 548 prefix = make(Prefix, 32+len(trie.prefix)) 549 copy(prefix, trie.prefix) 550 prefix = prefix[:len(trie.prefix)] 551 } else { 552 prefix = make(Prefix, 32+len(actualRootPrefix)) 553 copy(prefix, actualRootPrefix) 554 prefix = prefix[:len(actualRootPrefix)] 555 } 556 557 // Visit the root first. Not that this works for empty trie as well since 558 // in that case item == nil && len(children) == 0. 559 if trie.item != nil { 560 if err := visitor(prefix, trie.item); err != nil { 561 if err == SkipSubtree { 562 return nil 563 } 564 return err 565 } 566 } 567 568 // Then continue to the children. 569 return trie.children.walk(&prefix, visitor) 570} 571 572func (trie *Trie) longestCommonPrefixLength(prefix Prefix) (i int) { 573 for ; i < len(prefix) && i < len(trie.prefix) && prefix[i] == trie.prefix[i]; i++ { 574 } 575 return 576} 577 578func (trie *Trie) dump() string { 579 writer := &bytes.Buffer{} 580 trie.print(writer, 0) 581 return writer.String() 582} 583 584func (trie *Trie) print(writer io.Writer, indent int) { 585 fmt.Fprintf(writer, "%s%s %v\n", strings.Repeat(" ", indent), string(trie.prefix), trie.item) 586 trie.children.print(writer, indent+2) 587} 588 589// Errors ---------------------------------------------------------------------- 590 591var ( 592 SkipSubtree = errors.New("Skip this subtree") 593 ErrNilPrefix = errors.New("Nil prefix passed into a method call") 594) 595