github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/trie/domain/trie.go (about)

     1  package domain
     2  
     3  var (
     4  	_        uint8 = 0
     5  	last     uint8 = 1
     6  	wildcard uint8 = 2
     7  )
     8  
     9  type trie[T any] struct {
    10  	Value  T                   `json:"value"`
    11  	Symbol uint8               `json:"symbol"`
    12  	Child  map[string]*trie[T] `json:"child"`
    13  }
    14  
    15  func (d *trie[T]) child(s string, insert bool) (*trie[T], bool) {
    16  	if insert {
    17  		if d.Child == nil {
    18  			d.Child = make(map[string]*trie[T])
    19  		}
    20  
    21  		if d.Child[s] == nil {
    22  			d.Child[s] = &trie[T]{}
    23  		}
    24  	} else {
    25  		if d.Child == nil {
    26  			return nil, false
    27  		}
    28  	}
    29  
    30  	r, ok := d.Child[s]
    31  	return r, ok
    32  }
    33  
    34  func search[T any](root *trie[T], domain *fqdnReader) (t T, ok bool) {
    35  
    36  	first, asterisk := true, false
    37  
    38  	for domain.hasNext() {
    39  		r, cok := root.child(domain.str(), false)
    40  		switch cok {
    41  		case false:
    42  			if !first {
    43  				return
    44  			}
    45  
    46  			if asterisk {
    47  				domain.next()
    48  				continue
    49  			}
    50  
    51  			root, cok = root.child("*", false)
    52  			if !cok {
    53  				return
    54  			}
    55  
    56  			asterisk = true
    57  
    58  		case true:
    59  			switch r.Symbol {
    60  			case wildcard:
    61  				t, ok = r.Value, true
    62  			case last:
    63  				if domain.last() {
    64  					return r.Value, true
    65  				}
    66  			}
    67  
    68  			root = r
    69  			domain.next()
    70  			first = false
    71  		}
    72  	}
    73  
    74  	return
    75  }
    76  
    77  func insert[T any](node *trie[T], z *fqdnReader, mark T) {
    78  	for z.hasNext() {
    79  		if z.last() && z.str() == "*" {
    80  			node.Symbol, node.Value = wildcard, mark
    81  			break
    82  		}
    83  
    84  		node, _ = node.child(z.str(), true)
    85  
    86  		if z.last() {
    87  			node.Symbol, node.Value = last, mark
    88  		}
    89  
    90  		z.next()
    91  	}
    92  }
    93  
    94  type deleteElement[T any] struct {
    95  	str  string
    96  	node *trie[T]
    97  }
    98  
    99  func remove[T any](node *trie[T], domain *fqdnReader) {
   100  	// fmt.Println("remove", domain.domain)
   101  
   102  	nodes := []*deleteElement[T]{
   103  		{
   104  			str:  "root",
   105  			node: node,
   106  		},
   107  	}
   108  
   109  	for domain.hasNext() {
   110  		z, ok := node.child(domain.str(), false)
   111  		if !ok {
   112  			if domain.str() == "*" && node.Symbol == wildcard {
   113  				break
   114  			}
   115  			return
   116  		}
   117  
   118  		node = z
   119  		nodes = append(nodes, &deleteElement[T]{domain.str(), node})
   120  		domain.next()
   121  	}
   122  
   123  	if node.Symbol == 0 {
   124  		return
   125  	}
   126  
   127  	for i := len(nodes) - 1; i >= 0; i-- {
   128  		if len(nodes[i].node.Child) != 0 {
   129  			if i == len(nodes)-1 {
   130  				nodes[i].node.Symbol = 0
   131  			}
   132  			break
   133  		}
   134  
   135  		if len(nodes[i].node.Child) == 0 {
   136  			if i-1 > 0 {
   137  				delete(nodes[i-1].node.Child, nodes[i].str)
   138  			}
   139  		}
   140  	}
   141  
   142  	// for _, v := range nodes {
   143  	// 	fmt.Println(v.str, len(v.node.Child), maps.Keys(v.node.Child))
   144  	// }
   145  }