334 lines
7.5 KiB
Go
334 lines
7.5 KiB
Go
![]() |
package iradix
|
||
|
|
||
|
import (
|
||
|
"bytes"
|
||
|
|
||
|
"github.com/hashicorp/golang-lru/simplelru"
|
||
|
)
|
||
|
|
||
|
const (
|
||
|
// defaultModifiedCache is the default size of the modified node
|
||
|
// cache used per transaction. This is used to cache the updates
|
||
|
// to the nodes near the root, while the leaves do not need to be
|
||
|
// cached. This is important for very large transactions to prevent
|
||
|
// the modified cache from growing to be enormous.
|
||
|
defaultModifiedCache = 8192
|
||
|
)
|
||
|
|
||
|
// Tree implements an immutable radix tree. This can be treated as a
|
||
|
// Dictionary abstract data type. The main advantage over a standard
|
||
|
// hash map is prefix-based lookups and ordered iteration. The immutability
|
||
|
// means that it is safe to concurrently read from a Tree without any
|
||
|
// coordination.
|
||
|
type Tree struct {
|
||
|
root *Node
|
||
|
size int
|
||
|
}
|
||
|
|
||
|
// New returns an empty Tree
|
||
|
func New() *Tree {
|
||
|
t := &Tree{root: &Node{}}
|
||
|
return t
|
||
|
}
|
||
|
|
||
|
// Len is used to return the number of elements in the tree
|
||
|
func (t *Tree) Len() int {
|
||
|
return t.size
|
||
|
}
|
||
|
|
||
|
// Txn is a transaction on the tree. This transaction is applied
|
||
|
// atomically and returns a new tree when committed. A transaction
|
||
|
// is not thread safe, and should only be used by a single goroutine.
|
||
|
type Txn struct {
|
||
|
root *Node
|
||
|
size int
|
||
|
modified *simplelru.LRU
|
||
|
}
|
||
|
|
||
|
// Txn starts a new transaction that can be used to mutate the tree
|
||
|
func (t *Tree) Txn() *Txn {
|
||
|
txn := &Txn{
|
||
|
root: t.root,
|
||
|
size: t.size,
|
||
|
}
|
||
|
return txn
|
||
|
}
|
||
|
|
||
|
// writeNode returns a node to be modified, if the current
|
||
|
// node as already been modified during the course of
|
||
|
// the transaction, it is used in-place.
|
||
|
func (t *Txn) writeNode(n *Node) *Node {
|
||
|
// Ensure the modified set exists
|
||
|
if t.modified == nil {
|
||
|
lru, err := simplelru.NewLRU(defaultModifiedCache, nil)
|
||
|
if err != nil {
|
||
|
panic(err)
|
||
|
}
|
||
|
t.modified = lru
|
||
|
}
|
||
|
|
||
|
// If this node has already been modified, we can
|
||
|
// continue to use it during this transaction.
|
||
|
if _, ok := t.modified.Get(n); ok {
|
||
|
return n
|
||
|
}
|
||
|
|
||
|
// Copy the existing node
|
||
|
nc := new(Node)
|
||
|
if n.prefix != nil {
|
||
|
nc.prefix = make([]byte, len(n.prefix))
|
||
|
copy(nc.prefix, n.prefix)
|
||
|
}
|
||
|
if n.leaf != nil {
|
||
|
nc.leaf = new(leafNode)
|
||
|
*nc.leaf = *n.leaf
|
||
|
}
|
||
|
if len(n.edges) != 0 {
|
||
|
nc.edges = make([]edge, len(n.edges))
|
||
|
copy(nc.edges, n.edges)
|
||
|
}
|
||
|
|
||
|
// Mark this node as modified
|
||
|
t.modified.Add(n, nil)
|
||
|
return nc
|
||
|
}
|
||
|
|
||
|
// insert does a recursive insertion
|
||
|
func (t *Txn) insert(n *Node, k, search []byte, v interface{}) (*Node, interface{}, bool) {
|
||
|
// Handle key exhaution
|
||
|
if len(search) == 0 {
|
||
|
nc := t.writeNode(n)
|
||
|
if n.isLeaf() {
|
||
|
old := nc.leaf.val
|
||
|
nc.leaf.val = v
|
||
|
return nc, old, true
|
||
|
} else {
|
||
|
nc.leaf = &leafNode{
|
||
|
key: k,
|
||
|
val: v,
|
||
|
}
|
||
|
return nc, nil, false
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Look for the edge
|
||
|
idx, child := n.getEdge(search[0])
|
||
|
|
||
|
// No edge, create one
|
||
|
if child == nil {
|
||
|
e := edge{
|
||
|
label: search[0],
|
||
|
node: &Node{
|
||
|
leaf: &leafNode{
|
||
|
key: k,
|
||
|
val: v,
|
||
|
},
|
||
|
prefix: search,
|
||
|
},
|
||
|
}
|
||
|
nc := t.writeNode(n)
|
||
|
nc.addEdge(e)
|
||
|
return nc, nil, false
|
||
|
}
|
||
|
|
||
|
// Determine longest prefix of the search key on match
|
||
|
commonPrefix := longestPrefix(search, child.prefix)
|
||
|
if commonPrefix == len(child.prefix) {
|
||
|
search = search[commonPrefix:]
|
||
|
newChild, oldVal, didUpdate := t.insert(child, k, search, v)
|
||
|
if newChild != nil {
|
||
|
nc := t.writeNode(n)
|
||
|
nc.edges[idx].node = newChild
|
||
|
return nc, oldVal, didUpdate
|
||
|
}
|
||
|
return nil, oldVal, didUpdate
|
||
|
}
|
||
|
|
||
|
// Split the node
|
||
|
nc := t.writeNode(n)
|
||
|
splitNode := &Node{
|
||
|
prefix: search[:commonPrefix],
|
||
|
}
|
||
|
nc.replaceEdge(edge{
|
||
|
label: search[0],
|
||
|
node: splitNode,
|
||
|
})
|
||
|
|
||
|
// Restore the existing child node
|
||
|
modChild := t.writeNode(child)
|
||
|
splitNode.addEdge(edge{
|
||
|
label: modChild.prefix[commonPrefix],
|
||
|
node: modChild,
|
||
|
})
|
||
|
modChild.prefix = modChild.prefix[commonPrefix:]
|
||
|
|
||
|
// Create a new leaf node
|
||
|
leaf := &leafNode{
|
||
|
key: k,
|
||
|
val: v,
|
||
|
}
|
||
|
|
||
|
// If the new key is a subset, add to to this node
|
||
|
search = search[commonPrefix:]
|
||
|
if len(search) == 0 {
|
||
|
splitNode.leaf = leaf
|
||
|
return nc, nil, false
|
||
|
}
|
||
|
|
||
|
// Create a new edge for the node
|
||
|
splitNode.addEdge(edge{
|
||
|
label: search[0],
|
||
|
node: &Node{
|
||
|
leaf: leaf,
|
||
|
prefix: search,
|
||
|
},
|
||
|
})
|
||
|
return nc, nil, false
|
||
|
}
|
||
|
|
||
|
// delete does a recursive deletion
|
||
|
func (t *Txn) delete(parent, n *Node, search []byte) (*Node, *leafNode) {
|
||
|
// Check for key exhaution
|
||
|
if len(search) == 0 {
|
||
|
if !n.isLeaf() {
|
||
|
return nil, nil
|
||
|
}
|
||
|
|
||
|
// Remove the leaf node
|
||
|
nc := t.writeNode(n)
|
||
|
nc.leaf = nil
|
||
|
|
||
|
// Check if this node should be merged
|
||
|
if n != t.root && len(nc.edges) == 1 {
|
||
|
nc.mergeChild()
|
||
|
}
|
||
|
return nc, n.leaf
|
||
|
}
|
||
|
|
||
|
// Look for an edge
|
||
|
label := search[0]
|
||
|
idx, child := n.getEdge(label)
|
||
|
if child == nil || !bytes.HasPrefix(search, child.prefix) {
|
||
|
return nil, nil
|
||
|
}
|
||
|
|
||
|
// Consume the search prefix
|
||
|
search = search[len(child.prefix):]
|
||
|
newChild, leaf := t.delete(n, child, search)
|
||
|
if newChild == nil {
|
||
|
return nil, nil
|
||
|
}
|
||
|
|
||
|
// Copy this node
|
||
|
nc := t.writeNode(n)
|
||
|
|
||
|
// Delete the edge if the node has no edges
|
||
|
if newChild.leaf == nil && len(newChild.edges) == 0 {
|
||
|
nc.delEdge(label)
|
||
|
if n != t.root && len(nc.edges) == 1 && !nc.isLeaf() {
|
||
|
nc.mergeChild()
|
||
|
}
|
||
|
} else {
|
||
|
nc.edges[idx].node = newChild
|
||
|
}
|
||
|
return nc, leaf
|
||
|
}
|
||
|
|
||
|
// Insert is used to add or update a given key. The return provides
|
||
|
// the previous value and a bool indicating if any was set.
|
||
|
func (t *Txn) Insert(k []byte, v interface{}) (interface{}, bool) {
|
||
|
newRoot, oldVal, didUpdate := t.insert(t.root, k, k, v)
|
||
|
if newRoot != nil {
|
||
|
t.root = newRoot
|
||
|
}
|
||
|
if !didUpdate {
|
||
|
t.size++
|
||
|
}
|
||
|
return oldVal, didUpdate
|
||
|
}
|
||
|
|
||
|
// Delete is used to delete a given key. Returns the old value if any,
|
||
|
// and a bool indicating if the key was set.
|
||
|
func (t *Txn) Delete(k []byte) (interface{}, bool) {
|
||
|
newRoot, leaf := t.delete(nil, t.root, k)
|
||
|
if newRoot != nil {
|
||
|
t.root = newRoot
|
||
|
}
|
||
|
if leaf != nil {
|
||
|
t.size--
|
||
|
return leaf.val, true
|
||
|
}
|
||
|
return nil, false
|
||
|
}
|
||
|
|
||
|
// Root returns the current root of the radix tree within this
|
||
|
// transaction. The root is not safe across insert and delete operations,
|
||
|
// but can be used to read the current state during a transaction.
|
||
|
func (t *Txn) Root() *Node {
|
||
|
return t.root
|
||
|
}
|
||
|
|
||
|
// Get is used to lookup a specific key, returning
|
||
|
// the value and if it was found
|
||
|
func (t *Txn) Get(k []byte) (interface{}, bool) {
|
||
|
return t.root.Get(k)
|
||
|
}
|
||
|
|
||
|
// Commit is used to finalize the transaction and return a new tree
|
||
|
func (t *Txn) Commit() *Tree {
|
||
|
t.modified = nil
|
||
|
return &Tree{t.root, t.size}
|
||
|
}
|
||
|
|
||
|
// Insert is used to add or update a given key. The return provides
|
||
|
// the new tree, previous value and a bool indicating if any was set.
|
||
|
func (t *Tree) Insert(k []byte, v interface{}) (*Tree, interface{}, bool) {
|
||
|
txn := t.Txn()
|
||
|
old, ok := txn.Insert(k, v)
|
||
|
return txn.Commit(), old, ok
|
||
|
}
|
||
|
|
||
|
// Delete is used to delete a given key. Returns the new tree,
|
||
|
// old value if any, and a bool indicating if the key was set.
|
||
|
func (t *Tree) Delete(k []byte) (*Tree, interface{}, bool) {
|
||
|
txn := t.Txn()
|
||
|
old, ok := txn.Delete(k)
|
||
|
return txn.Commit(), old, ok
|
||
|
}
|
||
|
|
||
|
// Root returns the root node of the tree which can be used for richer
|
||
|
// query operations.
|
||
|
func (t *Tree) Root() *Node {
|
||
|
return t.root
|
||
|
}
|
||
|
|
||
|
// Get is used to lookup a specific key, returning
|
||
|
// the value and if it was found
|
||
|
func (t *Tree) Get(k []byte) (interface{}, bool) {
|
||
|
return t.root.Get(k)
|
||
|
}
|
||
|
|
||
|
// longestPrefix finds the length of the shared prefix
|
||
|
// of two strings
|
||
|
func longestPrefix(k1, k2 []byte) int {
|
||
|
max := len(k1)
|
||
|
if l := len(k2); l < max {
|
||
|
max = l
|
||
|
}
|
||
|
var i int
|
||
|
for i = 0; i < max; i++ {
|
||
|
if k1[i] != k2[i] {
|
||
|
break
|
||
|
}
|
||
|
}
|
||
|
return i
|
||
|
}
|
||
|
|
||
|
// concat two byte slices, returning a third new copy
|
||
|
func concat(a, b []byte) []byte {
|
||
|
c := make([]byte, len(a)+len(b))
|
||
|
copy(c, a)
|
||
|
copy(c[len(a):], b)
|
||
|
return c
|
||
|
}
|