157 lines
3.4 KiB
Go
157 lines
3.4 KiB
Go
package framework
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
)
|
|
|
|
type Trie struct {
|
|
root *node
|
|
}
|
|
|
|
func NewTrie() *Trie {
|
|
return &Trie{root: newNode("")}
|
|
}
|
|
|
|
func (t *Trie) FindRoute(uri string) []ControllerHandler {
|
|
uri = strings.ToUpper(uri)
|
|
uri = strings.TrimPrefix(uri, "/")
|
|
if uri == "" {
|
|
return t.root.handlers
|
|
}
|
|
|
|
found := t.root.findRoute(uri)
|
|
if found == nil {
|
|
return nil
|
|
}
|
|
|
|
return found.handlers
|
|
}
|
|
|
|
func (t *Trie) AddRouter(uri string, handlers []ControllerHandler) error {
|
|
uri = strings.ToUpper(uri)
|
|
uri = strings.TrimPrefix(uri, "/")
|
|
if uri == "" {
|
|
t.root.isLast = true
|
|
t.root.handlers = append(t.root.handlers, handlers...)
|
|
return nil
|
|
}
|
|
|
|
upperUri := strings.ToUpper(uri)
|
|
match := t.FindRoute(upperUri)
|
|
if match != nil {
|
|
// existing route
|
|
return fmt.Errorf("existing route for %q", uri)
|
|
}
|
|
|
|
// The route does not exist, add it to the tree
|
|
err := t.root.addRoute(upperUri, handlers)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type node struct {
|
|
isLast bool
|
|
segment string
|
|
handlers []ControllerHandler
|
|
children []*node
|
|
}
|
|
|
|
func newNode(segment string) *node {
|
|
return &node{segment: segment}
|
|
}
|
|
|
|
func isWildcard(s string) bool {
|
|
return strings.HasPrefix(s, ":")
|
|
}
|
|
|
|
// /user/name
|
|
// /user/:id/name
|
|
// /user/3/name
|
|
|
|
// findRoute finds the handler for the uri if exists.
|
|
//
|
|
// we suppose that uri passed here doesn't begin with "/"
|
|
func (n *node) findRoute(uri string) *node {
|
|
splitted := strings.SplitN(uri, "/", 2)
|
|
splittedLen := len(splitted)
|
|
|
|
if isWildcard(splitted[0]) {
|
|
// input is a wildcard, check if this endpoint has already children
|
|
// if so, return the first one which is not nil.
|
|
nbChildren := len(n.children)
|
|
if nbChildren > 0 && n.children[0].segment != splitted[0] {
|
|
// several nodes exist, return the first one that is not nil
|
|
return n.children[0]
|
|
}
|
|
}
|
|
|
|
// try to find the value in childre
|
|
for _, child := range n.children {
|
|
if isWildcard(child.segment) || child.segment == splitted[0] {
|
|
if splittedLen == 1 {
|
|
// This is the last value, do the check and return
|
|
if child.isLast {
|
|
// if isLast, that means we have already registered the endpoint
|
|
return child
|
|
} else {
|
|
// otherwise, take it as not registered
|
|
return nil
|
|
}
|
|
}
|
|
// More segments to check
|
|
return child.findRoute(splitted[1])
|
|
}
|
|
}
|
|
|
|
// nothing found in the children
|
|
return nil
|
|
}
|
|
|
|
func (n *node) addRoute(uri string, handlers []ControllerHandler) error {
|
|
splitted := strings.SplitN(uri, "/", 2)
|
|
splittedLen := len(splitted)
|
|
isLast := splittedLen == 1
|
|
|
|
// try to find the value in childre
|
|
for _, child := range n.children {
|
|
if isWildcard(child.segment) || child.segment == splitted[0] {
|
|
if isLast {
|
|
// This is the last value, do the check and return
|
|
if child.isLast {
|
|
// if isLast, that means we have already registered the endpoint
|
|
return errors.New("node exists")
|
|
} else {
|
|
// otherwise, set the child
|
|
child.isLast = true
|
|
child.handlers = append(child.handlers, handlers...)
|
|
return nil
|
|
}
|
|
}
|
|
// More segments to check
|
|
return child.addRoute(splitted[1], handlers)
|
|
}
|
|
}
|
|
|
|
// create a new node
|
|
new := newNode(splitted[0])
|
|
if isLast {
|
|
// this is the end
|
|
new.handlers = append(new.handlers, handlers...)
|
|
new.isLast = true
|
|
n.children = append(n.children, new)
|
|
return nil
|
|
}
|
|
// continue
|
|
new.isLast = false
|
|
err := new.addRoute(splitted[1], handlers)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n.children = append(n.children, new)
|
|
return nil
|
|
}
|