353 lines
6.9 KiB
Go
353 lines
6.9 KiB
Go
package iradix
|
|
|
|
import (
|
|
"bytes"
|
|
"sort"
|
|
)
|
|
|
|
// WalkFn is used when walking the tree. Takes a
|
|
// key and value, returning if iteration should
|
|
// be terminated.
|
|
type WalkFn func(k []byte, v interface{}) bool
|
|
|
|
// leafNode is used to represent a value
|
|
type leafNode struct {
|
|
mutateCh chan struct{}
|
|
key []byte
|
|
val interface{}
|
|
}
|
|
|
|
// edge is used to represent an edge node
|
|
type edge struct {
|
|
label byte
|
|
node *Node
|
|
}
|
|
|
|
// Node is an immutable node in the radix tree
|
|
type Node struct {
|
|
// mutateCh is closed if this node is modified
|
|
mutateCh chan struct{}
|
|
|
|
// leaf is used to store possible leaf
|
|
leaf *leafNode
|
|
|
|
// prefix is the common prefix we ignore
|
|
prefix []byte
|
|
|
|
// Edges should be stored in-order for iteration.
|
|
// We avoid a fully materialized slice to save memory,
|
|
// since in most cases we expect to be sparse
|
|
edges edges
|
|
}
|
|
|
|
func (n *Node) isLeaf() bool {
|
|
return n.leaf != nil
|
|
}
|
|
|
|
func (n *Node) addEdge(e edge) {
|
|
num := len(n.edges)
|
|
idx := sort.Search(num, func(i int) bool {
|
|
return n.edges[i].label >= e.label
|
|
})
|
|
n.edges = append(n.edges, e)
|
|
if idx != num {
|
|
copy(n.edges[idx+1:], n.edges[idx:num])
|
|
n.edges[idx] = e
|
|
}
|
|
}
|
|
|
|
func (n *Node) replaceEdge(e edge) {
|
|
num := len(n.edges)
|
|
idx := sort.Search(num, func(i int) bool {
|
|
return n.edges[i].label >= e.label
|
|
})
|
|
if idx < num && n.edges[idx].label == e.label {
|
|
n.edges[idx].node = e.node
|
|
return
|
|
}
|
|
panic("replacing missing edge")
|
|
}
|
|
|
|
func (n *Node) getEdge(label byte) (int, *Node) {
|
|
num := len(n.edges)
|
|
idx := sort.Search(num, func(i int) bool {
|
|
return n.edges[i].label >= label
|
|
})
|
|
if idx < num && n.edges[idx].label == label {
|
|
return idx, n.edges[idx].node
|
|
}
|
|
return -1, nil
|
|
}
|
|
|
|
func (n *Node) delEdge(label byte) {
|
|
num := len(n.edges)
|
|
idx := sort.Search(num, func(i int) bool {
|
|
return n.edges[i].label >= label
|
|
})
|
|
if idx < num && n.edges[idx].label == label {
|
|
copy(n.edges[idx:], n.edges[idx+1:])
|
|
n.edges[len(n.edges)-1] = edge{}
|
|
n.edges = n.edges[:len(n.edges)-1]
|
|
}
|
|
}
|
|
|
|
func (n *Node) GetWatch(k []byte) (<-chan struct{}, interface{}, bool) {
|
|
search := k
|
|
watch := n.mutateCh
|
|
for {
|
|
// Check for key exhaustion
|
|
if len(search) == 0 {
|
|
if n.isLeaf() {
|
|
return n.leaf.mutateCh, n.leaf.val, true
|
|
}
|
|
break
|
|
}
|
|
|
|
// Look for an edge
|
|
_, n = n.getEdge(search[0])
|
|
if n == nil {
|
|
break
|
|
}
|
|
|
|
// Update to the finest granularity as the search makes progress
|
|
watch = n.mutateCh
|
|
|
|
// Consume the search prefix
|
|
if bytes.HasPrefix(search, n.prefix) {
|
|
search = search[len(n.prefix):]
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
return watch, nil, false
|
|
}
|
|
|
|
func (n *Node) Get(k []byte) (interface{}, bool) {
|
|
_, val, ok := n.GetWatch(k)
|
|
return val, ok
|
|
}
|
|
|
|
// LongestPrefix is like Get, but instead of an
|
|
// exact match, it will return the longest prefix match.
|
|
func (n *Node) LongestPrefix(k []byte) ([]byte, interface{}, bool) {
|
|
var last *leafNode
|
|
search := k
|
|
for {
|
|
// Look for a leaf node
|
|
if n.isLeaf() {
|
|
last = n.leaf
|
|
}
|
|
|
|
// Check for key exhaution
|
|
if len(search) == 0 {
|
|
break
|
|
}
|
|
|
|
// Look for an edge
|
|
_, n = n.getEdge(search[0])
|
|
if n == nil {
|
|
break
|
|
}
|
|
|
|
// Consume the search prefix
|
|
if bytes.HasPrefix(search, n.prefix) {
|
|
search = search[len(n.prefix):]
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
if last != nil {
|
|
return last.key, last.val, true
|
|
}
|
|
return nil, nil, false
|
|
}
|
|
|
|
// Minimum is used to return the minimum value in the tree
|
|
func (n *Node) Minimum() ([]byte, interface{}, bool) {
|
|
for {
|
|
if n.isLeaf() {
|
|
return n.leaf.key, n.leaf.val, true
|
|
}
|
|
if len(n.edges) > 0 {
|
|
n = n.edges[0].node
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
return nil, nil, false
|
|
}
|
|
|
|
// Maximum is used to return the maximum value in the tree
|
|
func (n *Node) Maximum() ([]byte, interface{}, bool) {
|
|
for {
|
|
if num := len(n.edges); num > 0 {
|
|
n = n.edges[num-1].node
|
|
continue
|
|
}
|
|
if n.isLeaf() {
|
|
return n.leaf.key, n.leaf.val, true
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
return nil, nil, false
|
|
}
|
|
|
|
// Iterator is used to return an iterator at
|
|
// the given node to walk the tree
|
|
func (n *Node) Iterator() *Iterator {
|
|
return &Iterator{node: n}
|
|
}
|
|
|
|
// rawIterator is used to return a raw iterator at the given node to walk the
|
|
// tree.
|
|
func (n *Node) rawIterator() *rawIterator {
|
|
iter := &rawIterator{node: n}
|
|
iter.Next()
|
|
return iter
|
|
}
|
|
|
|
// Walk is used to walk the tree
|
|
func (n *Node) Walk(fn WalkFn) {
|
|
recursiveWalk(n, fn)
|
|
}
|
|
|
|
// WalkPrefix is used to walk the tree under a prefix
|
|
func (n *Node) WalkPrefix(prefix []byte, fn WalkFn) {
|
|
search := prefix
|
|
for {
|
|
// Check for key exhaution
|
|
if len(search) == 0 {
|
|
recursiveWalk(n, fn)
|
|
return
|
|
}
|
|
|
|
// Look for an edge
|
|
_, n = n.getEdge(search[0])
|
|
if n == nil {
|
|
break
|
|
}
|
|
|
|
// Consume the search prefix
|
|
if bytes.HasPrefix(search, n.prefix) {
|
|
search = search[len(n.prefix):]
|
|
|
|
} else if bytes.HasPrefix(n.prefix, search) {
|
|
// Child may be under our search prefix
|
|
recursiveWalk(n, fn)
|
|
return
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// WalkPath is used to walk the tree, but only visiting nodes
|
|
// from the root down to a given leaf. Where WalkPrefix walks
|
|
// all the entries *under* the given prefix, this walks the
|
|
// entries *above* the given prefix.
|
|
func (n *Node) WalkPath(path []byte, fn WalkFn) {
|
|
search := path
|
|
for {
|
|
// Visit the leaf values if any
|
|
if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
|
|
return
|
|
}
|
|
|
|
// Check for key exhaution
|
|
if len(search) == 0 {
|
|
return
|
|
}
|
|
|
|
// Look for an edge
|
|
_, n = n.getEdge(search[0])
|
|
if n == nil {
|
|
return
|
|
}
|
|
|
|
// Consume the search prefix
|
|
if bytes.HasPrefix(search, n.prefix) {
|
|
search = search[len(n.prefix):]
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
func (n *Node) Seek(prefix []byte) *Seeker {
|
|
search := prefix
|
|
p := &pos{n: n}
|
|
for {
|
|
// Check for key exhaution
|
|
if len(search) == 0 {
|
|
return &Seeker{p}
|
|
}
|
|
|
|
num := len(n.edges)
|
|
idx := sort.Search(num, func(i int) bool {
|
|
return n.edges[i].label >= search[0]
|
|
})
|
|
p.current = idx
|
|
if idx < len(n.edges) {
|
|
n = n.edges[idx].node
|
|
if bytes.HasPrefix(search, n.prefix) && len(n.edges) > 0 {
|
|
search = search[len(n.prefix):]
|
|
p.current++
|
|
p = &pos{n: n, prev: p}
|
|
continue
|
|
}
|
|
}
|
|
p.current++
|
|
return &Seeker{p}
|
|
}
|
|
}
|
|
|
|
type Seeker struct {
|
|
*pos
|
|
}
|
|
|
|
type pos struct {
|
|
n *Node
|
|
current int
|
|
prev *pos
|
|
isLeaf bool
|
|
}
|
|
|
|
func (s *Seeker) Next() (k []byte, v interface{}, ok bool) {
|
|
if s.current >= len(s.n.edges) {
|
|
if s.prev == nil {
|
|
return nil, nil, false
|
|
}
|
|
s.pos = s.prev
|
|
return s.Next()
|
|
}
|
|
|
|
edge := s.n.edges[s.current]
|
|
s.current++
|
|
if edge.node.leaf != nil && !s.isLeaf {
|
|
s.isLeaf = true
|
|
s.current--
|
|
return edge.node.leaf.key, edge.node.leaf.val, true
|
|
}
|
|
s.isLeaf = false
|
|
s.pos = &pos{n: edge.node, prev: s.pos}
|
|
return s.Next()
|
|
}
|
|
|
|
// recursiveWalk is used to do a pre-order walk of a node
|
|
// recursively. Returns true if the walk should be aborted
|
|
func recursiveWalk(n *Node, fn WalkFn) bool {
|
|
// Visit the leaf values if any
|
|
if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
|
|
return true
|
|
}
|
|
|
|
// Recurse on the children
|
|
for _, e := range n.edges {
|
|
if recursiveWalk(e.node, fn) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|