]>
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 | ||
47 | func (n *node) replaceEdge(e edge) { | |
48 | num := len(n.edges) | |
49 | idx := sort.Search(num, func(i int) bool { | |
50 | return n.edges[i].label >= e.label | |
51 | }) | |
52 | if idx < num && n.edges[idx].label == e.label { | |
53 | n.edges[idx].node = e.node | |
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 | } | |
201 | parent.replaceEdge(edge{ | |
202 | label: search[0], | |
203 | node: child, | |
204 | }) | |
205 | ||
206 | // Restore the existing node | |
207 | child.addEdge(edge{ | |
208 | label: n.prefix[commonPrefix], | |
209 | node: n, | |
210 | }) | |
211 | n.prefix = n.prefix[commonPrefix:] | |
212 | ||
213 | // Create a new leaf node | |
214 | leaf := &leafNode{ | |
215 | key: s, | |
216 | val: v, | |
217 | } | |
218 | ||
219 | // If the new key is a subset, add to to this node | |
220 | search = search[commonPrefix:] | |
221 | if len(search) == 0 { | |
222 | child.leaf = leaf | |
223 | return nil, false | |
224 | } | |
225 | ||
226 | // Create a new edge for the node | |
227 | child.addEdge(edge{ | |
228 | label: search[0], | |
229 | node: &node{ | |
230 | leaf: leaf, | |
231 | prefix: search, | |
232 | }, | |
233 | }) | |
234 | return nil, false | |
235 | } | |
236 | } | |
237 | ||
238 | // Delete is used to delete a key, returning the previous | |
239 | // value and if it was deleted | |
240 | func (t *Tree) Delete(s string) (interface{}, bool) { | |
241 | var parent *node | |
242 | var label byte | |
243 | n := t.root | |
244 | search := s | |
245 | for { | |
246 | // Check for key exhaution | |
247 | if len(search) == 0 { | |
248 | if !n.isLeaf() { | |
249 | break | |
250 | } | |
251 | goto DELETE | |
252 | } | |
253 | ||
254 | // Look for an edge | |
255 | parent = n | |
256 | label = search[0] | |
257 | n = n.getEdge(label) | |
258 | if n == nil { | |
259 | break | |
260 | } | |
261 | ||
262 | // Consume the search prefix | |
263 | if strings.HasPrefix(search, n.prefix) { | |
264 | search = search[len(n.prefix):] | |
265 | } else { | |
266 | break | |
267 | } | |
268 | } | |
269 | return nil, false | |
270 | ||
271 | DELETE: | |
272 | // Delete the leaf | |
273 | leaf := n.leaf | |
274 | n.leaf = nil | |
275 | t.size-- | |
276 | ||
277 | // Check if we should delete this node from the parent | |
278 | if parent != nil && len(n.edges) == 0 { | |
279 | parent.delEdge(label) | |
280 | } | |
281 | ||
282 | // Check if we should merge this node | |
283 | if n != t.root && len(n.edges) == 1 { | |
284 | n.mergeChild() | |
285 | } | |
286 | ||
287 | // Check if we should merge the parent's other child | |
288 | if parent != nil && parent != t.root && len(parent.edges) == 1 && !parent.isLeaf() { | |
289 | parent.mergeChild() | |
290 | } | |
291 | ||
292 | return leaf.val, true | |
293 | } | |
294 | ||
295 | func (n *node) mergeChild() { | |
296 | e := n.edges[0] | |
297 | child := e.node | |
298 | n.prefix = n.prefix + child.prefix | |
299 | n.leaf = child.leaf | |
300 | n.edges = child.edges | |
301 | } | |
302 | ||
303 | // Get is used to lookup a specific key, returning | |
304 | // the value and if it was found | |
305 | func (t *Tree) Get(s string) (interface{}, bool) { | |
306 | n := t.root | |
307 | search := s | |
308 | for { | |
309 | // Check for key exhaution | |
310 | if len(search) == 0 { | |
311 | if n.isLeaf() { | |
312 | return n.leaf.val, true | |
313 | } | |
314 | break | |
315 | } | |
316 | ||
317 | // Look for an edge | |
318 | n = n.getEdge(search[0]) | |
319 | if n == nil { | |
320 | break | |
321 | } | |
322 | ||
323 | // Consume the search prefix | |
324 | if strings.HasPrefix(search, n.prefix) { | |
325 | search = search[len(n.prefix):] | |
326 | } else { | |
327 | break | |
328 | } | |
329 | } | |
330 | return nil, false | |
331 | } | |
332 | ||
333 | // LongestPrefix is like Get, but instead of an | |
334 | // exact match, it will return the longest prefix match. | |
335 | func (t *Tree) LongestPrefix(s string) (string, interface{}, bool) { | |
336 | var last *leafNode | |
337 | n := t.root | |
338 | search := s | |
339 | for { | |
340 | // Look for a leaf node | |
341 | if n.isLeaf() { | |
342 | last = n.leaf | |
343 | } | |
344 | ||
345 | // Check for key exhaution | |
346 | if len(search) == 0 { | |
347 | break | |
348 | } | |
349 | ||
350 | // Look for an edge | |
351 | n = n.getEdge(search[0]) | |
352 | if n == nil { | |
353 | break | |
354 | } | |
355 | ||
356 | // Consume the search prefix | |
357 | if strings.HasPrefix(search, n.prefix) { | |
358 | search = search[len(n.prefix):] | |
359 | } else { | |
360 | break | |
361 | } | |
362 | } | |
363 | if last != nil { | |
364 | return last.key, last.val, true | |
365 | } | |
366 | return "", nil, false | |
367 | } | |
368 | ||
369 | // Minimum is used to return the minimum value in the tree | |
370 | func (t *Tree) Minimum() (string, interface{}, bool) { | |
371 | n := t.root | |
372 | for { | |
373 | if n.isLeaf() { | |
374 | return n.leaf.key, n.leaf.val, true | |
375 | } | |
376 | if len(n.edges) > 0 { | |
377 | n = n.edges[0].node | |
378 | } else { | |
379 | break | |
380 | } | |
381 | } | |
382 | return "", nil, false | |
383 | } | |
384 | ||
385 | // Maximum is used to return the maximum value in the tree | |
386 | func (t *Tree) Maximum() (string, interface{}, bool) { | |
387 | n := t.root | |
388 | for { | |
389 | if num := len(n.edges); num > 0 { | |
390 | n = n.edges[num-1].node | |
391 | continue | |
392 | } | |
393 | if n.isLeaf() { | |
394 | return n.leaf.key, n.leaf.val, true | |
395 | } | |
396 | break | |
397 | } | |
398 | return "", nil, false | |
399 | } | |
400 | ||
401 | // Walk is used to walk the tree | |
402 | func (t *Tree) Walk(fn WalkFn) { | |
403 | recursiveWalk(t.root, fn) | |
404 | } | |
405 | ||
406 | // WalkPrefix is used to walk the tree under a prefix | |
407 | func (t *Tree) WalkPrefix(prefix string, fn WalkFn) { | |
408 | n := t.root | |
409 | search := prefix | |
410 | for { | |
411 | // Check for key exhaution | |
412 | if len(search) == 0 { | |
413 | recursiveWalk(n, fn) | |
414 | return | |
415 | } | |
416 | ||
417 | // Look for an edge | |
418 | n = n.getEdge(search[0]) | |
419 | if n == nil { | |
420 | break | |
421 | } | |
422 | ||
423 | // Consume the search prefix | |
424 | if strings.HasPrefix(search, n.prefix) { | |
425 | search = search[len(n.prefix):] | |
426 | ||
427 | } else if strings.HasPrefix(n.prefix, search) { | |
428 | // Child may be under our search prefix | |
429 | recursiveWalk(n, fn) | |
430 | return | |
431 | } else { | |
432 | break | |
433 | } | |
434 | } | |
435 | ||
436 | } | |
437 | ||
438 | // WalkPath is used to walk the tree, but only visiting nodes | |
439 | // from the root down to a given leaf. Where WalkPrefix walks | |
440 | // all the entries *under* the given prefix, this walks the | |
441 | // entries *above* the given prefix. | |
442 | func (t *Tree) WalkPath(path string, fn WalkFn) { | |
443 | n := t.root | |
444 | search := path | |
445 | for { | |
446 | // Visit the leaf values if any | |
447 | if n.leaf != nil && fn(n.leaf.key, n.leaf.val) { | |
448 | return | |
449 | } | |
450 | ||
451 | // Check for key exhaution | |
452 | if len(search) == 0 { | |
453 | return | |
454 | } | |
455 | ||
456 | // Look for an edge | |
457 | n = n.getEdge(search[0]) | |
458 | if n == nil { | |
459 | return | |
460 | } | |
461 | ||
462 | // Consume the search prefix | |
463 | if strings.HasPrefix(search, n.prefix) { | |
464 | search = search[len(n.prefix):] | |
465 | } else { | |
466 | break | |
467 | } | |
468 | } | |
469 | } | |
470 | ||
471 | // recursiveWalk is used to do a pre-order walk of a node | |
472 | // recursively. Returns true if the walk should be aborted | |
473 | func recursiveWalk(n *node, fn WalkFn) bool { | |
474 | // Visit the leaf values if any | |
475 | if n.leaf != nil && fn(n.leaf.key, n.leaf.val) { | |
476 | return true | |
477 | } | |
478 | ||
479 | // Recurse on the children | |
480 | for _, e := range n.edges { | |
481 | if recursiveWalk(e.node, fn) { | |
482 | return true | |
483 | } | |
484 | } | |
485 | return false | |
486 | } | |
487 | ||
488 | // ToMap is used to walk the tree and convert it into a map | |
489 | func (t *Tree) ToMap() map[string]interface{} { | |
490 | out := make(map[string]interface{}, t.size) | |
491 | t.Walk(func(k string, v interface{}) bool { | |
492 | out[k] = v | |
493 | return false | |
494 | }) | |
495 | return out | |
496 | } |