github.com/cloudwego/kitex@v0.9.0/pkg/generic/descriptor/tree.go

/*
 * Copyright 2013 Julien Schmidt. All rights reserved.
 * Use of this source code is governed by a BSD-style license that can be found
 * in the LICENSE file.
 *
 * This file may have been modified by CloudWeGo authors. All CloudWeGo
 * Modifications are Copyright 2021 CloudWeGo Authors.
 */

package descriptor

import (
	"fmt"
	"net/url"
	"strings"
)

func countParams(path string) uint16 {
	var n uint
	for i := range []byte(path) {
		switch path[i] {
		case ':', '*':
			n++
		}
	}
	return uint16(n)
}

type nodeType uint8

const (
	static nodeType = iota // default
	param
	catchAll
	paramLabel = byte(':')
	anyLabel   = byte('*')
	slash      = "/"
	nilString  = ""
)

type (
	node struct {
		nType    nodeType
		label    byte
		prefix   string
		parent   *node
		children children
		// original path
		ppath string
		// param names
		pnames     []string
		function   *FunctionDescriptor
		paramChild *node
		anyChild   *node
		// isLeaf indicates that node does not have child routes
		isLeaf bool
	}
	children []*node
)

func checkPathValid(path string) {
	if path == nilString {
		panic("empty path")
	}
	if path[0] != '/' {
		panic("path must begin with '/'")
	}
	for i, c := range []byte(path) {
		switch c {
		case ':':
			if (i < len(path)-1 && path[i+1] == '/') || i == (len(path)-1) {
				panic("wildcards must be named with a non-empty name in path '" + path + "'")
			}
			i++
			for ; i < len(path) && path[i] != '/'; i++ {
				if path[i] == ':' || path[i] == '*' {
					panic("only one wildcard per path segment is allowed, find multi in path '" + path + "'")
				}
			}
		case '*':
			if i == len(path)-1 {
				panic("wildcards must be named with a non-empty name in path '" + path + "'")
			}
			if i > 0 && path[i-1] != '/' {
				panic(" no / before wildcards in path " + path)
			}
			for ; i < len(path); i++ {
				if path[i] == '/' {
					panic("catch-all routes are only allowed at the end of the path in path '" + path + "'")
				}
			}
		}
	}
}

// addRoute adds a node with the given function to the path.
// Not concurrency-safe!
func (n *node) addRoute(path string, function *FunctionDescriptor) {
	checkPathValid(path)

	var (
		pnames []string // Param names
		ppath  = path   // Pristine path
	)

	if function == nil {
		panic(fmt.Sprintf("adding route without handler function: %v", path))
	}

	// Add the front static route part of a non-static route
	for i, lcpIndex := 0, len(path); i < lcpIndex; i++ {
		// param route
		if path[i] == paramLabel {
			j := i + 1

			n.insert(path[:i], nil, static, nilString, nil)
			for ; i < lcpIndex && path[i] != '/'; i++ {
			}

			pnames = append(pnames, path[j:i])
			path = path[:j] + path[i:]
			i, lcpIndex = j, len(path)

			if i == lcpIndex {
				// path node is last fragment of route path. ie. `/users/:id`
				n.insert(path[:i], function, param, ppath, pnames)
				return
			} else {
				n.insert(path[:i], nil, param, nilString, pnames)
			}
		} else if path[i] == anyLabel {
			n.insert(path[:i], nil, static, nilString, nil)
			pnames = append(pnames, path[i+1:])
			n.insert(path[:i+1], function, catchAll, ppath, pnames)
			return
		}
	}
	n.insert(path, function, static, ppath, pnames)
}

