chore: avoid string reversal allocation in DomainSet.Has (#3049)

component/trie/domain_set.go

@@ -6,6 +6,7 @@ package trie
 import (
 	"sort"
 	"strings"
+	"unicode/utf8"
 
 	"github.com/metacubex/mihomo/common/utils"
 	"github.com/openacid/low/bitmap"
@@ -75,8 +76,15 @@ func (ss *DomainSet) Has(key string) bool {
 	if ss == nil {
 		return false
 	}
-	key = utils.Reverse(key)
-	key = strings.ToLower(key)
+	for i := 0; i < len(key); i++ {
+		if key[i] >= utf8.RuneSelf {
+			// The set is built with rune-wise reversal, which only matches
+			// byte-wise reversal for ASCII. Normalize the same way and
+			// byte-reverse so revLowerAt below observes it unchanged.
+			key = byteReverse(strings.ToLower(utils.Reverse(key)))
+			break
+		}
+	}
 	// no more labels in this node
 	// skip character matching
 	// go to next level
@@ -87,7 +95,7 @@ func (ss *DomainSet) Has(key string) bool {
 	stack := make([]wildcardCursor, 0)
 	for i := 0; i < len(key); i++ {
 	RESTART:
-		c := key[i]
+		c := revLowerAt(key, i)
 		for ; ; bmIdx++ {
 			if getBit(ss.labelBitmap, bmIdx) != 0 {
 				if len(stack) > 0 {
@@ -97,7 +105,7 @@ func (ss *DomainSet) Has(key string) bool {
 					nextNodeId := countZeros(ss.labelBitmap, ss.ranks, cursor.bmIdx+1)
 					nextBmIdx := selectIthOne(ss.labelBitmap, ss.ranks, ss.selects, nextNodeId-1) + 1
 					j := cursor.index
-					for ; j < len(key) && key[j] != domainStepByte; j++ {
+					for ; j < len(key) && revLowerAt(key, j) != domainStepByte; j++ {
 					}
 					if j == len(key) {
 						if getBit(ss.leaves, nextNodeId) != 0 {
@@ -137,6 +145,24 @@ func (ss *DomainSet) Has(key string) bool {
 
 }
 
+// revLowerAt returns the i-th byte of key read back to front, lowercased for
+// ASCII. It lets Has walk the reversed key without materializing it.
+func revLowerAt(key string, i int) byte {
+	c := key[len(key)-1-i]
+	if c >= 'A' && c <= 'Z' {
+		c += 'a' - 'A'
+	}
+	return c
+}
+
+func byteReverse(s string) string {
+	buf := make([]byte, len(s))
+	for i := 0; i < len(s); i++ {
+		buf[i] = s[len(s)-1-i]
+	}
+	return string(buf)
+}
+
 func (ss *DomainSet) keys(f func(key string) bool) {
 	var currentKey []byte
 	var traverse func(int, int) bool

component/trie/domain_set_test.go

@@ -1,9 +1,12 @@
 package trie_test
 
 import (
-	"golang.org/x/exp/slices"
+	"strconv"
+	"strings"
 	"testing"
 
+	"golang.org/x/exp/slices"
+
 	"github.com/metacubex/mihomo/component/trie"
 	"github.com/stretchr/testify/assert"
 )
@@ -106,3 +109,82 @@ func TestDomainSetWildcard(t *testing.T) {
 	assert.False(t, set.Has("test.test.test.qq.com"))
 	testDump(t, tree, set)
 }
+
+func TestDomainSetCase(t *testing.T) {
+	tree := trie.New[struct{}]()
+	for _, domain := range []string{"example.com", "+.mixed.example.org"} {
+		assert.NoError(t, tree.Insert(domain, struct{}{}))
+	}
+	set := tree.NewDomainSet()
+	assert.NotNil(t, set)
+	assert.True(t, set.Has("EXAMPLE.COM"))
+	assert.True(t, set.Has("ExAmPlE.cOm"))
+	assert.True(t, set.Has("WWW.MIXED.EXAMPLE.ORG"))
+	assert.False(t, set.Has("EXAMPLE.NET"))
+}
+
+// TestDomainSetUnicode covers keys that are not ASCII, which take a different
+// path than the byte-wise one because the set is built with rune-wise reversal.
+func TestDomainSetUnicode(t *testing.T) {
+	tree := trie.New[struct{}]()
+	for _, domain := range []string{"中文.example", "+.测试.cn"} {
+		assert.NoError(t, tree.Insert(domain, struct{}{}))
+	}
+	set := tree.NewDomainSet()
+	assert.NotNil(t, set)
+	assert.True(t, set.Has("中文.example"))
+	assert.True(t, set.Has("www.测试.cn"))
+	assert.False(t, set.Has("日本語.example"))
+}
+
+func TestDomainSetOversizedKey(t *testing.T) {
+	tree := trie.New[struct{}]()
+	assert.NoError(t, tree.Insert("+.example.com", struct{}{}))
+	set := tree.NewDomainSet()
+	assert.NotNil(t, set)
+
+	var builder strings.Builder
+	for builder.Len() < 300 {
+		builder.WriteString("label.")
+	}
+	assert.True(t, set.Has(builder.String()+"example.com"))
+	assert.False(t, set.Has(builder.String()+"example.net"))
+}
+
+func BenchmarkDomainSetHas(b *testing.B) {
+	tree := trie.New[struct{}]()
+	for i := 0; i < 10000; i++ {
+		assert.NoError(b, tree.Insert("+."+strconv.Itoa(i)+".example.com", struct{}{}))
+	}
+	set := tree.NewDomainSet()
+
+	// Keys are split by length because the Go compiler only keeps a
+	// non-constant sized allocation off the heap up to 32 bytes, so hostnames
+	// longer than that used to take a different path.
+	benchmarks := []struct {
+		name string
+		keys []string
+	}{
+		{"short", []string{
+			"www.4242.example.com",
+			"a.b.c.9999.example.com",
+			"no-such-host.example.net",
+			"WWW.1234.EXAMPLE.COM",
+		}},
+		{"long", []string{
+			"ec2-52-201-13-44.compute-1.4242.example.com",
+			"prod-eu-west-1-api.metrics.9999.example.com",
+			"abcdef123456.dualstack.us-east-1.elb.example.net",
+			"EC2-52-201-13-44.COMPUTE-1.1234.EXAMPLE.COM",
+		}},
+	}
+	for _, benchmark := range benchmarks {
+		b.Run(benchmark.name, func(b *testing.B) {
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				set.Has(benchmark.keys[i%len(benchmark.keys)])
+			}
+		})
+	}
+}