| package radix |
| |
| import ( |
| "sort" |
| "strings" |
| ) |
| |
| // WalkFn is used when walking the tree. Takes a |
| // key and value, returning if iteration should |
| // be terminated. |
| type WalkFn func(s string, v interface{}) bool |
| |
| // leafNode is used to represent a value |
| type leafNode struct { |
| key string |
| val interface{} |
| } |
| |
| // edge is used to represent an edge node |
| type edge struct { |
| label byte |
| node *node |
| } |
| |
| type node struct { |
| // leaf is used to store possible leaf |
| leaf *leafNode |
| |
| // prefix is the common prefix we ignore |
| prefix string |
| |
| // 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) { |
| n.edges = append(n.edges, e) |
| n.edges.Sort() |
| } |
| |
| 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) *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 n.edges[idx].node |
| } |
| return nil |
| } |
| |
| type edges []edge |
| |
| func (e edges) Len() int { |
| return len(e) |
| } |
| |
| func (e edges) Less(i, j int) bool { |
| return e[i].label < e[j].label |
| } |
| |
| func (e edges) Swap(i, j int) { |
| e[i], e[j] = e[j], e[i] |
| } |
| |
| func (e edges) Sort() { |
| sort.Sort(e) |
| } |
| |
| // Tree implements a 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, |
| type Tree struct { |
| root *node |
| size int |
| } |
| |
| // New returns an empty Tree |
| func New() *Tree { |
| return NewFromMap(nil) |
| } |
| |
| // NewFromMap returns a new tree containing the keys |
| // from an existing map |
| func NewFromMap(m map[string]interface{}) *Tree { |
| t := &Tree{root: &node{}} |
| for k, v := range m { |
| t.Insert(k, v) |
| } |
| return t |
| } |
| |
| // Len is used to return the number of elements in the tree |
| func (t *Tree) Len() int { |
| return t.size |
| } |
| |
| // longestPrefix finds the length of the shared prefix |
| // of two strings |
| func longestPrefix(k1, k2 string) 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 |
| } |
| |
| // Insert is used to add a newentry or update |
| // an existing entry. Returns if updated. |
| func (t *Tree) Insert(s string, v interface{}) (interface{}, bool) { |
| var parent *node |
| n := t.root |
| search := s |
| for { |
| // Handle key exhaution |
| if len(search) == 0 { |
| if n.isLeaf() { |
| old := n.leaf.val |
| n.leaf.val = v |
| return old, true |
| } else { |
| n.leaf = &leafNode{ |
| key: s, |
| val: v, |
| } |
| t.size++ |
| return nil, false |
| } |
| } |
| |
| // Look for the edge |
| parent = n |
| n = n.getEdge(search[0]) |
| |
| // No edge, create one |
| if n == nil { |
| e := edge{ |
| label: search[0], |
| node: &node{ |
| leaf: &leafNode{ |
| key: s, |
| val: v, |
| }, |
| prefix: search, |
| }, |
| } |
| parent.addEdge(e) |
| t.size++ |
| return nil, false |
| } |
| |
| // Determine longest prefix of the search key on match |
| commonPrefix := longestPrefix(search, n.prefix) |
| if commonPrefix == len(n.prefix) { |
| search = search[commonPrefix:] |
| continue |
| } |
| |
| // Split the node |
| t.size++ |
| child := &node{ |
| prefix: search[:commonPrefix], |
| } |
| parent.replaceEdge(edge{ |
| label: search[0], |
| node: child, |
| }) |
| |
| // Restore the existing node |
| child.addEdge(edge{ |
| label: n.prefix[commonPrefix], |
| node: n, |
| }) |
| n.prefix = n.prefix[commonPrefix:] |
| |
| // Create a new leaf node |
| leaf := &leafNode{ |
| key: s, |
| val: v, |
| } |
| |
| // If the new key is a subset, add to to this node |
| search = search[commonPrefix:] |
| if len(search) == 0 { |
| child.leaf = leaf |
| return nil, false |
| } |
| |
| // Create a new edge for the node |
| child.addEdge(edge{ |
| label: search[0], |
| node: &node{ |
| leaf: leaf, |
| prefix: search, |
| }, |
| }) |
| return nil, false |
| } |
| return nil, false |
| } |
| |
| // Delete is used to delete a key, returning the previous |
| // value and if it was deleted |
| func (t *Tree) Delete(s string) (interface{}, bool) { |
| n := t.root |
| search := s |
| for { |
| // Check for key exhaution |
| if len(search) == 0 { |
| if !n.isLeaf() { |
| break |
| } |
| goto DELETE |
| } |
| |
| // Look for an edge |
| n = n.getEdge(search[0]) |
| if n == nil { |
| break |
| } |
| |
| // Consume the search prefix |
| if strings.HasPrefix(search, n.prefix) { |
| search = search[len(n.prefix):] |
| } else { |
| break |
| } |
| } |
| return nil, false |
| |
| DELETE: |
| // Delete the leaf |
| leaf := n.leaf |
| n.leaf = nil |
| t.size-- |
| |
| // Check if we should merge this node |
| if len(n.edges) == 1 { |
| e := n.edges[0] |
| child := e.node |
| n.prefix = n.prefix + child.prefix |
| n.leaf = child.leaf |
| n.edges = child.edges |
| } |
| return leaf.val, true |
| } |
| |
| // Get is used to lookup a specific key, returning |
| // the value and if it was found |
| func (t *Tree) Get(s string) (interface{}, bool) { |
| n := t.root |
| search := s |
| for { |
| // Check for key exhaution |
| if len(search) == 0 { |
| if n.isLeaf() { |
| return n.leaf.val, true |
| } |
| break |
| } |
| |
| // Look for an edge |
| n = n.getEdge(search[0]) |
| if n == nil { |
| break |
| } |
| |
| // Consume the search prefix |
| if strings.HasPrefix(search, n.prefix) { |
| search = search[len(n.prefix):] |
| } else { |
| break |
| } |
| } |
| return nil, false |
| } |
| |
| // LongestPrefix is like Get, but instead of an |
| // exact match, it will return the longest prefix match. |
| func (t *Tree) LongestPrefix(s string) (string, interface{}, bool) { |
| var last *leafNode |
| n := t.root |
| search := s |
| 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 strings.HasPrefix(search, n.prefix) { |
| search = search[len(n.prefix):] |
| } else { |
| break |
| } |
| } |
| if last != nil { |
| return last.key, last.val, true |
| } |
| return "", nil, false |
| } |
| |
| // Minimum is used to return the minimum value in the tree |
| func (t *Tree) Minimum() (string, interface{}, bool) { |
| n := t.root |
| 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, false |
| } |
| |
| // Maximum is used to return the maximum value in the tree |
| func (t *Tree) Maximum() (string, interface{}, bool) { |
| n := t.root |
| 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, false |
| } |
| |
| // Walk is used to walk the tree |
| func (t *Tree) Walk(fn WalkFn) { |
| recursiveWalk(t.root, fn) |
| } |
| |
| // WalkPrefix is used to walk the tree under a prefix |
| func (t *Tree) WalkPrefix(prefix string, fn WalkFn) { |
| n := t.root |
| 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 strings.HasPrefix(search, n.prefix) { |
| search = search[len(n.prefix):] |
| |
| } else if strings.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 (t *Tree) WalkPath(path string, fn WalkFn) { |
| n := t.root |
| 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 strings.HasPrefix(search, n.prefix) { |
| search = search[len(n.prefix):] |
| } else { |
| break |
| } |
| } |
| } |
| |
| // 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 |
| } |
| |
| // ToMap is used to walk the tree and convert it into a map |
| func (t *Tree) ToMap() map[string]interface{} { |
| out := make(map[string]interface{}, t.size) |
| t.Walk(func(k string, v interface{}) bool { |
| out[k] = v |
| return false |
| }) |
| return out |
| } |