github.com/jxskiss/gopkg/v2@v2.14.9-0.20240514120614-899f3e7952b4/zlog/radix_tree.go (about)

     1  package zlog
     2  
     3  import (
     4  	"fmt"
     5  	"sort"
     6  	"strings"
     7  )
     8  
     9  type radixTree[T any] struct {
    10  	root radixNode[T]
    11  }
    12  
    13  func (p *radixTree[T]) search(name string) (T, bool) {
    14  	value, found := p.root.search(name)
    15  	return value, found
    16  }
    17  
    18  type radixNode[T any] struct {
    19  	prefix     string
    20  	value      *T
    21  	childNodes childNodes[T]
    22  }
    23  
    24  type childNodes[T any] struct {
    25  	labels []byte
    26  	nodes  []*radixNode[T]
    27  }
    28  
    29  func (nodes *childNodes[_]) Len() int { return len(nodes.labels) }
    30  
    31  func (nodes *childNodes[_]) Less(i, j int) bool { return nodes.labels[i] < nodes.labels[j] }
    32  
    33  func (nodes *childNodes[_]) Swap(i, j int) {
    34  	nodes.labels[i], nodes.labels[j] = nodes.labels[j], nodes.labels[i]
    35  	nodes.nodes[i], nodes.nodes[j] = nodes.nodes[j], nodes.nodes[i]
    36  }
    37  
    38  func (n *radixNode[T]) getValue() (T, bool) {
    39  	if n.value == nil {
    40  		var zero T
    41  		return zero, false
    42  	}
    43  	return *n.value, true
    44  }
    45  
    46  func (n *radixNode[T]) insert(name string, value T) {
    47  	if name == "" {
    48  		n.value = &value
    49  		return
    50  	}
    51  
    52  	firstChar := name[0]
    53  	for i, label := range n.childNodes.labels {
    54  		if firstChar == label {
    55  			// Split based on the common prefix of the existing node and the new one.
    56  			child, prefixSplit := n.splitCommonPrefix(i, name)
    57  			child.insert(name[prefixSplit:], value)
    58  			return
    59  		}
    60  	}
    61  
    62  	// No existing node starting with this letter, so create it.
    63  	child := &radixNode[T]{prefix: name, value: &value}
    64  	n.childNodes.labels = append(n.childNodes.labels, firstChar)
    65  	n.childNodes.nodes = append(n.childNodes.nodes, child)
    66  	sort.Sort(&n.childNodes)
    67  	return
    68  }
    69  
    70  func (n *radixNode[T]) splitCommonPrefix(existingChildIndex int, name string) (*radixNode[T], int) {
    71  	child := n.childNodes.nodes[existingChildIndex]
    72  
    73  	if strings.HasPrefix(name, child.prefix) {
    74  		// No split needs to be done. Rather, the new name shares the entire
    75  		// prefix with the existing node, so the new node is just a child of
    76  		// the existing one. Or the new name is the same as the existing name,
    77  		// which means that we just move on to the next token.
    78  		// Either way, this return accomplishes that.
    79  		return child, len(child.prefix)
    80  	}
    81  
    82  	i := longestPrefix(name, child.prefix)
    83  	commonPrefix := name[:i]
    84  	child.prefix = child.prefix[i:]
    85  
    86  	// Create a new intermediary node in the place of the existing node, with
    87  	// the existing node as a child.
    88  	newNode := &radixNode[T]{
    89  		prefix: commonPrefix,
    90  		childNodes: childNodes[T]{
    91  			labels: []byte{child.prefix[0]},
    92  			nodes:  []*radixNode[T]{child},
    93  		},
    94  	}
    95  	n.childNodes.nodes[existingChildIndex] = newNode
    96  	return newNode, i
    97  }
    98  
    99  func (n *radixNode[T]) search(name string) (value T, found bool) {
   100  	node := n
   101  	for {
   102  		nameLen := len(name)
   103  		if nameLen == 0 {
   104  			break
   105  		}
   106  		firstChar := name[0]
   107  		if child := node.getChild(firstChar); child != nil {
   108  			childPrefixLen := len(child.prefix)
   109  			if nameLen >= childPrefixLen && child.prefix == name[:childPrefixLen] {
   110  				name = name[len(child.prefix):]
   111  				node = child
   112  				continue
   113  			}
   114  		}
   115  		break
   116  	}
   117  	return node.getValue()
   118  }
   119  
   120  func (n *radixNode[T]) getChild(label byte) *radixNode[T] {
   121  	num := len(n.childNodes.labels)
   122  	i := sort.Search(num, func(i int) bool {
   123  		return n.childNodes.labels[i] >= label
   124  	})
   125  	if i < num && n.childNodes.labels[i] == label {
   126  		return n.childNodes.nodes[i]
   127  	}
   128  	return nil
   129  }
   130  
   131  func (n *radixNode[T]) dumpTree(prefix string) string {
   132  	var out string
   133  	if n.value != nil {
   134  		out += fmt.Sprintf("%s=%v\n", prefix+n.prefix, n.value)
   135  	}
   136  	for _, child := range n.childNodes.nodes {
   137  		out += child.dumpTree(prefix + n.prefix)
   138  	}
   139  	return out
   140  }
   141  
   142  // longestPrefix finds the length of the shared prefix of two strings.
   143  func longestPrefix(k1, k2 string) int {
   144  	max := len(k1)
   145  	if l := len(k2); l < max {
   146  		max = l
   147  	}
   148  	var i int
   149  	for i = 0; i < max; i++ {
   150  		if k1[i] != k2[i] {
   151  			break
   152  		}
   153  	}
   154  	return i
   155  }