node.go 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. package trie
  2. import "strings"
  3. // Node is the trie's node
  4. type Node[T any] struct {
  5. childMap map[string]*Node[T]
  6. childNode *Node[T] // optimize for only one child
  7. childStr string
  8. inited bool
  9. data T
  10. }
  11. func (n *Node[T]) getChild(s string) *Node[T] {
  12. if n.childMap == nil {
  13. if n.childNode != nil && n.childStr == s {
  14. return n.childNode
  15. }
  16. return nil
  17. }
  18. return n.childMap[s]
  19. }
  20. func (n *Node[T]) hasChild(s string) bool {
  21. return n.getChild(s) != nil
  22. }
  23. func (n *Node[T]) addChild(s string, child *Node[T]) {
  24. if n.childMap == nil {
  25. if n.childNode == nil {
  26. n.childStr = s
  27. n.childNode = child
  28. return
  29. }
  30. n.childMap = map[string]*Node[T]{}
  31. if n.childNode != nil {
  32. n.childMap[n.childStr] = n.childNode
  33. }
  34. n.childStr = ""
  35. n.childNode = nil
  36. }
  37. n.childMap[s] = child
  38. }
  39. func (n *Node[T]) getOrNewChild(s string) *Node[T] {
  40. node := n.getChild(s)
  41. if node == nil {
  42. node = newNode[T]()
  43. n.addChild(s, node)
  44. }
  45. return node
  46. }
  47. func (n *Node[T]) optimize() {
  48. if len(n.childStr) > 0 {
  49. n.childStr = strClone(n.childStr)
  50. }
  51. if n.childNode != nil {
  52. n.childNode.optimize()
  53. }
  54. if n.childMap == nil {
  55. return
  56. }
  57. switch len(n.childMap) {
  58. case 0:
  59. n.childMap = nil
  60. return
  61. case 1:
  62. for key := range n.childMap {
  63. n.childStr = key
  64. n.childNode = n.childMap[key]
  65. }
  66. n.childMap = nil
  67. n.optimize()
  68. return
  69. }
  70. children := make(map[string]*Node[T], len(n.childMap)) // avoid map reallocate memory
  71. for key := range n.childMap {
  72. child := n.childMap[key]
  73. if child == nil {
  74. continue
  75. }
  76. key = strClone(key)
  77. children[key] = child
  78. child.optimize()
  79. }
  80. n.childMap = children
  81. }
  82. func strClone(key string) string {
  83. switch key { // try to save string's memory
  84. case wildcard:
  85. key = wildcard
  86. case dotWildcard:
  87. key = dotWildcard
  88. case complexWildcard:
  89. key = complexWildcard
  90. case domainStep:
  91. key = domainStep
  92. default:
  93. key = strings.Clone(key)
  94. }
  95. return key
  96. }
  97. func (n *Node[T]) isEmpty() bool {
  98. if n == nil || n.inited == false {
  99. return true
  100. }
  101. return false
  102. }
  103. func (n *Node[T]) setData(data T) {
  104. n.data = data
  105. n.inited = true
  106. }
  107. func (n *Node[T]) getChildren() map[string]*Node[T] {
  108. if n.childMap == nil {
  109. if n.childNode != nil {
  110. m := make(map[string]*Node[T])
  111. m[n.childStr] = n.childNode
  112. return m
  113. }
  114. } else {
  115. return n.childMap
  116. }
  117. return nil
  118. }
  119. func (n *Node[T]) Data() T {
  120. return n.data
  121. }
  122. func newNode[T any]() *Node[T] {
  123. return &Node[T]{}
  124. }