[AI生成] 基于前缀树go实现路由匹配

时间复杂度概览
操作              时间复杂度    说明
添加路由       O(L)               L为路径的段数（即深度）
查找路由       O(L)               L（同上）
空间复杂度    O(N)             N为所有路由的段总数

 

package main

import (
	"errors"
	"fmt"
	"strings"
)

// Handler 路由处理函数类型（演示用，实际可替换为 http.HandlerFunc）
type Handler func()

// TrieNode 前缀树节点
type TrieNode struct {
	children      map[string]*TrieNode // 静态子节点
	paramChild    *TrieNode            // 参数子节点（如 :id）
	wildcardChild *TrieNode            // 通配符子节点（如 *filepath）
	paramName     string               // 参数名（例如 "id"）
	wildcardName  string               // 通配符变量名（例如 "filepath"）
	handler       Handler              // 路由绑定的处理函数
}

// NewTrieNode 创建新节点
func NewTrieNode() *TrieNode {
	return &TrieNode{
		children: make(map[string]*TrieNode),
	}
}

// AddRoute 添加路由
// path 必须以 '/' 开头，例如 "/user/:id", "/static/*filepath"
func (n *TrieNode) AddRoute(path string, handler Handler) error {
	// 清理路径：去除末尾斜杠（根路径除外）
	if path != "/" && strings.HasSuffix(path, "/") {
		path = path[:len(path)-1]
	}
	if path == "" {
		return errors.New("路径不能为空")
	}
	if path[0] != '/' {
		return errors.New("路径必须以 / 开头")
	}

	// 分割路径段
	parts := strings.Split(path[1:], "/")
	if len(parts) == 1 && parts[0] == "" { // 根路径
		parts = []string{}
	}

	return n.insert(parts, handler)
}

// insert 递归插入
func (n *TrieNode) insert(parts []string, handler Handler) error {
	if len(parts) == 0 {
		if n.handler != nil {
			return errors.New("路由冲突：该路径已有处理函数")
		}
		n.handler = handler
		return nil
	}

	part := parts[0]
	remaining := parts[1:]

	// 处理通配符（必须以 '*' 开头，且必须是最后一段）
	if strings.HasPrefix(part, "*") {
		if n.wildcardChild != nil {
			return errors.New("路由冲突：通配符子节点已存在")
		}
		if len(remaining) > 0 {
			return errors.New("通配符只能出现在路径最后一段")
		}
		node := NewTrieNode()
		node.wildcardName = part[1:] // 去掉 '*'
		n.wildcardChild = node
		return node.insert(remaining, handler)
	}

	// 处理参数（必须以 ':' 开头）
	if strings.HasPrefix(part, ":") {
		if n.paramChild != nil {
			// 允许覆盖参数名不同但位置相同的路由？这里做严格冲突检查
			if n.paramChild.paramName != part[1:] {
				return errors.New("路由冲突：同一位置已有不同名的参数")
			}
		} else {
			node := NewTrieNode()
			node.paramName = part[1:] // 去掉 ':'
			n.paramChild = node
		}
		return n.paramChild.insert(remaining, handler)
	}

	// 静态段
	child, ok := n.children[part]
	if !ok {
		child = NewTrieNode()
		n.children[part] = child
	}
	return child.insert(remaining, handler)
}

// Search 查找路由，返回处理函数和路径参数
func (n *TrieNode) Search(path string) (Handler, map[string]string) {
	// 清理路径，移除末尾斜杠（根除外）
	if path != "/" && strings.HasSuffix(path, "/") {
		path = path[:len(path)-1]
	}
	if path == "" {
		return nil, nil
	}
	if path[0] != '/' {
		return nil, nil
	}

	parts := strings.Split(path[1:], "/")
	if len(parts) == 1 && parts[0] == "" {
		parts = []string{}
	}
	params := make(map[string]string)
	node := n.findNode(parts, params)
	if node == nil {
		return nil, nil
	}
	return node.handler, params
}

// findNode 递归查找节点，同时填充 params
func (n *TrieNode) findNode(parts []string, params map[string]string) *TrieNode {
	if len(parts) == 0 {
		return n
	}

	part := parts[0]
	remaining := parts[1:]

	// 优先静态匹配
	if child, ok := n.children[part]; ok {
		if node := child.findNode(remaining, params); node != nil {
			return node
		}
	}

	// 其次参数匹配
	if n.paramChild != nil {
		params[n.paramChild.paramName] = part
		if node := n.paramChild.findNode(remaining, params); node != nil {
			return node
		}
		// 参数匹配失败需要回滚（本例中不会发生，因为 paramChild 没有其他分支）
		delete(params, n.paramChild.paramName)
	}

	// 最后通配符匹配（必须为最后一段）
	if n.wildcardChild != nil {
		// 通配符匹配剩余所有部分
		params[n.wildcardChild.wildcardName] = strings.Join(append([]string{part}, remaining...), "/")
		return n.wildcardChild // 通配符节点没有子节点
	}

	return nil
}

// 演示使用
func main() {
	root := NewTrieNode()

	// 添加路由
	root.AddRoute("/", func() { fmt.Println("Home") })
	root.AddRoute("/user/profile", func() { fmt.Println("User Profile") })
	root.AddRoute("/user/:id", func() { fmt.Println("Get User by ID") })
	root.AddRoute("/post/:pid/comment/:cid", func() { fmt.Println("Post Comment") })
	root.AddRoute("/static/*filepath", func() { fmt.Println("Static File") })

	// 测试用例
	testCases := []string{
		"/",
		"/user/profile",
		"/user/123",
		"/post/42/comment/7",
		"/static/css/style.css",
		"/notfound",
	}

	for _, tc := range testCases {
		handler, params := root.Search(tc)
		if handler == nil {
			fmt.Printf("路径 %-25s -> 未匹配\n", tc)
			continue
		}
		fmt.Printf("路径 %-25s -> 匹配成功, 参数: %v\n", tc, params)
		// 实际调用时可以使用 handler()
	}
}