From 128277023b20b592f3512f81d289ad234801f0c6 Mon Sep 17 00:00:00 2001 From: AhravDutta Date: Wed, 21 Jan 2026 20:46:17 -0800 Subject: [PATCH 1/2] Add SIMD-based root state skipping Add support for Go 1.26's experimental `simd/archsimd` package to accelerate root-state skipping when the pattern set has 16 or fewer distinct starting bytes. This feature is only available on --- README.md | 11 +++++++++ builder.go | 2 ++ go.mod | 2 +- prefilter.go | 44 ++++++++++++++++++++++++++++++++++++ prefilter_nosimd.go | 11 +++++++++ prefilter_simd_amd64.go | 49 +++++++++++++++++++++++++++++++++++++++++ stream.go | 7 ++++-- trie.go | 14 +++++++++++- 8 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 prefilter.go create mode 100644 prefilter_nosimd.go create mode 100644 prefilter_simd_amd64.go diff --git a/README.md b/README.md index 1616d4b..d7ae731 100644 --- a/README.md +++ b/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: diff --git a/builder.go b/builder.go index 9c27a1e..bd8c5b4 100644 --- a/builder.go +++ b/builder.go @@ -206,6 +206,8 @@ func (tb *TrieBuilder) Build() *Trie { } } + trie.initPrefilter() + return trie } diff --git a/go.mod b/go.mod index 61e342f..441716e 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,3 @@ module github.com/BobuSumisu/aho-corasick -go 1.23.1 +go 1.26 diff --git a/prefilter.go b/prefilter.go new file mode 100644 index 0000000..6284d88 --- /dev/null +++ b/prefilter.go @@ -0,0 +1,44 @@ +package ahocorasick + +type rootPrefilter struct { + bytes [16]byte + blocks [16][16]byte + count int + simd bool +} + +func (p *rootPrefilter) init(rootTrans [256]uint32) { + p.count = 0 + p.simd = false + + for b := 0; b < 256; b++ { + if rootTrans[b] != rootState { + if p.count == len(p.bytes) { + p.count = 0 + return + } + p.bytes[p.count] = byte(b) + p.count++ + } + } + + if p.count == 0 { + return + } + + for i := 0; i < p.count; i++ { + for j := 0; j < 16; j++ { + 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]) +} diff --git a/prefilter_nosimd.go b/prefilter_nosimd.go new file mode 100644 index 0000000..b81b1fc --- /dev/null +++ b/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 +} diff --git a/prefilter_simd_amd64.go b/prefilter_simd_amd64.go new file mode 100644 index 0000000..364e05d --- /dev/null +++ b/prefilter_simd_amd64.go @@ -0,0 +1,49 @@ +//go:build goexperiment.simd && amd64 + +package ahocorasick + +import ( + "math/bits" + + "simd/archsimd" +) + +func (p *rootPrefilter) enableSIMD() bool { + return p.count > 0 && archsimd.X86.AVX() +} + +func (p *rootPrefilter) nextCandidateSIMD(input []byte, start int) int { + if p.count == 0 { + return len(input) + } + + 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 { + return i + bits.TrailingZeros16(mask) + } + i += 16 + } + + for ; i < n; i++ { + b := input[i] + for j := 0; j < p.count; j++ { + if b == p.bytes[j] { + return i + } + } + } + + return n +} diff --git a/stream.go b/stream.go index 59c0918..16918ed 100644 --- a/stream.go +++ b/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 } diff --git a/trie.go b/trie.go index 776ec0c..88356bf 100644 --- a/trie.go +++ b/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++ } } From d75ada24affbd5967fb41211e112d4fefc5cf5a3 Mon Sep 17 00:00:00 2001 From: ahrav Date: Wed, 21 Jan 2026 21:58:21 -0800 Subject: [PATCH 2/2] Cleanup --- prefilter.go | 16 ++++++++++------ prefilter_simd_amd64.go | 5 +++++ 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/prefilter.go b/prefilter.go index 6284d88..bc56fca 100644 --- a/prefilter.go +++ b/prefilter.go @@ -1,19 +1,22 @@ 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 - blocks [16][16]byte - count int - simd bool + 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 := 0; b < 256; b++ { + 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 } @@ -26,8 +29,9 @@ func (p *rootPrefilter) init(rootTrans [256]uint32) { return } + // Pre-broadcast each candidate byte for SIMD comparisons. for i := 0; i < p.count; i++ { - for j := 0; j < 16; j++ { + for j := range 16 { p.blocks[i][j] = p.bytes[i] } } diff --git a/prefilter_simd_amd64.go b/prefilter_simd_amd64.go index 364e05d..3baa9ca 100644 --- a/prefilter_simd_amd64.go +++ b/prefilter_simd_amd64.go @@ -12,11 +12,14 @@ 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]) @@ -31,11 +34,13 @@ func (p *rootPrefilter) nextCandidateSIMD(input []byte, start int) int { 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++ {