func (n *node) insert(path string, function *FunctionDescriptor, t nodeType, ppath string, pnames []string) {
	currentNode := n
	search := path

	for {
		searchLen := len(search)
		prefixLen := len(currentNode.prefix)
		lcpLen := 0

		max := prefixLen
		if searchLen < max {
			max = searchLen
		}
		for ; lcpLen < max && search[lcpLen] == currentNode.prefix[lcpLen]; lcpLen++ {
		}

		if lcpLen == 0 {
			currentNode.label = search[0]
			currentNode.prefix = search
			if function != nil {
				currentNode.nType = t
				currentNode.function = function
				currentNode.ppath = ppath
				currentNode.pnames = pnames
			}
			currentNode.isLeaf = currentNode.children == nil && currentNode.paramChild == nil && currentNode.anyChild == nil
		} else if lcpLen < prefixLen {
			// Split node
			n := newNode(
				currentNode.nType,
				currentNode.prefix[lcpLen:],
				currentNode,
				currentNode.children,
				currentNode.function,
				currentNode.ppath,
				currentNode.pnames,
				currentNode.paramChild,
				currentNode.anyChild,
			)
			// Update parent path for all children to new node
			for _, child := range currentNode.children {
				child.parent = n
			}
			if currentNode.paramChild != nil {
				currentNode.paramChild.parent = n
			}
			if currentNode.anyChild != nil {
				currentNode.anyChild.parent = n
			}

			// Reset parent node
			currentNode.nType = static
			currentNode.label = currentNode.prefix[0]
			currentNode.prefix = currentNode.prefix[:lcpLen]
			currentNode.children = nil
			currentNode.function = nil
			currentNode.ppath = nilString
			currentNode.pnames = nil
			currentNode.paramChild = nil
			currentNode.anyChild = nil
			currentNode.isLeaf = false

			// Only Static children could reach here
			currentNode.children = append(currentNode.children, n)

			if lcpLen == searchLen {
				// At parent node
				currentNode.nType = t
				currentNode.function = function
				currentNode.ppath = ppath
				currentNode.pnames = pnames
			} else {
				// Create child node
				n = newNode(t, search[lcpLen:], currentNode, nil, function, ppath, pnames, nil, nil)
				// Only Static children could reach here
				currentNode.children = append(currentNode.children, n)
			}
			currentNode.isLeaf = currentNode.children == nil && currentNode.paramChild == nil && currentNode.anyChild == nil
		} else if lcpLen < searchLen {
			search = search[lcpLen:]
			c := currentNode.findChildWithLabel(search[0])
			if c != nil {
				// Go deeper
				currentNode = c
				continue
			}
			// Create child node
			n := newNode(t, search, currentNode, nil, function, ppath, pnames, nil, nil)
			switch t {
			case static:
				currentNode.children = append(currentNode.children, n)
			case param:
				currentNode.paramChild = n
			case catchAll:
				currentNode.anyChild = n
			}
			currentNode.isLeaf = currentNode.children == nil && currentNode.paramChild == nil && currentNode.anyChild == nil
		} else {
			// Node already exists
			if currentNode.function != nil && function != nil {
				panic("handlers are already registered for path '" + ppath + "'")
			}

			if function != nil {
				currentNode.function = function
				currentNode.ppath = ppath
				if len(currentNode.pnames) == 0 {
					currentNode.pnames = pnames
				}
			}
		}
		return
	}
}

