595 lines
14 KiB
Go
595 lines
14 KiB
Go
// Copyright (c) 2014 The go-patricia AUTHORS
|
|
//
|
|
// Use of this source code is governed by The MIT License
|
|
// that can be found in the LICENSE file.
|
|
|
|
package patricia
|
|
|
|
import (
|
|
"bytes"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"strings"
|
|
)
|
|
|
|
//------------------------------------------------------------------------------
|
|
// Trie
|
|
//------------------------------------------------------------------------------
|
|
|
|
const (
|
|
DefaultMaxPrefixPerNode = 10
|
|
DefaultMaxChildrenPerSparseNode = 8
|
|
)
|
|
|
|
type (
|
|
Prefix []byte
|
|
Item interface{}
|
|
VisitorFunc func(prefix Prefix, item Item) error
|
|
)
|
|
|
|
// Trie is a generic patricia trie that allows fast retrieval of items by prefix.
|
|
// and other funky stuff.
|
|
//
|
|
// Trie is not thread-safe.
|
|
type Trie struct {
|
|
prefix Prefix
|
|
item Item
|
|
|
|
maxPrefixPerNode int
|
|
maxChildrenPerSparseNode int
|
|
|
|
children childList
|
|
}
|
|
|
|
// Public API ------------------------------------------------------------------
|
|
|
|
type Option func(*Trie)
|
|
|
|
// Trie constructor.
|
|
func NewTrie(options ...Option) *Trie {
|
|
trie := &Trie{}
|
|
|
|
for _, opt := range options {
|
|
opt(trie)
|
|
}
|
|
|
|
if trie.maxPrefixPerNode <= 0 {
|
|
trie.maxPrefixPerNode = DefaultMaxPrefixPerNode
|
|
}
|
|
if trie.maxChildrenPerSparseNode <= 0 {
|
|
trie.maxChildrenPerSparseNode = DefaultMaxChildrenPerSparseNode
|
|
}
|
|
|
|
trie.children = newSparseChildList(trie.maxChildrenPerSparseNode)
|
|
return trie
|
|
}
|
|
|
|
func MaxPrefixPerNode(value int) Option {
|
|
return func(trie *Trie) {
|
|
trie.maxPrefixPerNode = value
|
|
}
|
|
}
|
|
|
|
func MaxChildrenPerSparseNode(value int) Option {
|
|
return func(trie *Trie) {
|
|
trie.maxChildrenPerSparseNode = value
|
|
}
|
|
}
|
|
|
|
// Item returns the item stored in the root of this trie.
|
|
func (trie *Trie) Item() Item {
|
|
return trie.item
|
|
}
|
|
|
|
// Insert inserts a new item into the trie using the given prefix. Insert does
|
|
// not replace existing items. It returns false if an item was already in place.
|
|
func (trie *Trie) Insert(key Prefix, item Item) (inserted bool) {
|
|
return trie.put(key, item, false)
|
|
}
|
|
|
|
// Set works much like Insert, but it always sets the item, possibly replacing
|
|
// the item previously inserted.
|
|
func (trie *Trie) Set(key Prefix, item Item) {
|
|
trie.put(key, item, true)
|
|
}
|
|
|
|
// Get returns the item located at key.
|
|
//
|
|
// This method is a bit dangerous, because Get can as well end up in an internal
|
|
// node that is not really representing any user-defined value. So when nil is
|
|
// a valid value being used, it is not possible to tell if the value was inserted
|
|
// into the tree by the user or not. A possible workaround for this is not to use
|
|
// nil interface as a valid value, even using zero value of any type is enough
|
|
// to prevent this bad behaviour.
|
|
func (trie *Trie) Get(key Prefix) (item Item) {
|
|
_, node, found, leftover := trie.findSubtree(key)
|
|
if !found || len(leftover) != 0 {
|
|
return nil
|
|
}
|
|
return node.item
|
|
}
|
|
|
|
// Match returns what Get(prefix) != nil would return. The same warning as for
|
|
// Get applies here as well.
|
|
func (trie *Trie) Match(prefix Prefix) (matchedExactly bool) {
|
|
return trie.Get(prefix) != nil
|
|
}
|
|
|
|
// MatchSubtree returns true when there is a subtree representing extensions
|
|
// to key, that is if there are any keys in the tree which have key as prefix.
|
|
func (trie *Trie) MatchSubtree(key Prefix) (matched bool) {
|
|
_, _, matched, _ = trie.findSubtree(key)
|
|
return
|
|
}
|
|
|
|
// Visit calls visitor on every node containing a non-nil item
|
|
// in alphabetical order.
|
|
//
|
|
// If an error is returned from visitor, the function stops visiting the tree
|
|
// and returns that error, unless it is a special error - SkipSubtree. In that
|
|
// case Visit skips the subtree represented by the current node and continues
|
|
// elsewhere.
|
|
func (trie *Trie) Visit(visitor VisitorFunc) error {
|
|
return trie.walk(nil, visitor)
|
|
}
|
|
|
|
func (trie *Trie) size() int {
|
|
n := 0
|
|
|
|
trie.walk(nil, func(prefix Prefix, item Item) error {
|
|
n++
|
|
return nil
|
|
})
|
|
|
|
return n
|
|
}
|
|
|
|
func (trie *Trie) total() int {
|
|
return 1 + trie.children.total()
|
|
}
|
|
|
|
// VisitSubtree works much like Visit, but it only visits nodes matching prefix.
|
|
func (trie *Trie) VisitSubtree(prefix Prefix, visitor VisitorFunc) error {
|
|
// Nil prefix not allowed.
|
|
if prefix == nil {
|
|
panic(ErrNilPrefix)
|
|
}
|
|
|
|
// Empty trie must be handled explicitly.
|
|
if trie.prefix == nil {
|
|
return nil
|
|
}
|
|
|
|
// Locate the relevant subtree.
|
|
_, root, found, leftover := trie.findSubtree(prefix)
|
|
if !found {
|
|
return nil
|
|
}
|
|
prefix = append(prefix, leftover...)
|
|
|
|
// Visit it.
|
|
return root.walk(prefix, visitor)
|
|
}
|
|
|
|
// VisitPrefixes visits only nodes that represent prefixes of key.
|
|
// To say the obvious, returning SkipSubtree from visitor makes no sense here.
|
|
func (trie *Trie) VisitPrefixes(key Prefix, visitor VisitorFunc) error {
|
|
// Nil key not allowed.
|
|
if key == nil {
|
|
panic(ErrNilPrefix)
|
|
}
|
|
|
|
// Empty trie must be handled explicitly.
|
|
if trie.prefix == nil {
|
|
return nil
|
|
}
|
|
|
|
// Walk the path matching key prefixes.
|
|
node := trie
|
|
prefix := key
|
|
offset := 0
|
|
for {
|
|
// Compute what part of prefix matches.
|
|
common := node.longestCommonPrefixLength(key)
|
|
key = key[common:]
|
|
offset += common
|
|
|
|
// Partial match means that there is no subtree matching prefix.
|
|
if common < len(node.prefix) {
|
|
return nil
|
|
}
|
|
|
|
// Call the visitor.
|
|
if item := node.item; item != nil {
|
|
if err := visitor(prefix[:offset], item); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
if len(key) == 0 {
|
|
// This node represents key, we are finished.
|
|
return nil
|
|
}
|
|
|
|
// There is some key suffix left, move to the children.
|
|
child := node.children.next(key[0])
|
|
if child == nil {
|
|
// There is nowhere to continue, return.
|
|
return nil
|
|
}
|
|
|
|
node = child
|
|
}
|
|
}
|
|
|
|
// Delete deletes the item represented by the given prefix.
|
|
//
|
|
// True is returned if the matching node was found and deleted.
|
|
func (trie *Trie) Delete(key Prefix) (deleted bool) {
|
|
// Nil prefix not allowed.
|
|
if key == nil {
|
|
panic(ErrNilPrefix)
|
|
}
|
|
|
|
// Empty trie must be handled explicitly.
|
|
if trie.prefix == nil {
|
|
return false
|
|
}
|
|
|
|
// Find the relevant node.
|
|
path, found, _ := trie.findSubtreePath(key)
|
|
if !found {
|
|
return false
|
|
}
|
|
|
|
node := path[len(path)-1]
|
|
var parent *Trie
|
|
if len(path) != 1 {
|
|
parent = path[len(path)-2]
|
|
}
|
|
|
|
// If the item is already set to nil, there is nothing to do.
|
|
if node.item == nil {
|
|
return false
|
|
}
|
|
|
|
// Delete the item.
|
|
node.item = nil
|
|
|
|
// Initialise i before goto.
|
|
// Will be used later in a loop.
|
|
i := len(path) - 1
|
|
|
|
// In case there are some child nodes, we cannot drop the whole subtree.
|
|
// We can try to compact nodes, though.
|
|
if node.children.length() != 0 {
|
|
goto Compact
|
|
}
|
|
|
|
// In case we are at the root, just reset it and we are done.
|
|
if parent == nil {
|
|
node.reset()
|
|
return true
|
|
}
|
|
|
|
// We can drop a subtree.
|
|
// Find the first ancestor that has its value set or it has 2 or more child nodes.
|
|
// That will be the node where to drop the subtree at.
|
|
for ; i >= 0; i-- {
|
|
if current := path[i]; current.item != nil || current.children.length() >= 2 {
|
|
break
|
|
}
|
|
}
|
|
|
|
// Handle the case when there is no such node.
|
|
// In other words, we can reset the whole tree.
|
|
if i == -1 {
|
|
path[0].reset()
|
|
return true
|
|
}
|
|
|
|
// We can just remove the subtree here.
|
|
node = path[i]
|
|
if i == 0 {
|
|
parent = nil
|
|
} else {
|
|
parent = path[i-1]
|
|
}
|
|
// i+1 is always a valid index since i is never pointing to the last node.
|
|
// The loop above skips at least the last node since we are sure that the item
|
|
// is set to nil and it has no children, othewise we would be compacting instead.
|
|
node.children.remove(path[i+1].prefix[0])
|
|
|
|
Compact:
|
|
// The node is set to the first non-empty ancestor,
|
|
// so try to compact since that might be possible now.
|
|
if compacted := node.compact(); compacted != node {
|
|
if parent == nil {
|
|
*node = *compacted
|
|
} else {
|
|
parent.children.replace(node.prefix[0], compacted)
|
|
*parent = *parent.compact()
|
|
}
|
|
}
|
|
|
|
return true
|
|
}
|
|
|
|
// DeleteSubtree finds the subtree exactly matching prefix and deletes it.
|
|
//
|
|
// True is returned if the subtree was found and deleted.
|
|
func (trie *Trie) DeleteSubtree(prefix Prefix) (deleted bool) {
|
|
// Nil prefix not allowed.
|
|
if prefix == nil {
|
|
panic(ErrNilPrefix)
|
|
}
|
|
|
|
// Empty trie must be handled explicitly.
|
|
if trie.prefix == nil {
|
|
return false
|
|
}
|
|
|
|
// Locate the relevant subtree.
|
|
parent, root, found, _ := trie.findSubtree(prefix)
|
|
if !found {
|
|
return false
|
|
}
|
|
|
|
// If we are in the root of the trie, reset the trie.
|
|
if parent == nil {
|
|
root.reset()
|
|
return true
|
|
}
|
|
|
|
// Otherwise remove the root node from its parent.
|
|
parent.children.remove(root.prefix[0])
|
|
return true
|
|
}
|
|
|
|
// Internal helper methods -----------------------------------------------------
|
|
|
|
func (trie *Trie) empty() bool {
|
|
return trie.item == nil && trie.children.length() == 0
|
|
}
|
|
|
|
func (trie *Trie) reset() {
|
|
trie.prefix = nil
|
|
trie.children = newSparseChildList(trie.maxPrefixPerNode)
|
|
}
|
|
|
|
func (trie *Trie) put(key Prefix, item Item, replace bool) (inserted bool) {
|
|
// Nil prefix not allowed.
|
|
if key == nil {
|
|
panic(ErrNilPrefix)
|
|
}
|
|
|
|
var (
|
|
common int
|
|
node *Trie = trie
|
|
child *Trie
|
|
)
|
|
|
|
if node.prefix == nil {
|
|
if len(key) <= trie.maxPrefixPerNode {
|
|
node.prefix = key
|
|
goto InsertItem
|
|
}
|
|
node.prefix = key[:trie.maxPrefixPerNode]
|
|
key = key[trie.maxPrefixPerNode:]
|
|
goto AppendChild
|
|
}
|
|
|
|
for {
|
|
// Compute the longest common prefix length.
|
|
common = node.longestCommonPrefixLength(key)
|
|
key = key[common:]
|
|
|
|
// Only a part matches, split.
|
|
if common < len(node.prefix) {
|
|
goto SplitPrefix
|
|
}
|
|
|
|
// common == len(node.prefix) since never (common > len(node.prefix))
|
|
// common == len(former key) <-> 0 == len(key)
|
|
// -> former key == node.prefix
|
|
if len(key) == 0 {
|
|
goto InsertItem
|
|
}
|
|
|
|
// Check children for matching prefix.
|
|
child = node.children.next(key[0])
|
|
if child == nil {
|
|
goto AppendChild
|
|
}
|
|
node = child
|
|
}
|
|
|
|
SplitPrefix:
|
|
// Split the prefix if necessary.
|
|
child = new(Trie)
|
|
*child = *node
|
|
*node = *NewTrie()
|
|
node.prefix = child.prefix[:common]
|
|
child.prefix = child.prefix[common:]
|
|
child = child.compact()
|
|
node.children = node.children.add(child)
|
|
|
|
AppendChild:
|
|
// Keep appending children until whole prefix is inserted.
|
|
// This loop starts with empty node.prefix that needs to be filled.
|
|
for len(key) != 0 {
|
|
child := NewTrie()
|
|
if len(key) <= trie.maxPrefixPerNode {
|
|
child.prefix = key
|
|
node.children = node.children.add(child)
|
|
node = child
|
|
goto InsertItem
|
|
} else {
|
|
child.prefix = key[:trie.maxPrefixPerNode]
|
|
key = key[trie.maxPrefixPerNode:]
|
|
node.children = node.children.add(child)
|
|
node = child
|
|
}
|
|
}
|
|
|
|
InsertItem:
|
|
// Try to insert the item if possible.
|
|
if replace || node.item == nil {
|
|
node.item = item
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (trie *Trie) compact() *Trie {
|
|
// Only a node with a single child can be compacted.
|
|
if trie.children.length() != 1 {
|
|
return trie
|
|
}
|
|
|
|
child := trie.children.head()
|
|
|
|
// If any item is set, we cannot compact since we want to retain
|
|
// the ability to do searching by key. This makes compaction less usable,
|
|
// but that simply cannot be avoided.
|
|
if trie.item != nil || child.item != nil {
|
|
return trie
|
|
}
|
|
|
|
// Make sure the combined prefixes fit into a single node.
|
|
if len(trie.prefix)+len(child.prefix) > trie.maxPrefixPerNode {
|
|
return trie
|
|
}
|
|
|
|
// Concatenate the prefixes, move the items.
|
|
child.prefix = append(trie.prefix, child.prefix...)
|
|
if trie.item != nil {
|
|
child.item = trie.item
|
|
}
|
|
|
|
return child
|
|
}
|
|
|
|
func (trie *Trie) findSubtree(prefix Prefix) (parent *Trie, root *Trie, found bool, leftover Prefix) {
|
|
// Find the subtree matching prefix.
|
|
root = trie
|
|
for {
|
|
// Compute what part of prefix matches.
|
|
common := root.longestCommonPrefixLength(prefix)
|
|
prefix = prefix[common:]
|
|
|
|
// We used up the whole prefix, subtree found.
|
|
if len(prefix) == 0 {
|
|
found = true
|
|
leftover = root.prefix[common:]
|
|
return
|
|
}
|
|
|
|
// Partial match means that there is no subtree matching prefix.
|
|
if common < len(root.prefix) {
|
|
leftover = root.prefix[common:]
|
|
return
|
|
}
|
|
|
|
// There is some prefix left, move to the children.
|
|
child := root.children.next(prefix[0])
|
|
if child == nil {
|
|
// There is nowhere to continue, there is no subtree matching prefix.
|
|
return
|
|
}
|
|
|
|
parent = root
|
|
root = child
|
|
}
|
|
}
|
|
|
|
func (trie *Trie) findSubtreePath(prefix Prefix) (path []*Trie, found bool, leftover Prefix) {
|
|
// Find the subtree matching prefix.
|
|
root := trie
|
|
var subtreePath []*Trie
|
|
for {
|
|
// Append the current root to the path.
|
|
subtreePath = append(subtreePath, root)
|
|
|
|
// Compute what part of prefix matches.
|
|
common := root.longestCommonPrefixLength(prefix)
|
|
prefix = prefix[common:]
|
|
|
|
// We used up the whole prefix, subtree found.
|
|
if len(prefix) == 0 {
|
|
path = subtreePath
|
|
found = true
|
|
leftover = root.prefix[common:]
|
|
return
|
|
}
|
|
|
|
// Partial match means that there is no subtree matching prefix.
|
|
if common < len(root.prefix) {
|
|
leftover = root.prefix[common:]
|
|
return
|
|
}
|
|
|
|
// There is some prefix left, move to the children.
|
|
child := root.children.next(prefix[0])
|
|
if child == nil {
|
|
// There is nowhere to continue, there is no subtree matching prefix.
|
|
return
|
|
}
|
|
|
|
root = child
|
|
}
|
|
}
|
|
|
|
func (trie *Trie) walk(actualRootPrefix Prefix, visitor VisitorFunc) error {
|
|
var prefix Prefix
|
|
// Allocate a bit more space for prefix at the beginning.
|
|
if actualRootPrefix == nil {
|
|
prefix = make(Prefix, 32+len(trie.prefix))
|
|
copy(prefix, trie.prefix)
|
|
prefix = prefix[:len(trie.prefix)]
|
|
} else {
|
|
prefix = make(Prefix, 32+len(actualRootPrefix))
|
|
copy(prefix, actualRootPrefix)
|
|
prefix = prefix[:len(actualRootPrefix)]
|
|
}
|
|
|
|
// Visit the root first. Not that this works for empty trie as well since
|
|
// in that case item == nil && len(children) == 0.
|
|
if trie.item != nil {
|
|
if err := visitor(prefix, trie.item); err != nil {
|
|
if err == SkipSubtree {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
}
|
|
|
|
// Then continue to the children.
|
|
return trie.children.walk(&prefix, visitor)
|
|
}
|
|
|
|
func (trie *Trie) longestCommonPrefixLength(prefix Prefix) (i int) {
|
|
for ; i < len(prefix) && i < len(trie.prefix) && prefix[i] == trie.prefix[i]; i++ {
|
|
}
|
|
return
|
|
}
|
|
|
|
func (trie *Trie) dump() string {
|
|
writer := &bytes.Buffer{}
|
|
trie.print(writer, 0)
|
|
return writer.String()
|
|
}
|
|
|
|
func (trie *Trie) print(writer io.Writer, indent int) {
|
|
fmt.Fprintf(writer, "%s%s %v\n", strings.Repeat(" ", indent), string(trie.prefix), trie.item)
|
|
trie.children.print(writer, indent+2)
|
|
}
|
|
|
|
// Errors ----------------------------------------------------------------------
|
|
|
|
var (
|
|
SkipSubtree = errors.New("Skip this subtree")
|
|
ErrNilPrefix = errors.New("Nil prefix passed into a method call")
|
|
)
|