chore: optimized initialization
This commit is contained in:
parent
53b1250cd8
commit
a6ee3348df
@ -25,7 +25,7 @@ func ValidAndSplitDomain(domain string) ([]string, bool) {
|
|||||||
if domain != "" && domain[len(domain)-1] == '.' {
|
if domain != "" && domain[len(domain)-1] == '.' {
|
||||||
return nil, false
|
return nil, false
|
||||||
}
|
}
|
||||||
|
domain=strings.ToLower(domain)
|
||||||
parts := strings.Split(domain, domainStep)
|
parts := strings.Split(domain, domainStep)
|
||||||
if len(parts) == 1 {
|
if len(parts) == 1 {
|
||||||
if parts[0] == "" {
|
if parts[0] == "" {
|
||||||
@ -123,6 +123,30 @@ func (t *DomainTrie[T]) Optimize() {
|
|||||||
t.root.optimize()
|
t.root.optimize()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (t *DomainTrie[T]) Foreach(print func(domain string, data T)) {
|
||||||
|
for key, data := range t.root.getChildren() {
|
||||||
|
recursion([]string{key}, data, print)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func recursion[T any](items []string, node *Node[T], fn func(domain string, data T)) {
|
||||||
|
for key, data := range node.getChildren() {
|
||||||
|
newItems := append([]string{key}, items...)
|
||||||
|
if data != nil && data.inited {
|
||||||
|
domain := joinDomain(newItems)
|
||||||
|
if domain[0] == domainStepByte {
|
||||||
|
domain = complexWildcard + domain
|
||||||
|
}
|
||||||
|
fn(domain, data.Data())
|
||||||
|
}
|
||||||
|
recursion(newItems, data, fn)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func joinDomain(items []string) string {
|
||||||
|
return strings.Join(items, domainStep)
|
||||||
|
}
|
||||||
|
|
||||||
// New returns a new, empty Trie.
|
// New returns a new, empty Trie.
|
||||||
func New[T any]() *DomainTrie[T] {
|
func New[T any]() *DomainTrie[T] {
|
||||||
return &DomainTrie[T]{root: newNode[T]()}
|
return &DomainTrie[T]{root: newNode[T]()}
|
||||||
|
@ -105,3 +105,23 @@ func TestTrie_WildcardBoundary(t *testing.T) {
|
|||||||
|
|
||||||
assert.NotNil(t, tree.Search("example.com"))
|
assert.NotNil(t, tree.Search("example.com"))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestTrie_Foreach(t *testing.T) {
|
||||||
|
tree := New[netip.Addr]()
|
||||||
|
domainList := []string{
|
||||||
|
"google.com",
|
||||||
|
"stun.*.*.*",
|
||||||
|
"test.*.google.com",
|
||||||
|
"+.baidu.com",
|
||||||
|
"*.baidu.com",
|
||||||
|
"*.*.baidu.com",
|
||||||
|
}
|
||||||
|
for _, domain := range domainList {
|
||||||
|
tree.Insert(domain, localIP)
|
||||||
|
}
|
||||||
|
count := 0
|
||||||
|
tree.Foreach(func(domain string, data netip.Addr) {
|
||||||
|
count++
|
||||||
|
})
|
||||||
|
assert.Equal(t, 7, count)
|
||||||
|
}
|
||||||
|
@ -116,6 +116,18 @@ func (n *Node[T]) setData(data T) {
|
|||||||
n.inited = true
|
n.inited = true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (n *Node[T]) getChildren() map[string]*Node[T] {
|
||||||
|
if n.childMap == nil {
|
||||||
|
if n.childNode != nil {
|
||||||
|
m := make(map[string]*Node[T])
|
||||||
|
m[n.childStr] = n.childNode
|
||||||
|
return m
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
return n.childMap
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
func (n *Node[T]) Data() T {
|
func (n *Node[T]) Data() T {
|
||||||
return n.data
|
return n.data
|
||||||
}
|
}
|
||||||
|
@ -25,38 +25,14 @@ type DomainSet struct {
|
|||||||
|
|
||||||
// NewDomainSet creates a new *DomainSet struct, from a slice of sorted strings.
|
// NewDomainSet creates a new *DomainSet struct, from a slice of sorted strings.
|
||||||
func NewDomainSet(keys []string) *DomainSet {
|
func NewDomainSet(keys []string) *DomainSet {
|
||||||
filter := make(map[string]struct{}, len(keys))
|
domainTrie := New[struct{}]()
|
||||||
|
for _, domain := range keys {
|
||||||
|
domainTrie.Insert(domain, struct{}{})
|
||||||
|
}
|
||||||
reserveDomains := make([]string, 0, len(keys))
|
reserveDomains := make([]string, 0, len(keys))
|
||||||
insert := func(domain string) {
|
domainTrie.Foreach(func(domain string, data struct{}) {
|
||||||
reserveDomain := utils.Reverse(domain)
|
reserveDomains = append(reserveDomains, utils.Reverse(domain))
|
||||||
reserveDomain = strings.ToLower(reserveDomain)
|
})
|
||||||
if _, ok := filter[reserveDomain]; !ok {
|
|
||||||
filter[reserveDomain] = struct{}{}
|
|
||||||
domains := make([]string, 0, len(reserveDomains))
|
|
||||||
if strings.HasSuffix(reserveDomain, ".+") {
|
|
||||||
for _, domain := range reserveDomains {
|
|
||||||
if !strings.HasPrefix(domain, reserveDomain[0:len(reserveDomain)-2]) {
|
|
||||||
domains = append(domains, domain)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
reserveDomains = domains
|
|
||||||
}
|
|
||||||
reserveDomains = append(reserveDomains, reserveDomain)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
for _, key := range keys {
|
|
||||||
items, ok := ValidAndSplitDomain(key)
|
|
||||||
if !ok {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
if items[0] == complexWildcard {
|
|
||||||
domain := strings.Join(items[1:], domainStep)
|
|
||||||
insert(domain)
|
|
||||||
}
|
|
||||||
|
|
||||||
domain := strings.Join(items, domainStep)
|
|
||||||
insert(domain)
|
|
||||||
}
|
|
||||||
sort.Slice(reserveDomains, func(i, j int) bool {
|
sort.Slice(reserveDomains, func(i, j int) bool {
|
||||||
return len(reserveDomains[i]) < len(reserveDomains[j])
|
return len(reserveDomains[i]) < len(reserveDomains[j])
|
||||||
})
|
})
|
||||||
|
Reference in New Issue
Block a user