]>
Commit | Line | Data |
---|---|---|
15c0b25d AP |
1 | package radix |
2 | ||
3 | import ( | |
4 | "sort" | |
5 | "strings" | |
6 | ) | |
7 | ||
8 | // WalkFn is used when walking the tree. Takes a | |
9 | // key and value, returning if iteration should | |
10 | // be terminated. | |
11 | type WalkFn func(s string, v interface{}) bool | |
12 | ||
13 | // leafNode is used to represent a value | |
14 | type leafNode struct { | |
15 | key string | |
16 | val interface{} | |
17 | } | |
18 | ||
19 | // edge is used to represent an edge node | |
20 | type edge struct { | |
21 | label byte | |
22 | node *node | |
23 | } | |
24 | ||
25 | type node struct { | |
26 | // leaf is used to store possible leaf | |
27 | leaf *leafNode | |
28 | ||
29 | // prefix is the common prefix we ignore | |
30 | prefix string | |
31 | ||
32 | // Edges should be stored in-order for iteration. | |
33 | // We avoid a fully materialized slice to save memory, | |
34 | // since in most cases we expect to be sparse | |
35 | edges edges | |
36 | } | |
37 | ||
38 | func (n *node) isLeaf() bool { | |
39 | return n.leaf != nil | |
40 | } | |
41 | ||
42 | func (n *node) addEdge(e edge) { | |
43 | n.edges = append(n.edges, e) | |
44 | n.edges.Sort() | |
45 | } | |
46 | ||
107c1cdb | 47 | func (n *node) updateEdge(label byte, node *node) { |
15c0b25d AP |
48 | num := len(n.edges) |
49 | idx := sort.Search(num, func(i int) bool { | |
107c1cdb | 50 | return n.edges[i].label >= label |
15c0b25d | 51 | }) |
107c1cdb ND |
52 | if idx < num && n.edges[idx].label == label { |
53 | n.edges[idx].node = node | |
15c0b25d AP |
54 | return |
55 | } | |
56 | panic("replacing missing edge") | |
57 | } | |
58 | ||
59 | func (n *node) getEdge(label byte) *node { | |
60 | num := len(n.edges) | |
61 | idx := sort.Search(num, func(i int) bool { | |
62 | return n.edges[i].label >= label | |
63 | }) | |
64 | if idx < num && n.edges[idx].label == label { | |
65 | return n.edges[idx].node | |
66 | } | |
67 | return nil | |
68 | } | |
69 | ||
70 | func (n *node) delEdge(label byte) { | |
71 | num := len(n.edges) | |
72 | idx := sort.Search(num, func(i int) bool { | |
73 | return n.edges[i].label >= label | |
74 | }) | |
75 | if idx < num && n.edges[idx].label == label { | |
76 | copy(n.edges[idx:], n.edges[idx+1:]) | |
77 | n.edges[len(n.edges)-1] = edge{} | |
78 | n.edges = n.edges[:len(n.edges)-1] | |
79 | } | |
80 | } | |
81 | ||
82 | type edges []edge | |
83 | ||
84 | func (e edges) Len() int { | |
85 | return len(e) | |
86 | } | |
87 | ||
88 | func (e edges) Less(i, j int) bool { | |
89 | return e[i].label < e[j].label | |
90 | } | |
91 | ||
92 | func (e edges) Swap(i, j int) { | |
93 | e[i], e[j] = e[j], e[i] | |
94 | } | |
95 | ||
96 | func (e edges) Sort() { | |
97 | sort.Sort(e) | |
98 | } | |
99 | ||
100 | // Tree implements a radix tree. This can be treated as a | |
101 | // Dictionary abstract data type. The main advantage over | |
102 | // a standard hash map is prefix-based lookups and | |
103 | // ordered iteration, | |
104 | type Tree struct { | |
105 | root *node | |
106 | size int | |
107 | } | |
108 | ||
109 | // New returns an empty Tree | |
110 | func New() *Tree { | |
111 | return NewFromMap(nil) | |
112 | } | |
113 | ||
114 | // NewFromMap returns a new tree containing the keys | |
115 | // from an existing map | |
116 | func NewFromMap(m map[string]interface{}) *Tree { | |
117 | t := &Tree{root: &node{}} | |
118 | for k, v := range m { | |
119 | t.Insert(k, v) | |
120 | } | |
121 | return t | |
122 | } | |
123 | ||
124 | // Len is used to return the number of elements in the tree | |
125 | func (t *Tree) Len() int { | |
126 | return t.size | |
127 | } | |
128 | ||
129 | // longestPrefix finds the length of the shared prefix | |
130 | // of two strings | |
131 | func longestPrefix(k1, k2 string) int { | |
132 | max := len(k1) | |
133 | if l := len(k2); l < max { | |
134 | max = l | |
135 | } | |
136 | var i int | |
137 | for i = 0; i < max; i++ { | |
138 | if k1[i] != k2[i] { | |
139 | break | |
140 | } | |
141 | } | |
142 | return i | |
143 | } | |
144 | ||
145 | // Insert is used to add a newentry or update | |
146 | // an existing entry. Returns if updated. | |
147 | func (t *Tree) Insert(s string, v interface{}) (interface{}, bool) { | |
148 | var parent *node | |
149 | n := t.root | |
150 | search := s | |
151 | for { | |
152 | // Handle key exhaution | |
153 | if len(search) == 0 { | |
154 | if n.isLeaf() { | |
155 | old := n.leaf.val | |
156 | n.leaf.val = v | |
157 | return old, true | |
158 | } | |
159 | ||
160 | n.leaf = &leafNode{ | |
161 | key: s, | |
162 | val: v, | |
163 | } | |
164 | t.size++ | |
165 | return nil, false | |
166 | } | |
167 | ||
168 | // Look for the edge | |
169 | parent = n | |
170 | n = n.getEdge(search[0]) | |
171 | ||
172 | // No edge, create one | |
173 | if n == nil { | |
174 | e := edge{ | |
175 | label: search[0], | |
176 | node: &node{ | |
177 | leaf: &leafNode{ | |
178 | key: s, | |
179 | val: v, | |
180 | }, | |
181 | prefix: search, | |
182 | }, | |
183 | } | |
184 | parent.addEdge(e) | |
185 | t.size++ | |
186 | return nil, false | |
187 | } | |
188 | ||
189 | // Determine longest prefix of the search key on match | |
190 | commonPrefix := longestPrefix(search, n.prefix) | |
191 | if commonPrefix == len(n.prefix) { | |
192 | search = search[commonPrefix:] | |
193 | continue | |
194 | } | |
195 | ||
196 | // Split the node | |
197 | t.size++ | |
198 | child := &node{ | |
199 | prefix: search[:commonPrefix], | |
200 | } | |
107c1cdb | 201 | parent.updateEdge(search[0], child) |
15c0b25d AP |
202 | |
203 | // Restore the existing node | |
204 | child.addEdge(edge{ | |
205 | label: n.prefix[commonPrefix], | |
206 | node: n, | |
207 | }) | |
208 | n.prefix = n.prefix[commonPrefix:] | |
209 | ||
210 | // Create a new leaf node | |
211 | leaf := &leafNode{ | |
212 | key: s, | |
213 | val: v, | |
214 | } | |
215 | ||
216 | // If the new key is a subset, add to to this node | |
217 | search = search[commonPrefix:] | |
218 | if len(search) == 0 { | |
219 | child.leaf = leaf | |
220 | return nil, false | |
221 | } | |
222 | ||
223 | // Create a new edge for the node | |
224 | child.addEdge(edge{ | |
225 | label: search[0], | |
226 | node: &node{ | |
227 | leaf: leaf, | |
228 | prefix: search, | |
229 | }, | |
230 | }) | |
231 | return nil, false | |
232 | } | |
233 | } | |
234 | ||
235 | // Delete is used to delete a key, returning the previous | |
236 | // value and if it was deleted | |
237 | func (t *Tree) Delete(s string) (interface{}, bool) { | |
238 | var parent *node | |
239 | var label byte | |
240 | n := t.root | |
241 | search := s | |
242 | for { | |
243 | // Check for key exhaution | |
244 | if len(search) == 0 { | |
245 | if !n.isLeaf() { | |
246 | break | |
247 | } | |
248 | goto DELETE | |
249 | } | |
250 | ||
251 | // Look for an edge | |
252 | parent = n | |
253 | label = search[0] | |
254 | n = n.getEdge(label) | |
255 | if n == nil { | |
256 | break | |
257 | } | |
258 | ||
259 | // Consume the search prefix | |
260 | if strings.HasPrefix(search, n.prefix) { | |
261 | search = search[len(n.prefix):] | |
262 | } else { | |
263 | break | |
264 | } | |
265 | } | |
266 | return nil, false | |
267 | ||
268 | DELETE: | |
269 | // Delete the leaf | |
270 | leaf := n.leaf | |
271 | n.leaf = nil | |
272 | t.size-- | |
273 | ||
274 | // Check if we should delete this node from the parent | |
275 | if parent != nil && len(n.edges) == 0 { | |
276 | parent.delEdge(label) | |
277 | } | |
278 | ||
279 | // Check if we should merge this node | |
280 | if n != t.root && len(n.edges) == 1 { | |
281 | n.mergeChild() | |
282 | } | |
283 | ||
284 | // Check if we should merge the parent's other child | |
285 | if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() { | |
286 | parent.mergeChild() | |
287 | } | |
288 | ||
289 | return leaf.val, true | |
290 | } | |
291 | ||
107c1cdb ND |
292 | // DeletePrefix is used to delete the subtree under a prefix |
293 | // Returns how many nodes were deleted | |
294 | // Use this to delete large subtrees efficiently | |
295 | func (t *Tree) DeletePrefix(s string) int { | |
296 | return t.deletePrefix(nil, t.root, s) | |
297 | } | |
298 | ||
299 | // delete does a recursive deletion | |
300 | func (t *Tree) deletePrefix(parent, n *node, prefix string) int { | |
301 | // Check for key exhaustion | |
302 | if len(prefix) == 0 { | |
303 | // Remove the leaf node | |
304 | subTreeSize := 0 | |
305 | //recursively walk from all edges of the node to be deleted | |
306 | recursiveWalk(n, func(s string, v interface{}) bool { | |
307 | subTreeSize++ | |
308 | return false | |
309 | }) | |
310 | if n.isLeaf() { | |
311 | n.leaf = nil | |
312 | } | |
313 | n.edges = nil // deletes the entire subtree | |
314 | ||
315 | // Check if we should merge the parent's other child | |
316 | if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() { | |
317 | parent.mergeChild() | |
318 | } | |
319 | t.size -= subTreeSize | |
320 | return subTreeSize | |
321 | } | |
322 | ||
323 | // Look for an edge | |
324 | label := prefix[0] | |
325 | child := n.getEdge(label) | |
326 | if child == nil || (!strings.HasPrefix(child.prefix, prefix) && !strings.HasPrefix(prefix, child.prefix)) { | |
327 | return 0 | |
328 | } | |
329 | ||
330 | // Consume the search prefix | |
331 | if len(child.prefix) > len(prefix) { | |
332 | prefix = prefix[len(prefix):] | |
333 | } else { | |
334 | prefix = prefix[len(child.prefix):] | |
335 | } | |
336 | return t.deletePrefix(n, child, prefix) | |
337 | } | |
338 | ||
15c0b25d AP |
339 | func (n *node) mergeChild() { |
340 | e := n.edges[0] | |
341 | child := e.node | |
342 | n.prefix = n.prefix + child.prefix | |
343 | n.leaf = child.leaf | |
344 | n.edges = child.edges | |
345 | } | |
346 | ||
347 | // Get is used to lookup a specific key, returning | |
348 | // the value and if it was found | |
349 | func (t *Tree) Get(s string) (interface{}, bool) { | |
350 | n := t.root | |
351 | search := s | |
352 | for { | |
353 | // Check for key exhaution | |
354 | if len(search) == 0 { | |
355 | if n.isLeaf() { | |
356 | return n.leaf.val, true | |
357 | } | |
358 | break | |
359 | } | |
360 | ||
361 | // Look for an edge | |
362 | n = n.getEdge(search[0]) | |
363 | if n == nil { | |
364 | break | |
365 | } | |
366 | ||
367 | // Consume the search prefix | |
368 | if strings.HasPrefix(search, n.prefix) { | |
369 | search = search[len(n.prefix):] | |
370 | } else { | |
371 | break | |
372 | } | |
373 | } | |
374 | return nil, false | |
375 | } | |
376 | ||
377 | // LongestPrefix is like Get, but instead of an | |
378 | // exact match, it will return the longest prefix match. | |
379 | func (t *Tree) LongestPrefix(s string) (string, interface{}, bool) { | |
380 | var last *leafNode | |
381 | n := t.root | |
382 | search := s | |
383 | for { | |
384 | // Look for a leaf node | |
385 | if n.isLeaf() { | |
386 | last = n.leaf | |
387 | } | |
388 | ||
389 | // Check for key exhaution | |
390 | if len(search) == 0 { | |
391 | break | |
392 | } | |
393 | ||
394 | // Look for an edge | |
395 | n = n.getEdge(search[0]) | |
396 | if n == nil { | |
397 | break | |
398 | } | |
399 | ||
400 | // Consume the search prefix | |
401 | if strings.HasPrefix(search, n.prefix) { | |
402 | search = search[len(n.prefix):] | |
403 | } else { | |
404 | break | |
405 | } | |
406 | } | |
407 | if last != nil { | |
408 | return last.key, last.val, true | |
409 | } | |
410 | return "", nil, false | |
411 | } | |
412 | ||
413 | // Minimum is used to return the minimum value in the tree | |
414 | func (t *Tree) Minimum() (string, interface{}, bool) { | |
415 | n := t.root | |
416 | for { | |
417 | if n.isLeaf() { | |
418 | return n.leaf.key, n.leaf.val, true | |
419 | } | |
420 | if len(n.edges) > 0 { | |
421 | n = n.edges[0].node | |
422 | } else { | |
423 | break | |
424 | } | |
425 | } | |
426 | return "", nil, false | |
427 | } | |
428 | ||
429 | // Maximum is used to return the maximum value in the tree | |
430 | func (t *Tree) Maximum() (string, interface{}, bool) { | |
431 | n := t.root | |
432 | for { | |
433 | if num := len(n.edges); num > 0 { | |
434 | n = n.edges[num-1].node | |
435 | continue | |
436 | } | |
437 | if n.isLeaf() { | |
438 | return n.leaf.key, n.leaf.val, true | |
439 | } | |
440 | break | |
441 | } | |
442 | return "", nil, false | |
443 | } | |
444 | ||
445 | // Walk is used to walk the tree | |
446 | func (t *Tree) Walk(fn WalkFn) { | |
447 | recursiveWalk(t.root, fn) | |
448 | } | |
449 | ||
450 | // WalkPrefix is used to walk the tree under a prefix | |
451 | func (t *Tree) WalkPrefix(prefix string, fn WalkFn) { | |
452 | n := t.root | |
453 | search := prefix | |
454 | for { | |
455 | // Check for key exhaution | |
456 | if len(search) == 0 { | |
457 | recursiveWalk(n, fn) | |
458 | return | |
459 | } | |
460 | ||
461 | // Look for an edge | |
462 | n = n.getEdge(search[0]) | |
463 | if n == nil { | |
464 | break | |
465 | } | |
466 | ||
467 | // Consume the search prefix | |
468 | if strings.HasPrefix(search, n.prefix) { | |
469 | search = search[len(n.prefix):] | |
470 | ||
471 | } else if strings.HasPrefix(n.prefix, search) { | |
472 | // Child may be under our search prefix | |
473 | recursiveWalk(n, fn) | |
474 | return | |
475 | } else { | |
476 | break | |
477 | } | |
478 | } | |
479 | ||
480 | } | |
481 | ||
482 | // WalkPath is used to walk the tree, but only visiting nodes | |
483 | // from the root down to a given leaf. Where WalkPrefix walks | |
484 | // all the entries *under* the given prefix, this walks the | |
485 | // entries *above* the given prefix. | |
486 | func (t *Tree) WalkPath(path string, fn WalkFn) { | |
487 | n := t.root | |
488 | search := path | |
489 | for { | |
490 | // Visit the leaf values if any | |
491 | if n.leaf != nil && fn(n.leaf.key, n.leaf.val) { | |
492 | return | |
493 | } | |
494 | ||
495 | // Check for key exhaution | |
496 | if len(search) == 0 { | |
497 | return | |
498 | } | |
499 | ||
500 | // Look for an edge | |
501 | n = n.getEdge(search[0]) | |
502 | if n == nil { | |
503 | return | |
504 | } | |
505 | ||
506 | // Consume the search prefix | |
507 | if strings.HasPrefix(search, n.prefix) { | |
508 | search = search[len(n.prefix):] | |
509 | } else { | |
510 | break | |
511 | } | |
512 | } | |
513 | } | |
514 | ||
515 | // recursiveWalk is used to do a pre-order walk of a node | |
516 | // recursively. Returns true if the walk should be aborted | |
517 | func recursiveWalk(n *node, fn WalkFn) bool { | |
518 | // Visit the leaf values if any | |
519 | if n.leaf != nil && fn(n.leaf.key, n.leaf.val) { | |
520 | return true | |
521 | } | |
522 | ||
523 | // Recurse on the children | |
524 | for _, e := range n.edges { | |
525 | if recursiveWalk(e.node, fn) { | |
526 | return true | |
527 | } | |
528 | } | |
529 | return false | |
530 | } | |
531 | ||
532 | // ToMap is used to walk the tree and convert it into a map | |
533 | func (t *Tree) ToMap() map[string]interface{} { | |
534 | out := make(map[string]interface{}, t.size) | |
535 | t.Walk(func(k string, v interface{}) bool { | |
536 | out[k] = v | |
537 | return false | |
538 | }) | |
539 | return out | |
540 | } |