Router/treeNode.go

package router

import "strings"

// Node types
const (
	separator = '/'
	parameter = ':'
	wildcard  = '*'
)

// A node on our radix tree
type TreeNode[T any] struct {
	prefix    string
	data      T
	children  []*TreeNode[T]
	parameter *TreeNode[T]
	wildcard  *TreeNode[T]
	indexes   []uint8
	start     uint8
	end       uint8
	kind      byte
}

// Splits the node at the given index and inserts a new child
// node with the given path and data. If path is empty, it will
// not create another child node and instead assign the data
// directly to the node.
func (node *TreeNode[T]) split(index int, path string, data T) {
	// Create split node with the remaining string
	splitNode := node.clone(node.prefix[index:])

	// The existing data must be removed
	node.reset(node.prefix[:index])

	// If the path is empty, it means we don't create a 2nd child node.
	// Just assign the data for the existing node and store a single child node.
	if path == "" {
		node.data = data
		node.addChild(splitNode)
		return
	}

	node.addChild(splitNode)

	// Create new nodes with the remaining path
	node.append(path, data)
}

// Clone the node with a new prefix
func (node *TreeNode[T]) clone(prefix string) *TreeNode[T] {
	return &TreeNode[T]{
		prefix:    prefix,
		data:      node.data,
		children:  node.children,
		parameter: node.parameter,
		wildcard:  node.wildcard,
		indexes:   node.indexes,
		start:     node.start,
		end:       node.end,
		kind:      node.kind,
	}
}

// Reset the node, set the prefix
func (node *TreeNode[T]) reset(prefix string) {
	var empty T
	node.prefix = prefix
	node.data = empty
	node.children = nil
	node.parameter = nil
	node.wildcard = nil
	node.indexes = nil
	node.start = 0
	node.end = 0
	node.kind = 0
}

// Append the given path to the tree
func (node *TreeNode[T]) append(path string, data T) {
	// At this point, all we know is that somewhere
	// in the remaining string we have parameters.
	// node: /user|
	// path: /user|/:userid
	for {
		if path == "" {
			node.data = data
			return
		}

		paramStart := strings.IndexByte(path, parameter)

		if paramStart == -1 {
			paramStart = strings.IndexByte(path, wildcard)
		}

		// If it's a static route we are adding,
		// just add the remainder as a normal node.
		if paramStart == -1 {
			// If the node itself doesn't have a prefix (root node),
			// don't add a child and use the node itself.
			if node.prefix == "" {
				node.prefix = path
				node.data = data
				node.addTrailingSlash(data)
				return
			}

			child := &TreeNode[T]{
				prefix: path,
				data:   data,
			}

			node.addChild(child)
			child.addTrailingSlash(data)
			return
		}

		// If we're directly in front of a parameter,
		// add a parameter node.
		if paramStart == 0 {
			paramEnd := strings.IndexByte(path, separator)

			if paramEnd == -1 {
				paramEnd = len(path)
			}

			child := &TreeNode[T]{
				prefix: path[1:paramEnd],
				kind:   path[paramStart],
			}

			switch child.kind {
			case parameter:
				child.addTrailingSlash(data)
				node.parameter = child
				node = child
				path = path[paramEnd:]
				continue

			case wildcard:
				child.data = data
				node.wildcard = child
				return
			}
		}

		// We know there's a parameter, but not directly at the start.

		// If the node itself doesn't have a prefix (root node),
		// don't add a child and use the node itself.
		if node.prefix == "" {
			node.prefix = path[:paramStart]
			path = path[paramStart:]
			continue
		}

		// Add a normal node with the path before the parameter start.
		child := &TreeNode[T]{
			prefix: path[:paramStart],
		}

		// Allow trailing slashes to return
		// the same content as their parent node.
		if child.prefix == "/" {
			child.data = node.data
		}

		node.addChild(child)
		node = child
		path = path[paramStart:]
	}
}

// Add a child tree node
func (node *TreeNode[T]) addChild(child *TreeNode[T]) {
	if len(node.children) == 0 {
		node.children = append(node.children, nil)
	}

	firstChar := child.prefix[0]

	switch {
	case node.start == 0:
		node.start = firstChar
		node.indexes = []uint8{0}
		node.end = node.start + uint8(len(node.indexes))

	case firstChar < node.start:
		diff := node.start - firstChar
		newindexes := make([]uint8, diff+uint8(len(node.indexes)))
		copy(newindexes[diff:], node.indexes)
		node.start = firstChar
		node.indexes = newindexes
		node.end = node.start + uint8(len(node.indexes))

	case firstChar >= node.end:
		diff := firstChar - node.end + 1
		newindexes := make([]uint8, diff+uint8(len(node.indexes)))
		copy(newindexes, node.indexes)
		node.indexes = newindexes
		node.end = node.start + uint8(len(node.indexes))
	}

	index := node.indexes[firstChar-node.start]

	if index == 0 {
		node.indexes[firstChar-node.start] = uint8(len(node.children))
		node.children = append(node.children, child)
		return
	}

	node.children[index] = child
}

func (node *TreeNode[T]) addTrailingSlash(data T) {
	if strings.HasSuffix(node.prefix, "/") || node.kind == wildcard || (separator >= node.start && separator < node.end && node.indexes[separator-node.start] != 0) {
		return
	}

	node.addChild(&TreeNode[T]{
		prefix: "/",
		data:   data,
	})
}

// Traverses the tree and calls the given function on every node.
func (node *TreeNode[T]) each(callback func(*TreeNode[T])) {
	callback(node)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.each(callback)
	}

	if node.parameter != nil {
		node.parameter.each(callback)
	}

	if node.wildcard != nil {
		node.wildcard.each(callback)
	}
}

// Called when the node was fully parsed and needs to decide the next control flow.
// finish is only called from `tree.Add`.
func (node *TreeNode[T]) finish(path string, data T, i int, offset int) (*TreeNode[T], int, Flow) {
	char := path[i]

	if char >= node.start && char < node.end {
		index := node.indexes[char-node.start]

		if index != 0 {
			node = node.children[index]
			offset = i
			return node, offset, flowNext
		}
	}

	// No fitting children found, does this node even contain a prefix yet?
	// If no prefix is set, this is the starting node.
	if node.prefix == "" {
		node.append(path[i:], data)
		return node, offset, flowStop
	}

	// node: /user/|:id
	// path: /user/|:id/profile
	if node.parameter != nil && path[i] == parameter {
		node = node.parameter
		offset = i
		return node, offset, flowBegin
	}

	node.append(path[i:], data)
	return node, offset, flowStop
}