-
Notifications
You must be signed in to change notification settings - Fork 52
/
tree.go
131 lines (109 loc) · 2.38 KB
/
tree.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package znet
import (
"strings"
)
type (
// Tree records node
Tree struct {
root *Node
parameters Parameters
routes map[string]*Node
}
// Node records any URL params, and executes an end handler.
Node struct {
key string
// path records a request path
path string
handle HandlerFunc
// depth records Node's depth
depth int
// children records Node's children node
children map[string]*Node
// isPattern flag
isPattern bool
// middleware records middleware stack
middleware []HandlerFunc
}
)
func NewNode(key string, depth int) *Node {
return &Node{
key: key,
depth: depth,
children: make(map[string]*Node),
}
}
func NewTree() *Tree {
return &Tree{
root: NewNode("/", 1),
routes: make(map[string]*Node),
}
}
func (t *Tree) Add(pattern string, handle HandlerFunc, middleware ...HandlerFunc) {
var currentNode = t.root
if pattern != currentNode.key {
pattern = strings.TrimPrefix(pattern, "/")
res := strings.Split(pattern, "/")
for _, key := range res {
node, ok := currentNode.children[key]
if !ok {
node = NewNode(key, currentNode.depth+1)
if len(middleware) > 0 {
node.middleware = append(node.middleware, middleware...)
}
currentNode.children[key] = node
}
currentNode = node
}
}
if len(middleware) > 0 && currentNode.depth == 1 {
currentNode.middleware = append(currentNode.middleware, middleware...)
}
currentNode.handle = handle
currentNode.isPattern = true
currentNode.path = pattern
if routeName := t.parameters.routeName; routeName != "" {
t.routes[routeName] = currentNode
}
}
func (t *Tree) Find(pattern string, isRegex bool) (nodes []*Node) {
var (
node = t.root
queue []*Node
)
if pattern == node.path {
nodes = append(nodes, node)
return
}
if !isRegex {
pattern = strings.TrimPrefix(pattern, "/")
}
res := strings.Split(pattern, "/")
for _, key := range res {
child, ok := node.children[key]
if !ok && isRegex {
break
}
if !ok && !isRegex {
return
}
if pattern == child.path && !isRegex {
nodes = append(nodes, child)
return
}
node = child
}
queue = append(queue, node)
for len(queue) > 0 {
var queueTemp []*Node
for _, n := range queue {
if n.isPattern {
nodes = append(nodes, n)
}
for _, childNode := range n.children {
queueTemp = append(queueTemp, childNode)
}
}
queue = queueTemp
}
return
}