chore: optimized initialization

component/trie/domain.go

@@ -25,7 +25,7 @@ func ValidAndSplitDomain(domain string) ([]string, bool) {
 	if domain != "" && domain[len(domain)-1] == '.' {
 		return nil, false
 	}
-
+	domain=strings.ToLower(domain)
 	parts := strings.Split(domain, domainStep)
 	if len(parts) == 1 {
 		if parts[0] == "" {
@@ -123,6 +123,30 @@ func (t *DomainTrie[T]) 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.
 func New[T any]() *DomainTrie[T] {
 	return &DomainTrie[T]{root: newNode[T]()}

component/trie/domain_test.go

@@ -105,3 +105,23 @@ func TestTrie_WildcardBoundary(t *testing.T) {
 
 	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)
+}

component/trie/node.go

@@ -116,6 +116,18 @@ func (n *Node[T]) setData(data T) {
 	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 {
 	return n.data
 }

component/trie/sskv.go

@@ -25,38 +25,14 @@ type DomainSet struct {
 
 // NewDomainSet creates a new *DomainSet struct, from a slice of sorted strings.
 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))
-	insert := func(domain string) {
-		reserveDomain := 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)
-	}
+	domainTrie.Foreach(func(domain string, data struct{}) {
+		reserveDomains = append(reserveDomains, utils.Reverse(domain))
+	})
 	sort.Slice(reserveDomains, func(i, j int) bool {
 		return len(reserveDomains[i]) < len(reserveDomains[j])
 	})