ahocorasick.gno
8.66 Kb · 276 lines
1// Package ahocorasick implements Aho-Corasick multi-pattern search as a pure,
2// reusable package.
3//
4// Searching for k patterns by running a single-pattern scan k times costs
5// O(k*n): every extra word in the dictionary is another full pass over the
6// text. Aho-Corasick builds one automaton from the whole set, a trie of the
7// patterns plus a failure link on every node pointing at the longest proper
8// suffix that is also a prefix of some pattern, and then reads the text once.
9// Building the trie and the failure links costs O(total pattern length) and
10// scanning costs O(n + matches), whatever k is, and the text cursor never moves
11// backwards.
12//
13// One deliberate departure from the paper: each node stores the full list of
14// patterns matching there, merged from its failure target at build time, rather
15// than a link to be walked at every offset. That makes a reported match O(1)
16// instead of a walk up the failure chain, and it costs build memory in exactly
17// one shape, a set whose patterns nest ("a", "aa", "aaa", ...), where the
18// merged lists are quadratic in k rather than linear. MaxPatterns is what
19// bounds that, and nothing else does.
20//
21// That last property is why it is worth having on chain: the worst case equals
22// the best case, so a dictionary and a text chosen by an attacker cost the same
23// as a friendly pair. A naive k-pass scan is an input away from quadratic.
24//
25// It is [p/moul/x/daily/kmp](/p/moul/x/daily/kmp/v0) generalised: KMP's failure
26// table is this automaton for a one-pattern set.
27//
28// Operates on BYTES, not runes. gno strings are UTF-8 and a match offset is a
29// byte offset, which is the right unit for slicing. UTF-8 is self-synchronizing,
30// so a pattern that is itself valid UTF-8 can only match on a rune boundary of
31// valid UTF-8 text; a pattern that is a bare continuation byte can and will
32// match inside a rune, and slicing on that offset splits it.
33//
34// A live demo of this package is at
35// [r/moul/x/daily/ahocorasickdemo](/r/moul/x/daily/ahocorasickdemo/v0).
36package ahocorasick
37
38// Bounds on what New will build, so an on-chain caller cannot be handed an
39// automaton whose size it did not choose. At most MaxPatterns*MaxPatternLen+1
40// nodes exist, and the build is linear in the total pattern length.
41const (
42 MaxPatterns = 256
43 MaxPatternLen = 256
44)
45
46// Match is one occurrence of one pattern. Start and End are byte offsets into
47// the text, half-open, so text[Start:End] is the matched run. Index is the
48// pattern's slot in the slice given to New, which is how a caller maps a match
49// back to whatever it stored alongside the pattern.
50type Match struct {
51 Start int
52 End int
53 Index int
54 Pattern string
55}
56
57// edge is one labelled transition. A node keeps its edges sorted by label so a
58// lookup is a binary search over the alphabet actually used, not a 256-entry
59// row per node: a dictionary of words touches a few dozen distinct bytes and
60// paying 1 KiB a node for the rest is how this structure gets expensive.
61type edge struct {
62 label byte
63 next int
64}
65
66type node struct {
67 edges []edge
68 fail int
69 out []int // pattern slots matching here, longest first
70}
71
72// Matcher is a built automaton. It is immutable once New returns, so one
73// Matcher can serve any number of scans.
74type Matcher struct {
75 patterns []string
76 nodes []node
77}
78
79// New builds the automaton for patterns. Empty patterns are dropped: the empty
80// string matches at every offset, which is a different question from the one
81// this package answers.
82//
83// Returns nil when there is nothing to build (no pattern survives the drop) or
84// when the set exceeds the bounds above. A nil Matcher answers every method,
85// so a caller may hand the result straight to a scan without a check.
86func New(patterns []string) *Matcher {
87 accepted := make([]string, 0, len(patterns))
88 for _, p := range patterns {
89 if p == "" {
90 continue
91 }
92 if len(p) > MaxPatternLen {
93 return nil
94 }
95 accepted = append(accepted, p)
96 }
97 if len(accepted) == 0 || len(accepted) > MaxPatterns {
98 return nil
99 }
100
101 m := &Matcher{patterns: accepted, nodes: []node{{fail: 0}}}
102 for i, p := range accepted {
103 cur := 0
104 for j := 0; j < len(p); j++ {
105 next, ok := m.child(cur, p[j])
106 if !ok {
107 m.nodes = append(m.nodes, node{fail: 0})
108 next = len(m.nodes) - 1
109 m.addEdge(cur, p[j], next)
110 }
111 cur = next
112 }
113 // Same depth, so same length: ordering by slot keeps a duplicated
114 // pattern reporting in the caller's order.
115 m.nodes[cur].out = append(m.nodes[cur].out, i)
116 }
117 m.link()
118 return m
119}
120
121// link computes every failure link, then merges each node's outputs with its
122// failure target's. The merge is what makes a scan O(1) per reported match
123// instead of a walk up the failure chain at every offset, and the breadth-first
124// order is what makes it correct: a node's failure target is always shallower,
125// so it has already been merged.
126//
127// It is also where the memory goes: a nested set makes the merged lists
128// quadratic in the number of patterns. See the package doc.
129func (m *Matcher) link() {
130 queue := make([]int, 0, len(m.nodes))
131 for _, e := range m.nodes[0].edges {
132 m.nodes[e.next].fail = 0
133 queue = append(queue, e.next)
134 }
135 for len(queue) > 0 {
136 u := queue[0]
137 queue = queue[1:]
138 for _, e := range m.nodes[u].edges {
139 v := e.next
140 f := m.nodes[u].fail
141 for {
142 if w, ok := m.child(f, e.label); ok {
143 m.nodes[v].fail = w
144 break
145 }
146 if f == 0 {
147 m.nodes[v].fail = 0
148 break
149 }
150 f = m.nodes[f].fail
151 }
152 inherited := m.nodes[m.nodes[v].fail].out
153 if len(inherited) > 0 {
154 merged := make([]int, 0, len(m.nodes[v].out)+len(inherited))
155 merged = append(merged, m.nodes[v].out...)
156 merged = append(merged, inherited...)
157 m.nodes[v].out = merged
158 }
159 queue = append(queue, v)
160 }
161 }
162}
163
164// child returns the node reached from n by label, if there is one.
165func (m *Matcher) child(n int, label byte) (int, bool) {
166 edges := m.nodes[n].edges
167 lo, hi := 0, len(edges)
168 for lo < hi {
169 mid := (lo + hi) / 2
170 if edges[mid].label < label {
171 lo = mid + 1
172 } else {
173 hi = mid
174 }
175 }
176 if lo < len(edges) && edges[lo].label == label {
177 return edges[lo].next, true
178 }
179 return 0, false
180}
181
182// addEdge inserts a transition, keeping the edge list sorted by label.
183func (m *Matcher) addEdge(n int, label byte, next int) {
184 edges := m.nodes[n].edges
185 pos := len(edges)
186 for i, e := range edges {
187 if e.label > label {
188 pos = i
189 break
190 }
191 }
192 edges = append(edges, edge{})
193 copy(edges[pos+1:], edges[pos:])
194 edges[pos] = edge{label: label, next: next}
195 m.nodes[n].edges = edges
196}
197
198// step advances the automaton one byte, following failure links until the byte
199// is consumed or the root refuses it. This is the whole scan: the text index
200// only ever moves forward.
201func (m *Matcher) step(cur int, c byte) int {
202 for {
203 if w, ok := m.child(cur, c); ok {
204 return w
205 }
206 if cur == 0 {
207 return 0
208 }
209 cur = m.nodes[cur].fail
210 }
211}
212
213// FindAll returns every occurrence of every pattern, including overlapping and
214// nested ones, ordered by end offset. At one end offset the longest match comes
215// first, because the scan reports the node it reached before the shorter
216// suffixes hanging off its failure chain, so a caller taking the first match at
217// each offset gets the greedy one.
218func (m *Matcher) FindAll(text string) []Match {
219 if m == nil {
220 return nil
221 }
222 var out []Match
223 cur := 0
224 for i := 0; i < len(text); i++ {
225 cur = m.step(cur, text[i])
226 for _, slot := range m.nodes[cur].out {
227 p := m.patterns[slot]
228 out = append(out, Match{Start: i + 1 - len(p), End: i + 1, Index: slot, Pattern: p})
229 }
230 }
231 return out
232}
233
234// FindFirst returns the first match FindAll would return, and stops the scan
235// there. The one to reach for when the question is "does this text contain
236// anything from the list, and where".
237func (m *Matcher) FindFirst(text string) (Match, bool) {
238 if m == nil {
239 return Match{}, false
240 }
241 cur := 0
242 for i := 0; i < len(text); i++ {
243 cur = m.step(cur, text[i])
244 if out := m.nodes[cur].out; len(out) > 0 {
245 p := m.patterns[out[0]]
246 return Match{Start: i + 1 - len(p), End: i + 1, Index: out[0], Pattern: p}, true
247 }
248 }
249 return Match{}, false
250}
251
252// Contains reports whether any pattern occurs in text.
253func (m *Matcher) Contains(text string) bool {
254 _, ok := m.FindFirst(text)
255 return ok
256}
257
258// Patterns returns the accepted patterns, in the caller's order, as a copy.
259func (m *Matcher) Patterns() []string {
260 if m == nil {
261 return nil
262 }
263 out := make([]string, len(m.patterns))
264 copy(out, m.patterns)
265 return out
266}
267
268// Nodes returns the size of the automaton: one node per distinct pattern
269// prefix, plus the root. The number a caller sizing a dictionary wants, since
270// it is what the build actually costs.
271func (m *Matcher) Nodes() int {
272 if m == nil {
273 return 0
274 }
275 return len(m.nodes)
276}
