coredemo/framework/trie.go

package framework

import (
    "errors"
    "strings"
)

type Tree struct {
    root *node // 根节点
}

func NewTree() *Tree {
    return &Tree{root: newNode()}
}

// 增加路由节点
/*
/book/list
/book/:id (冲突)
/book/:id/name
/book/:student/age
/:user/name
/:user/name/:age(冲突)
*/
func (tree *Tree) AddRouter(uri string, handler ControllerHandler) error {
    n := tree.root
    // 确认路由是否冲突
    if n.matchNode(uri) != nil {
        return errors.New("route exists: " + uri)
    }

    segments := strings.Split(uri, "/")
    for index, segment := range segments {
        if !isWildSegment(segment) {
            segment = strings.ToUpper(segment)
        }
        isLast := index == len(segments)-1

        var objNode *node

        childNodes := n.filterChildNodes(segment)
        // 如果有匹配的子节点
        if len(childNodes) > 0 {
            // 如果有segment 相同的子节点，则选择这个子节点
            for _, cNode := range childNodes {
                if cNode.segment == segment {
                    objNode = cNode
                    break
                }
            }
        }

        if objNode == nil {
            // 创建一个当前 node 的节点
            cnode := newNode()
            cnode.segment = segment
            if isLast {
                cnode.isLast = true
                cnode.handler = handler
            }
            n.childes = append(n.childes, cnode)
            objNode = cnode
        }

        n = objNode
    }

    return nil
}

func (tree *Tree) FindHandler(uri string) ControllerHandler {
    matchNode := tree.root.matchNode(uri)
    if matchNode == nil {
        return nil
    }
    return matchNode.handler
}

// =====================================================================================================================

// 代表节点
type node struct {
    isLast  bool              // 代表这个节点是否可以成为最终的路由规则。该节点是否能成为一个独立的uri, 是否自身就是一个终极节点
    segment string            // uri 中的字符串，代表这个节点表示的路由中某个段的字符串
    handler ControllerHandler // 代表这个节点中包含的控制器，用于最终加载调用
    childes []*node           // 代表这个节点下的子节点
}

// ---------------------------------------------------------------------------------------------------------------------

func newNode() *node {
    return &node{
        isLast:  false,
        segment: "",
        handler: nil,
        childes: nil,
    }
}

// 判断路由是否在节点的所有子节点树中存在了
func (n *node) matchNode(uri string) *node {
    // 使用分隔符将uri 切割为两部分
    segments := strings.SplitN(uri, "/", 2)
    // 第一个部分用于匹配下一层子节点
    segment := segments[0]
    if !isWildSegment(segment) {
        segment = strings.ToUpper(segment)
    }

    // 匹配符合的下一层子节点
    cNodes := n.filterChildNodes(segment)
    if cNodes == nil || len(cNodes) == 0 {
        // 如果当前子节点没有一个符合，那么说明这个 uri 一定是之前不存在，直接返回 nil
        return nil
    }

    // 如果只有一个segment, 则是最后一个标记
    if len(segments) == 1 {
        // 如果segment 已经是最后一个节点，判断这些 cnode 是否有 isLast 标志
        for _, tn := range cNodes {
            if tn.isLast {
                return tn
            }
        }
        // 都不是最后一个节点
        return nil
    }

    // 如果有 2 个 segment, 递归每个子节点继续进行查找
    for _, tn := range cNodes {
        tnMatch := tn.matchNode(segments[1])
        if tnMatch != nil {
            return tnMatch
        }
    }

    return nil
}

// 过滤下一层满足 segment 规则的子节点
func (n *node) filterChildNodes(segment string) []*node {
    if len(n.childes) == 0 {
        return nil
    }

    // 如果 segment 是通配符，则所有下一层子节点都满足需求
    if isWildSegment(segment) {
        return n.childes
    }

    nodes := make([]*node, 0, len(n.childes))
    // 过滤所有的下一层子节点
    for _, cNode := range n.childes {
        if isWildSegment(cNode.segment) {
            // 如果下一层子节点有通配符，则满足需求
            nodes = append(nodes, cNode)
        } else if cNode.segment == segment {
            // 如果下一层子节点没有通配符，但文本完全匹配，则满足需求
            nodes = append(nodes, cNode)
        }
    }
    return nodes
}

// 判断一个 segment 是否是通用 segment， 即以 ：开头
func isWildSegment(segment string) bool {
    return strings.HasPrefix(segment, ":")
}