spark-server/internal/pkg/sensitive/matcher.go
2026-09-29 16:37:36 +08:00

112 lines
2.3 KiB
Go

package sensitive
// Entry is one enabled dictionary pattern (normalized word).
type Entry struct {
ID int64
Word string
RuneLen int
}
// Match is a hit on the normalized stream [Start, End) rune indices.
type Match struct {
EntryID int64
Start int
End int
}
// Matcher is an Aho-Corasick automaton over normalized patterns.
type Matcher struct {
entries []Entry
root *acNode
}
type acNode struct {
next map[rune]*acNode
fail *acNode
outputs []int // indices into entries
}
// Build constructs a matcher from patterns. Empty patterns are skipped.
func Build(entries []Entry) *Matcher {
m := &Matcher{entries: make([]Entry, 0, len(entries)), root: &acNode{next: map[rune]*acNode{}}}
for i := range entries {
e := entries[i]
w := e.Word
if w == "" {
continue
}
idx := len(m.entries)
rl := len([]rune(w))
e.RuneLen = rl
m.entries = append(m.entries, e)
m.addPattern(w, rl, idx)
}
m.buildFails()
return m
}
func (m *Matcher) addPattern(word string, _ int, outIdx int) {
node := m.root
for _, r := range word {
child := node.next[r]
if child == nil {
child = &acNode{next: map[rune]*acNode{}}
node.next[r] = child
}
node = child
}
node.outputs = append(node.outputs, outIdx)
}
func (m *Matcher) buildFails() {
queue := make([]*acNode, 0, 16)
for _, child := range m.root.next {
child.fail = m.root
queue = append(queue, child)
}
for qi := 0; qi < len(queue); qi++ {
cur := queue[qi]
for r, child := range cur.next {
queue = append(queue, child)
f := cur.fail
for f != nil && f.next[r] == nil {
f = f.fail
}
if f == nil {
child.fail = m.root
} else {
child.fail = f.next[r]
}
child.outputs = append(child.outputs, child.fail.outputs...)
}
}
}
// FindAll returns every match in norm (inclusive start, exclusive end).
func (m *Matcher) FindAll(norm []rune) []Match {
if m == nil || len(norm) == 0 {
return nil
}
var out []Match
node := m.root
for i, r := range norm {
for node != m.root && node.next[r] == nil {
node = node.fail
}
next := node.next[r]
if next == nil {
next = m.root
}
node = next
for _, idx := range node.outputs {
ent := m.entries[idx]
start := i + 1 - ent.RuneLen
if start < 0 {
start = 0
}
out = append(out, Match{EntryID: ent.ID, Start: start, End: i + 1})
}
}
return out
}