github.com/kelleygo/clashcore@v1.0.2/component/trie/node.go (about)

     1  package trie
     2  
     3  import "strings"
     4  
     5  // Node is the trie's node
     6  type Node[T any] struct {
     7  	childMap  map[string]*Node[T]
     8  	childNode *Node[T] // optimize for only one child
     9  	childStr  string
    10  	inited    bool
    11  	data      T
    12  }
    13  
    14  func (n *Node[T]) getChild(s string) *Node[T] {
    15  	if n.childMap == nil {
    16  		if n.childNode != nil && n.childStr == s {
    17  			return n.childNode
    18  		}
    19  		return nil
    20  	}
    21  	return n.childMap[s]
    22  }
    23  
    24  func (n *Node[T]) hasChild(s string) bool {
    25  	return n.getChild(s) != nil
    26  }
    27  
    28  func (n *Node[T]) addChild(s string, child *Node[T]) {
    29  	if n.childMap == nil {
    30  		if n.childNode == nil {
    31  			n.childStr = s
    32  			n.childNode = child
    33  			return
    34  		}
    35  		n.childMap = map[string]*Node[T]{}
    36  		if n.childNode != nil {
    37  			n.childMap[n.childStr] = n.childNode
    38  		}
    39  		n.childStr = ""
    40  		n.childNode = nil
    41  	}
    42  
    43  	n.childMap[s] = child
    44  }
    45  
    46  func (n *Node[T]) getOrNewChild(s string) *Node[T] {
    47  	node := n.getChild(s)
    48  	if node == nil {
    49  		node = newNode[T]()
    50  		n.addChild(s, node)
    51  	}
    52  	return node
    53  }
    54  
    55  func (n *Node[T]) optimize() {
    56  	if len(n.childStr) > 0 {
    57  		n.childStr = strClone(n.childStr)
    58  	}
    59  	if n.childNode != nil {
    60  		n.childNode.optimize()
    61  	}
    62  	if n.childMap == nil {
    63  		return
    64  	}
    65  	switch len(n.childMap) {
    66  	case 0:
    67  		n.childMap = nil
    68  		return
    69  	case 1:
    70  		for key := range n.childMap {
    71  			n.childStr = key
    72  			n.childNode = n.childMap[key]
    73  		}
    74  		n.childMap = nil
    75  		n.optimize()
    76  		return
    77  	}
    78  	children := make(map[string]*Node[T], len(n.childMap)) // avoid map reallocate memory
    79  	for key := range n.childMap {
    80  		child := n.childMap[key]
    81  		if child == nil {
    82  			continue
    83  		}
    84  		key = strClone(key)
    85  		children[key] = child
    86  		child.optimize()
    87  	}
    88  	n.childMap = children
    89  }
    90  
    91  func strClone(key string) string {
    92  	switch key { // try to save string's memory
    93  	case wildcard:
    94  		key = wildcard
    95  	case dotWildcard:
    96  		key = dotWildcard
    97  	case complexWildcard:
    98  		key = complexWildcard
    99  	case domainStep:
   100  		key = domainStep
   101  	default:
   102  		key = strings.Clone(key)
   103  	}
   104  	return key
   105  }
   106  
   107  func (n *Node[T]) isEmpty() bool {
   108  	if n == nil || n.inited == false {
   109  		return true
   110  	}
   111  	return false
   112  }
   113  
   114  func (n *Node[T]) setData(data T) {
   115  	n.data = data
   116  	n.inited = true
   117  }
   118  
   119  func (n *Node[T]) getChildren() map[string]*Node[T] {
   120  	if n.childMap == nil {
   121  		if n.childNode != nil {
   122  			m := make(map[string]*Node[T])
   123  			m[n.childStr] = n.childNode
   124  			return m
   125  		}
   126  	} else {
   127  		return n.childMap
   128  	}
   129  	return nil
   130  }
   131  func (n *Node[T]) Data() T {
   132  	return n.data
   133  }
   134  
   135  func newNode[T any]() *Node[T] {
   136  	return &Node[T]{}
   137  }