205 lines
4.1 KiB
Go
Raw Normal View History

2023-07-09 12:42:33 +02:00
package router
2023-07-09 17:46:17 +02:00
// Tree represents a radix tree.
2023-07-18 17:14:34 +02:00
type Tree[T any] struct {
2023-09-02 09:19:11 +02:00
root treeNode[T]
2023-07-09 12:42:33 +02:00
}
2023-07-09 17:46:17 +02:00
// Add adds a new element to the tree.
func (tree *Tree[T]) Add(path string, data T) {
2023-07-09 12:42:33 +02:00
// Search tree for equal parts until we can no longer proceed
i := 0
offset := 0
node := &tree.root
for {
begin:
switch node.kind {
case parameter:
// This only occurs when the same parameter based route is added twice.
// node: /post/:id|
// path: /post/:id|
if i == len(path) {
node.data = data
return
}
// When we hit a separator, we'll search for a fitting child.
if path[i] == separator {
2024-03-13 13:04:03 +01:00
node, offset, _ = node.end(path, data, i, offset)
goto next
2023-07-09 12:42:33 +02:00
}
default:
if i == len(path) {
// The path already exists.
// node: /blog|
// path: /blog|
if i-offset == len(node.prefix) {
node.data = data
return
}
// The path ended but the node prefix is longer.
// node: /blog|feed
// path: /blog|
node.split(i-offset, "", data)
return
}
// The node we just checked is entirely included in our path.
// node: /|
// path: /|blog
if i-offset == len(node.prefix) {
2024-03-13 13:04:03 +01:00
var control flow
2023-07-09 12:42:33 +02:00
node, offset, control = node.end(path, data, i, offset)
switch control {
2024-03-13 13:04:03 +01:00
case flowStop:
2023-07-09 12:42:33 +02:00
return
2024-03-13 13:04:03 +01:00
case flowBegin:
2023-07-09 12:42:33 +02:00
goto begin
2024-03-13 13:04:03 +01:00
case flowNext:
2023-07-09 12:42:33 +02:00
goto next
}
}
// We got a conflict.
// node: /b|ag
// path: /b|riefcase
if path[i] != node.prefix[i-offset] {
node.split(i-offset, path[i:], data)
return
}
}
next:
i++
}
}
2023-07-09 17:46:17 +02:00
// Lookup finds the data for the given path.
2024-03-14 23:26:59 +01:00
func (tree *Tree[T]) Lookup(path string) (T, []Parameter) {
var params []Parameter
2023-07-09 17:46:17 +02:00
data := tree.LookupNoAlloc(path, func(key string, value string) {
2024-03-14 23:26:59 +01:00
params = append(params, Parameter{key, value})
2023-07-09 17:46:17 +02:00
})
return data, params
}
// LookupNoAlloc finds the data for the given path without using any memory allocations.
func (tree *Tree[T]) LookupNoAlloc(path string, addParameter func(key string, value string)) T {
2023-07-09 12:42:33 +02:00
var (
2025-03-01 14:27:26 +01:00
i uint
parameterPath string
wildcardPath string
parameter *treeNode[T]
wildcard *treeNode[T]
node = &tree.root
2023-07-09 12:42:33 +02:00
)
2024-03-14 23:26:59 +01:00
// Skip the first loop iteration if the starting characters are equal
if len(path) > 0 && len(node.prefix) > 0 && path[0] == node.prefix[0] {
i = 1
}
2023-07-09 12:42:33 +02:00
begin:
// Search tree for equal parts until we can no longer proceed
2024-03-14 23:26:59 +01:00
for i < uint(len(path)) {
2023-07-09 12:42:33 +02:00
// The node we just checked is entirely included in our path.
// node: /|
// path: /|blog
2024-03-15 09:52:34 +01:00
if i == uint(len(node.prefix)) {
2023-07-09 12:42:33 +02:00
if node.wildcard != nil {
2024-03-14 23:26:59 +01:00
wildcard = node.wildcard
2024-03-15 09:52:34 +01:00
wildcardPath = path[i:]
2023-07-09 12:42:33 +02:00
}
2025-03-01 14:27:26 +01:00
parameter = node.parameter
parameterPath = path[i:]
2023-07-09 12:42:33 +02:00
char := path[i]
if char >= node.startIndex && char < node.endIndex {
index := node.indices[char-node.startIndex]
if index != 0 {
node = node.children[index]
2024-03-15 09:52:34 +01:00
path = path[i:]
i = 1
2023-07-09 12:42:33 +02:00
continue
}
}
// node: /|:id
// path: /|blog
if node.parameter != nil {
node = node.parameter
2024-03-15 09:52:34 +01:00
path = path[i:]
i = 1
2023-07-09 12:42:33 +02:00
2024-03-15 00:07:44 +01:00
for i < uint(len(path)) {
2023-07-09 12:42:33 +02:00
// node: /:id|/posts
// path: /123|/posts
if path[i] == separator {
2024-03-15 09:52:34 +01:00
addParameter(node.prefix, path[:i])
2023-07-09 12:42:33 +02:00
index := node.indices[separator-node.startIndex]
node = node.children[index]
2024-03-15 09:52:34 +01:00
path = path[i:]
i = 1
2023-07-09 12:42:33 +02:00
goto begin
}
i++
}
2024-03-15 00:07:44 +01:00
2024-03-15 09:52:34 +01:00
addParameter(node.prefix, path[:i])
2024-03-15 00:07:44 +01:00
return node.data
2023-07-09 12:42:33 +02:00
}
// node: /|*any
// path: /|image.png
2024-03-14 23:26:59 +01:00
goto notFound
2023-07-09 12:42:33 +02:00
}
// We got a conflict.
// node: /b|ag
// path: /b|riefcase
2024-03-15 09:52:34 +01:00
if path[i] != node.prefix[i] {
2024-03-14 23:26:59 +01:00
goto notFound
2023-07-09 12:42:33 +02:00
}
i++
}
2024-03-14 23:26:59 +01:00
// node: /blog|
// path: /blog|
2024-03-15 09:52:34 +01:00
if i == uint(len(node.prefix)) {
2024-03-14 23:26:59 +01:00
return node.data
}
// node: /|*any
// path: /|image.png
notFound:
2025-03-01 14:27:26 +01:00
if parameter != nil {
addParameter(parameter.prefix, parameterPath)
return parameter.data
}
2024-03-14 23:26:59 +01:00
if wildcard != nil {
2024-03-15 09:52:34 +01:00
addParameter(wildcard.prefix, wildcardPath)
2024-03-14 23:26:59 +01:00
return wildcard.data
}
var empty T
return empty
2023-07-09 12:42:33 +02:00
}
2023-09-06 16:52:22 +02:00
// Map binds all handlers to a new one provided by the callback.
func (tree *Tree[T]) Map(transform func(T) T) {
2023-07-09 12:42:33 +02:00
tree.root.each(func(node *treeNode[T]) {
2023-07-18 17:14:34 +02:00
node.data = transform(node.data)
2023-07-09 12:42:33 +02:00
})
}