patricia.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180
  1. // Copyright 2016 The CMux Authors. All rights reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
  12. // implied. See the License for the specific language governing
  13. // permissions and limitations under the License.
  14. package cmux
  15. import (
  16. "bytes"
  17. "io"
  18. )
  19. // patriciaTree is a simple patricia tree that handles []byte instead of string
  20. // and cannot be changed after instantiation.
  21. type patriciaTree struct {
  22. root *ptNode
  23. maxDepth int // max depth of the tree.
  24. }
  25. func newPatriciaTree(bs ...[]byte) *patriciaTree {
  26. max := 0
  27. for _, b := range bs {
  28. if max < len(b) {
  29. max = len(b)
  30. }
  31. }
  32. return &patriciaTree{
  33. root: newNode(bs),
  34. maxDepth: max + 1,
  35. }
  36. }
  37. func newPatriciaTreeString(strs ...string) *patriciaTree {
  38. b := make([][]byte, len(strs))
  39. for i, s := range strs {
  40. b[i] = []byte(s)
  41. }
  42. return newPatriciaTree(b...)
  43. }
  44. func (t *patriciaTree) matchPrefix(r io.Reader) bool {
  45. buf := make([]byte, t.maxDepth)
  46. n, _ := io.ReadFull(r, buf)
  47. return t.root.match(buf[:n], true)
  48. }
  49. func (t *patriciaTree) match(r io.Reader) bool {
  50. buf := make([]byte, t.maxDepth)
  51. n, _ := io.ReadFull(r, buf)
  52. return t.root.match(buf[:n], false)
  53. }
  54. type ptNode struct {
  55. prefix []byte
  56. next map[byte]*ptNode
  57. terminal bool
  58. }
  59. func newNode(strs [][]byte) *ptNode {
  60. if len(strs) == 0 {
  61. return &ptNode{
  62. prefix: []byte{},
  63. terminal: true,
  64. }
  65. }
  66. if len(strs) == 1 {
  67. return &ptNode{
  68. prefix: strs[0],
  69. terminal: true,
  70. }
  71. }
  72. p, strs := splitPrefix(strs)
  73. n := &ptNode{
  74. prefix: p,
  75. }
  76. nexts := make(map[byte][][]byte)
  77. for _, s := range strs {
  78. if len(s) == 0 {
  79. n.terminal = true
  80. continue
  81. }
  82. nexts[s[0]] = append(nexts[s[0]], s[1:])
  83. }
  84. n.next = make(map[byte]*ptNode)
  85. for first, rests := range nexts {
  86. n.next[first] = newNode(rests)
  87. }
  88. return n
  89. }
  90. func splitPrefix(bss [][]byte) (prefix []byte, rest [][]byte) {
  91. if len(bss) == 0 || len(bss[0]) == 0 {
  92. return prefix, bss
  93. }
  94. if len(bss) == 1 {
  95. return bss[0], [][]byte{{}}
  96. }
  97. for i := 0; ; i++ {
  98. var cur byte
  99. eq := true
  100. for j, b := range bss {
  101. if len(b) <= i {
  102. eq = false
  103. break
  104. }
  105. if j == 0 {
  106. cur = b[i]
  107. continue
  108. }
  109. if cur != b[i] {
  110. eq = false
  111. break
  112. }
  113. }
  114. if !eq {
  115. break
  116. }
  117. prefix = append(prefix, cur)
  118. }
  119. rest = make([][]byte, 0, len(bss))
  120. for _, b := range bss {
  121. rest = append(rest, b[len(prefix):])
  122. }
  123. return prefix, rest
  124. }
  125. func (n *ptNode) match(b []byte, prefix bool) bool {
  126. l := len(n.prefix)
  127. if l > 0 {
  128. if l > len(b) {
  129. l = len(b)
  130. }
  131. if !bytes.Equal(b[:l], n.prefix) {
  132. return false
  133. }
  134. }
  135. if n.terminal && (prefix || len(n.prefix) == len(b)) {
  136. return true
  137. }
  138. if l >= len(b) {
  139. return false
  140. }
  141. nextN, ok := n.next[b[l]]
  142. if !ok {
  143. return false
  144. }
  145. if l == len(b) {
  146. b = b[l:l]
  147. } else {
  148. b = b[l+1:]
  149. }
  150. return nextN.match(b, prefix)
  151. }