[Experimental] SIMD nonsense

README.md

@@ -67,6 +67,17 @@ builder.LoadStrings("strings.txt")
 Both functions expects a text file with one pattern per line. `LoadPatterns` expects the pattern to
 be in hexadecimal form.
 
+## SIMD (experimental, Go 1.26)
+
+This library can use Go's experimental `simd/archsimd` package to accelerate
+root-state skipping when the pattern set has 16 or fewer distinct starting
+bytes. This is only available on amd64 with AVX and requires building with Go
+1.26 and `GOEXPERIMENT=simd` enabled.
+
+Example:
+
+    GOEXPERIMENT=simd gotip test ./...
+
 ## Storing
 
 Use `Encode` to store a `Trie` in gzip compressed binary format:

builder.go

@@ -206,6 +206,8 @@ func (tb *TrieBuilder) Build() *Trie {
 		}
 	}
 
+	trie.initPrefilter()
+
 	return trie
 }
 

go.mod

@@ -1,3 +1,3 @@
 module github.com/BobuSumisu/aho-corasick
 
-go 1.23.1
+go 1.26

prefilter.go

@@ -0,0 +1,48 @@
+package ahocorasick
+
+// rootPrefilter accelerates scanning at the root state by skipping bytes that
+// cannot start any pattern. It tracks up to 16 candidate bytes for SIMD use.
+type rootPrefilter struct {
+	bytes  [16]byte     // Candidate bytes that transition away from root.
+	blocks [16][16]byte // SIMD broadcast blocks for each candidate byte.
+	count  int          // Number of candidates in bytes/blocks.
+	simd   bool         // Whether SIMD scanning is enabled for this trie.
+}
+
+func (p *rootPrefilter) init(rootTrans [256]uint32) {
+	p.count = 0
+	p.simd = false
+
+	for b := range 256 {
+		if rootTrans[b] != rootState {
+			if p.count == len(p.bytes) {
+				// Too many candidates for the SIMD prefilter; disable it.
+				p.count = 0
+				return
+			}
+			p.bytes[p.count] = byte(b)
+			p.count++
+		}
+	}
+
+	if p.count == 0 {
+		return
+	}
+
+	// Pre-broadcast each candidate byte for SIMD comparisons.
+	for i := 0; i < p.count; i++ {
+		for j := range 16 {
+			p.blocks[i][j] = p.bytes[i]
+		}
+	}
+
+	p.simd = p.enableSIMD()
+}
+
+func (tr *Trie) initPrefilter() {
+	if len(tr.failTrans) <= int(rootState) {
+		tr.prefilter = rootPrefilter{}
+		return
+	}
+	tr.prefilter.init(tr.failTrans[rootState])
+}

prefilter_nosimd.go

@@ -0,0 +1,11 @@
+//go:build !goexperiment.simd || !amd64
+
+package ahocorasick
+
+func (p *rootPrefilter) enableSIMD() bool {
+	return false
+}
+
+func (p *rootPrefilter) nextCandidateSIMD(_ []byte, start int) int {
+	return start
+}

prefilter_simd_amd64.go

@@ -0,0 +1,54 @@
+//go:build goexperiment.simd && amd64
+
+package ahocorasick
+
+import (
+	"math/bits"
+
+	"simd/archsimd"
+)
+
+func (p *rootPrefilter) enableSIMD() bool {
+	return p.count > 0 && archsimd.X86.AVX()
+}
+
+// nextCandidateSIMD returns the next position at or after start that could
+// transition from the root state, or len(input) if none are found.
+func (p *rootPrefilter) nextCandidateSIMD(input []byte, start int) int {
+	if p.count == 0 {
+		return len(input)
+	}
+
+	// Load the broadcasted candidates once per call.
+	var needles [16]archsimd.Uint8x16
+	for i := 0; i < p.count; i++ {
+		needles[i] = archsimd.LoadUint8x16(&p.blocks[i])
+	}
+
+	i := start
+	n := len(input)
+	for i+16 <= n {
+		hay := archsimd.LoadUint8x16Slice(input[i : i+16])
+		var mask uint16
+		for j := 0; j < p.count; j++ {
+			mask |= hay.Equal(needles[j]).ToBits()
+		}
+		if mask != 0 {
+			// First set bit is the earliest candidate in this block.
+			return i + bits.TrailingZeros16(mask)
+		}
+		i += 16
+	}
+
+	// Scalar tail for any remaining bytes.
+	for ; i < n; i++ {
+		b := input[i]
+		for j := 0; j < p.count; j++ {
+			if b == p.bytes[j] {
+				return i
+			}
+		}
+	}
+
+	return n
+}

stream.go

@@ -128,7 +128,7 @@ func (dec *decoder) decode() (*Trie, error) {
 		return nil, err
 	}
 
-	return &Trie{
+	trie := &Trie{
 		failTrans: failTrans,
 		dictLink:  dictLink,
 		dict:      dict,
@@ -139,5 +139,8 @@ func (dec *decoder) decode() (*Trie, error) {
 		matchStructPool: sync.Pool{
 			New: func() any { return new(Match) },
 		},
-	}, nil
+	}
+	trie.initPrefilter()
+
+	return trie, nil
 }

trie.go

@@ -19,6 +19,7 @@ type Trie struct {
 
 	matchPool       sync.Pool // Pool for match slice pointers
 	matchStructPool sync.Pool // Pool for Match structs
+	prefilter       rootPrefilter
 }
 
 // Walk calls this function on any match, giving the end position, length of the matched bytes,
@@ -33,11 +34,20 @@ func (tr *Trie) Walk(input []byte, fn WalkFn) {
 	dict := tr.dict
 	pattern := tr.pattern
 	dictLink := tr.dictLink
+	prefilter := &tr.prefilter
 
 	s := rootState
 
 	inputLen := len(input)
-	for i := range inputLen {
+	for i := 0; i < inputLen; {
+		if s == rootState && prefilter.simd {
+			next := prefilter.nextCandidateSIMD(input, i)
+			if next >= inputLen {
+				return
+			}
+			i = next
+		}
+
 		s = failTrans[s][input[i]]
 
 		ds := dict[s]
@@ -52,6 +62,8 @@ func (tr *Trie) Walk(input []byte, fn WalkFn) {
 				}
 			}
 		}
+
+		i++
 	}
 }