使用前缀树实现支持动态参数和通配符的路由匹配

前缀树（Trie）是一种高效的字符串查找数据结构，在实现路由匹配功能时非常有用。本文将从前缀树的基本原理出发，讲解如何在 Go 中实现一个支持动态参数和通配符的路由匹配器，并附带单元测试。

什么是前缀树？

前缀树是一种树形结构，用于高效地存储和检索字符串集合中的单词。其特点是：

每个节点代表一个字符。
从根节点到某一节点的路径表示一个字符串的前缀。
前缀相同的字符串共享路径，从而节省存储空间。

例如，字符串集合 ["go", "golang", "gopher"] 构建的前缀树如下：

        (root)
         |
         g
         |
         o
        / \
       l   p
       |    \
       a     h
       |      \
       n       e
       |        \
       g         r

在路由匹配中，前缀树的节点不仅存储路径，还可以存储动态参数（如 :id）和通配符（如 *），以支持更灵活的匹配规则。

实现路由匹配器

数据结构设计

以下是前缀树的节点和树的定义：

// TrieNode 定义前缀树的节点
type TrieNode struct {
	children   map[string]*TrieNode       // 子节点
	isEnd      bool                       // 是否为路径结尾
	handler    func(params map[string]string) // 路由匹配的处理函数
	isParam    bool                       // 是否为动态参数节点
	isWildcard bool                       // 是否为通配符节点
}

// Trie 定义前缀树
type Trie struct {
	root *TrieNode
}

// NewTrie 初始化前缀树
func NewTrie() *Trie {
	return &Trie{root: &TrieNode{children: make(map[string]*TrieNode)}}
}

插入路径和处理函数

通过 Insert 方法将路径和对应的处理函数插入到前缀树中：

// Insert 插入路径及其处理函数
func (t *Trie) Insert(path string, handler func(params map[string]string)) {
	segments := strings.Split(strings.Trim(path, "/"), "/") // 按路径分段
	node := t.root
	for _, segment := range segments {
		isParam := strings.HasPrefix(segment, ":")
		isWildcard := segment == "*"

		// 根据动态参数和通配符处理
		key := segment
		if isParam {
			key = ":"
		} else if isWildcard {
			key = "*"
		}

		if _, exists := node.children[key]; !exists {
			node.children[key] = &TrieNode{
				children:   make(map[string]*TrieNode),
				isParam:    isParam,
				isWildcard: isWildcard,
			}
		}
		node = node.children[key]
	}
	node.isEnd = true
	node.handler = handler
}

路径匹配逻辑

通过 Match 方法从前缀树中查找与路径匹配的节点，并执行其处理函数：

// Match 匹配路径并执行处理函数
func (t *Trie) Match(path string) (bool, map[string]string) {
	segments := strings.Split(strings.Trim(path, "/"), "/")
	node := t.root
	params := make(map[string]string)

	for _, segment := range segments {
		if nextNode, exists := node.children[segment]; exists {
			node = nextNode
		} else if paramNode, exists := node.children[":"]; exists {
			// 动态参数匹配
			node = paramNode
			params[paramNodeKey(paramNode)] = segment
		} else if wildcardNode, exists := node.children["*"]; exists {
			// 通配符匹配
			node = wildcardNode
			break // 通配符节点匹配剩余路径，直接退出
		} else {
			return false, nil
		}
	}

	if node.isEnd && node.handler != nil {
		node.handler(params)
		return true, params
	}
	return false, nil
}

// paramNodeKey 返回动态参数节点的键值
func paramNodeKey(node *TrieNode) string {
	for key := range node.children {
		if node.children[key].isParam {
			return key
		}
	}
	return ""
}

单元测试

以下是针对上述功能的单元测试代码：

func TestTrie(t *testing.T) {
	trie := NewTrie()

	// 插入路径和处理函数
	trie.Insert("/home", func(params map[string]string) {
		fmt.Println("Welcome to Home!")
	})
	trie.Insert("/about", func(params map[string]string) {
		fmt.Println("About Us Page")
	})
	trie.Insert("/api/users/:id", func(params map[string]string) {
		fmt.Printf("User API - User ID: %s\n", params[":id"])
	})
	trie.Insert("/static/*", func(params map[string]string) {
		fmt.Println("Static files matched")
	})

	// 定义测试用例
	tests := []struct {
		path        string
		expectFound bool
		expectParams map[string]string
	}{
		{"/home", true, map[string]string{}},
		{"/about", true, map[string]string{}},
		{"/api/users/123", true, map[string]string{":id": "123"}},
		{"/static/js/app.js", true, map[string]string{}},
		{"/notfound", false, nil},
	}

	// 执行测试
	for _, tt := range tests {
		found, params := trie.Match(tt.path)
		if found != tt.expectFound {
			t.Errorf("Path: %s - Expected found: %v, got: %v", tt.path, tt.expectFound, found)
		}
		if len(params) != len(tt.expectParams) {
			t.Errorf("Path: %s - Expected params: %v, got: %v", tt.path, tt.expectParams, params)
		}
		for k, v := range tt.expectParams {
			if params[k] != v {
				t.Errorf("Path: %s - Expected param %s: %s, got: %s", tt.path, k, v, params[k])
			}
		}
	}
}

运行和测试

将代码保存为 trie.go 和 trie_test.go 两个文件。
在终端运行以下命令执行单元测试：

go test

如果所有测试通过，将输出类似以下内容：

ok  	command-line-arguments	0.003s

总结

通过本文，你学会了：

前缀树的基本原理。
如何使用 Go 实现一个支持动态参数和通配符的路由匹配器。
如何编写单元测试验证功能。

这一实现为开发高效的路由匹配器奠定了基础，适用于 Web 框架、API 网关等场景。如果有更多需求，比如正则匹配或更多复杂规则，可以进一步扩展此代码！