go-web/framework/trie.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) *node {
	uri = strings.ToUpper(uri)
	uri = strings.TrimPrefix(uri, "/")
	if uri == "" {
		return t.root
	}

	found := t.root.findRoute(uri)
	if found == nil {
		return nil
	}

	return found
}

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
	parent   *node
}

func (n *node) ParseParamsFromEndNode(uri string) map[string]string {
	ret := map[string]string{}
	uri = strings.ToUpper(uri)
	uri = strings.TrimPrefix(uri, "/")
	if uri == "" {
		// root
		return ret
	}
	segments := strings.Split(uri, "/")
	cnt := len(segments)
	cur := n

	for i := cnt - 1; i >= 0; i-- {
		if isWildcard(cur.segment) {
			// set params
			ret[cur.segment[1:]] = segments[i]
		}
		cur = cur.parent
	}
	return ret
}

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])
	new.parent = n
	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
}