// Returns the function registered with the given path (key). The values of
// wildcards are saved to a map.
// If no function can be found, a TSR (trailing slash redirect) recommendation is
// made if a function exists with an extra (without the) trailing slash for the
// given path.
func (n *node) getValue(path string, params func() *Params, unescape bool) (function *FunctionDescriptor, ps *Params, tsr bool) {
	var (
		cn          = n    // current node
		search      = path // current path
		searchIndex = 0
		paramIndex  int
	)

	backtrackToNextNodeType := func(fromNodeType nodeType) (nextNodeType nodeType, valid bool) {
		previous := cn
		cn = previous.parent
		valid = cn != nil

		// Next node type by priority
		if previous.nType == catchAll {
			nextNodeType = static
		} else {
			nextNodeType = previous.nType + 1
		}

		if fromNodeType == static {
			// when backtracking is done from static type block we did not change search so nothing to restore
			return
		}

		// restore search to value it was before we move to current node we are backtracking from.
		if previous.nType == static {
			searchIndex -= len(previous.prefix)
		} else {
			paramIndex--
			// for param/any node.prefix value is always `:`/`*` so we cannot deduce searchIndex from that and must use pValue
			// for that index as it would also contain part of path we cut off before moving into node we are backtracking from
			searchIndex -= len(ps.params[paramIndex].Value)
			ps.params = ps.params[:paramIndex]
		}
		search = path[searchIndex:]
		return
	}

	// search order: static > param > any
	for {
		if cn.nType == static {
			if len(search) >= len(cn.prefix) && cn.prefix == search[:len(cn.prefix)] {
				// Continue search
				search = search[len(cn.prefix):]
				searchIndex = searchIndex + len(cn.prefix)
			} else {
				// not equal
				if (len(cn.prefix) == len(search)+1) &&
					(cn.prefix[len(search)]) == '/' && cn.prefix[:len(search)] == search && (cn.function != nil || cn.anyChild != nil) {
					tsr = true
				}
				// No matching prefix, let's backtrack to the first possible alternative node of the decision path
				nk, ok := backtrackToNextNodeType(static)
				if !ok {
					return // No other possibilities on the decision path
				} else if nk == param {
					goto Param
				} else {
					// Not found (this should never be possible for static node we are looking currently)
					break
				}
			}
		}
		if search == nilString && cn.function != nil {
			function = cn.function
			break
		}

		// Static node
		if search != nilString {
			// If it can execute that logic, there is handler registered on the current node and search is `/`.
			if search == "/" && cn.function != nil {
				tsr = true
			}
			if child := cn.findChild(search[0]); child != nil {
				cn = child
				continue
			}
		}

		if search == nilString {
			if cd := cn.findChild('/'); cd != nil && (cd.function != nil || cd.anyChild != nil) {
				tsr = true
			}
		}

	Param:
		// Param node
		if child := cn.paramChild; search != nilString && child != nil {
			cn = child
			i := strings.Index(search, slash)
			if i == -1 {
				i = len(search)
			}
			if ps == nil {
				ps = params()
			}
			val := search[:i]
			if unescape {
				if v, err := url.QueryUnescape(val); err == nil {
					val = v
				}
			}
			ps.params = ps.params[:paramIndex+1]
			ps.params[paramIndex].Value = val
			paramIndex++
			search = search[i:]
			searchIndex = searchIndex + i
			if search == nilString {
				if cd := cn.findChild('/'); cd != nil && (cd.function != nil || cd.anyChild != nil) {
					tsr = true
				}
			}
			continue
		}
	Any:
		// Any node
		if child := cn.anyChild; child != nil {
			// If any node is found, use remaining path for paramValues
			cn = child
			if ps == nil {
				ps = params()
			}
			index := len(cn.pnames) - 1
			val := search
			if unescape {
				if v, err := url.QueryUnescape(val); err == nil {
					val = v
				}
			}
			ps.params = ps.params[:paramIndex+1]
			ps.params[index].Value = val
			// update indexes/search in case we need to backtrack when no handler match is found
			paramIndex++
			searchIndex += len(search)
			search = nilString
			function = cn.function
			break
		}

		// Let's backtrack to the first possible alternative node of the decision path
		nk, ok := backtrackToNextNodeType(catchAll)
		if !ok {
			break // No other possibilities on the decision path
		} else if nk == param {
			goto Param
		} else if nk == catchAll {
			goto Any
		} else {
			// Not found
			break
		}
	}

	if cn != nil {
		for i, name := range cn.pnames {
			ps.params[i].Key = name
		}
	}

	return
}

func (n *node) findChild(l byte) *node {
	for _, c := range n.children {
		if c.label == l {
			return c
		}
	}
	return nil
}

func (n *node) findChildWithLabel(l byte) *node {
	for _, c := range n.children {
		if c.label == l {
			return c
		}
	}
	if l == paramLabel {
		return n.paramChild
	}
	if l == anyLabel {
		return n.anyChild
	}
	return nil
}

func newNode(t nodeType, pre string, p *node, child children, f *FunctionDescriptor, ppath string, pnames []string, paramChildren, anyChildren *node) *node {
	return &node{
		nType:      t,
		label:      pre[0],
		prefix:     pre,
		parent:     p,
		children:   child,
		ppath:      ppath,
		pnames:     pnames,
		function:   f,
		paramChild: paramChildren,
		anyChild:   anyChildren,
		isLeaf:     child == nil && paramChildren == nil && anyChildren == nil,
	}
}