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

     1  package cidr
     2  
     3  import (
     4  	"math"
     5  	"net"
     6  )
     7  
     8  type Trie[T any] struct {
     9  	last  bool
    10  	mark  T
    11  	left  *Trie[T] // 0
    12  	right *Trie[T] // 1
    13  }
    14  
    15  // Insert insert node to tree
    16  func (t *Trie[T]) Insert(ip net.IP, maskSize int, mark T) {
    17  	r := t
    18  	for i := range ip {
    19  		for b := byte(128); b != 0; b = b >> 1 {
    20  			if ip[i]&b != 0 {
    21  				if r.right == nil {
    22  					r.right = new(Trie[T])
    23  				}
    24  				r = r.right
    25  			} else {
    26  				if r.left == nil {
    27  					r.left = new(Trie[T])
    28  				}
    29  				r = r.left
    30  			}
    31  
    32  			if i*8+int(math.Log2(float64(128/b)))+1 == maskSize {
    33  				r.mark = mark
    34  				r.last = true
    35  				r.left = new(Trie[T])
    36  				r.right = new(Trie[T])
    37  				return
    38  			}
    39  		}
    40  	}
    41  }
    42  
    43  // Search search from trie tree
    44  func (t *Trie[T]) Search(ip net.IP) (mark T, ok bool) {
    45  	r := t
    46  out:
    47  	for i := range ip {
    48  		for b := byte(128); b != 0; b = b >> 1 {
    49  			if ip[i]&b != 0 { // bit = 1
    50  				r = r.right
    51  			} else { // bit = 0
    52  				r = r.left
    53  			}
    54  			if r == nil {
    55  				break out
    56  			}
    57  			if r.last {
    58  				mark, ok = r.mark, true
    59  			}
    60  		}
    61  	}
    62  	return
    63  }
    64  
    65  // Search search from trie tree
    66  func (t *Trie[T]) Remove(ip net.IP) {
    67  	//TODO
    68  }
    69  
    70  // NewTrieTree create a new trie tree
    71  func NewTrieTree[T any]() Trie[T] { return Trie[T]{} }