forgejo/vendor/github.com/toqueteos/trie/trie.go

103 lines
2.4 KiB
Go
Raw Normal View History

// Package trie is an implementation of a trie (prefix tree) data structure over byte slices. It provides a
// small and simple API for usage as a set as well as a 'Node' API for walking the trie.
package trie
// A Trie is a a prefix tree.
type Trie struct {
root *Node
}
// New construct a new, empty Trie ready for use.
func New() *Trie {
return &Trie{
root: &Node{},
}
}
// Insert puts b into the Trie. It returns true if the element was not previously in t.
func (t *Trie) Insert(b []byte) bool {
n := t.root
for _, c := range b {
next, ok := n.Walk(c)
if !ok {
next = &Node{}
n.branches[c] = next
n.hasChildren = true
}
n = next
}
if n.terminal {
return false
}
n.terminal = true
return true
}
// Contains checks t for membership of b.
func (t *Trie) Contains(b []byte) bool {
n := t.root
for _, c := range b {
next, ok := n.Walk(c)
if !ok {
return false
}
n = next
}
return n.terminal
}
// PrefixIndex walks through `b` until a prefix is found (terminal node) or it is exhausted.
func (t *Trie) PrefixIndex(b []byte) int {
var idx int
n := t.root
for _, c := range b {
next, ok := n.Walk(c)
if !ok {
return -1
}
if next.terminal {
return idx
}
n = next
idx++
}
if !n.terminal {
idx = -1
}
return idx
}
// Root returns the root node of a Trie. A valid Trie (i.e., constructed with New), always has a non-nil root
// node.
func (t *Trie) Root() *Node {
return t.root
}
// A Node represents a logical vertex in the trie structure.
type Node struct {
branches [256]*Node
terminal bool
hasChildren bool
}
// Walk returns the node reached along edge c, if one exists. The ok value indicates whether such a node
// exist.
func (n *Node) Walk(c byte) (next *Node, ok bool) {
next = n.branches[int(c)]
return next, (next != nil)
}
// Terminal indicates whether n is terminal in the trie (that is, whether the path from the root to n
// represents an element in the set). For instance, if the root node is terminal, then []byte{} is in the
// trie.
func (n *Node) Terminal() bool {
return n.terminal
}
// Leaf indicates whether n is a leaf node in the trie (that is, whether it has children). A leaf node must be
// terminal (else it would not exist). Logically, if n is a leaf node then the []byte represented by the path
// from the root to n is not a proper prefix of any element of the trie.
func (n *Node) Leaf() bool {
return !n.hasChildren
}