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(nc, 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 }