iter.go 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. package iradix
  2. import "bytes"
  3. // Iterator is used to iterate over a set of nodes
  4. // in pre-order
  5. type Iterator struct {
  6. node *Node
  7. stack []edges
  8. }
  9. // SeekPrefixWatch is used to seek the iterator to a given prefix
  10. // and returns the watch channel of the finest granularity
  11. func (i *Iterator) SeekPrefixWatch(prefix []byte) (watch <-chan struct{}) {
  12. // Wipe the stack
  13. i.stack = nil
  14. n := i.node
  15. watch = n.mutateCh
  16. search := prefix
  17. for {
  18. // Check for key exhaution
  19. if len(search) == 0 {
  20. i.node = n
  21. return
  22. }
  23. // Look for an edge
  24. _, n = n.getEdge(search[0])
  25. if n == nil {
  26. i.node = nil
  27. return
  28. }
  29. // Update to the finest granularity as the search makes progress
  30. watch = n.mutateCh
  31. // Consume the search prefix
  32. if bytes.HasPrefix(search, n.prefix) {
  33. search = search[len(n.prefix):]
  34. } else if bytes.HasPrefix(n.prefix, search) {
  35. i.node = n
  36. return
  37. } else {
  38. i.node = nil
  39. return
  40. }
  41. }
  42. }
  43. // SeekPrefix is used to seek the iterator to a given prefix
  44. func (i *Iterator) SeekPrefix(prefix []byte) {
  45. i.SeekPrefixWatch(prefix)
  46. }
  47. // Next returns the next node in order
  48. func (i *Iterator) Next() ([]byte, interface{}, bool) {
  49. // Initialize our stack if needed
  50. if i.stack == nil && i.node != nil {
  51. i.stack = []edges{
  52. edges{
  53. edge{node: i.node},
  54. },
  55. }
  56. }
  57. for len(i.stack) > 0 {
  58. // Inspect the last element of the stack
  59. n := len(i.stack)
  60. last := i.stack[n-1]
  61. elem := last[0].node
  62. // Update the stack
  63. if len(last) > 1 {
  64. i.stack[n-1] = last[1:]
  65. } else {
  66. i.stack = i.stack[:n-1]
  67. }
  68. // Push the edges onto the frontier
  69. if len(elem.edges) > 0 {
  70. i.stack = append(i.stack, elem.edges)
  71. }
  72. // Return the leaf values if any
  73. if elem.leaf != nil {
  74. return elem.leaf.key, elem.leaf.val, true
  75. }
  76. }
  77. return nil, nil, false
  78. }