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 }