perf(std text aho-corasick): 13x faster build for 1M+ state automata
ober
aed1665f1eb7100892dd018536477318d58101b4
--- a/lib/std/text/aho-corasick.ss +++ b/lib/std/text/aho-corasick.ss @@ -65,24 +65,40 @@ (def (make-ac patterns) (when (null? patterns) (error 'make-ac "ruleset must contain at least one pattern")) - (let-values ([(nodes num-states) (build-trie patterns)]) - (compute-fail-and-output! nodes num-states) - (let-values ([(goto output) (build-dense-tables nodes num-states)]) - (ac-internal-make num-states (length patterns) goto output)))) + (let ([prof? (getenv "AC_PROFILE")]) + (let ([t0 (real-time)]) + (let-values ([(root num-states) (build-trie patterns)]) + (when prof? + (printf " [ac] build-trie: ~a ms (~a states)~n" + (- (real-time) t0) num-states)) + (let ([t1 (real-time)]) + (let-values ([(goto output) (build-fail-and-dense! root num-states)]) + (when prof? + (printf " [ac] build-fail-and-dense: ~a ms~n" + (- (real-time) t1))) + (ac-internal-make num-states (length patterns) goto output))))))) - ;; A trie node is mutable during build only. Layout: #(goto-ht fail term). - ;; goto-ht: eqv-hashtable from byte (0..255) -> state-id (fixnum) - ;; fail: state-id of the failure link (0 == root) - ;; term: list of (pattern-id . pattern-len) + ;; A trie node is mutable during build only. Layout: + ;; #(goto-alist fail term id) + ;; goto-alist: list of (byte . child-node-ref) + ;; Set to #f after consumption during BFS to free GC. + ;; fail: state-id of the failure link (0 == root) + ;; term: list of (pattern-id . pattern-len) + ;; id: state-id, assigned in BFS order during phase 2 + ;; + ;; Per-node alist (instead of a Chez eqv-hashtable) is the big win: + ;; most trie nodes have 1-2 children, and an eqv-hashtable allocation + ;; (with its internal arrays) is hundreds of times more expensive than + ;; a cons cell. Across 1M+ nodes that dominates the trie phase. (def (new-node) - (vector (make-eqv-hashtable) 0 '())) + (vector '() 0 '() -1)) - ;; Phase 1: walk every pattern into the trie, return (values nodes count). - ;; nodes is a vector of trie-nodes indexed by state-id; root is at 0. + ;; Phase 1: walk every pattern into the trie. Returns (values root count). + ;; State ids are NOT assigned here -- that happens in BFS order so the + ;; fail-link / dense-table phase can rely on the ordering invariant. (def (build-trie patterns) - (let ([by-id (make-eqv-hashtable)] - [next-id (box 1)]) - (hashtable-set! by-id 0 (new-node)) + (let ([root (new-node)] + [count (box 1)]) (for-each (lambda (entry) (let ([id (car entry)] [bv (cdr entry)]) @@ -91,108 +107,104 @@ (let ([m (bytevector-length bv)]) (when (fxzero? m) (error 'make-ac "empty pattern not allowed" id)) - (let walk ([state 0] [i 0]) + (let walk ([node root] [i 0]) (cond [(fx= i m) - (let ([n (hashtable-ref by-id state #f)]) - (vector-set! n 2 - (cons (cons id m) (vector-ref n 2))))] + (vector-set! node 2 + (cons (cons id m) (vector-ref node 2)))] [else - (let* ([b (bytevector-u8-ref bv i)] - [n (hashtable-ref by-id state #f)] - [g (vector-ref n 0)] - [nx (hashtable-ref g b #f)]) - (cond - [nx (walk nx (fx+ i 1))] - [else - (let ([new-state (unbox next-id)]) - (set-box! next-id (fx+ new-state 1)) - (hashtable-set! by-id new-state (new-node)) - (hashtable-set! g b new-state) - (walk new-state (fx+ i 1)))]))]))))) + (let ([b (bytevector-u8-ref bv i)]) + (let scan ([alist (vector-ref node 0)]) + (cond + [(null? alist) + (let ([new-node (new-node)]) + (set-box! count (fx+ (unbox count) 1)) + (vector-set! node 0 + (cons (cons b new-node) (vector-ref node 0))) + (walk new-node (fx+ i 1)))] + [(fx= (car (car alist)) b) + (walk (cdr (car alist)) (fx+ i 1))] + [else (scan (cdr alist))])))]))))) patterns) - (let* ([n (unbox next-id)] - [v (make-vector n #f)]) - (let loop ([i 0]) - (when (fx< i n) - (vector-set! v i (hashtable-ref by-id i #f)) - (loop (fx+ i 1)))) - (values v n)))) + (values root (unbox count)))) - ;; Phase 2: BFS from root. For each non-root state t reached via byte - ;; b from parent s, fail(t) = goto*(fail(s), b) where goto* walks the - ;; failure chain. output(t) = output(t) ++ output(fail(t)). - (def (compute-fail-and-output! nodes num-states) - (let ([queue (make-simple-queue)]) - (let-values ([(_keys vals) (hashtable-entries - (vector-ref (vector-ref nodes 0) 0))]) - (let ([n (vector-length vals)]) - (let loop ([i 0]) - (when (fx< i n) - (queue-enqueue! queue (vector-ref vals i)) - (loop (fx+ i 1)))))) + ;; Phases 2+3 combined: BFS from root, assigning state ids, computing + ;; fail links, building the dense (state*256 + byte) -> next-state + ;; transition table in one pass. + ;; + ;; Key invariant: row(s) = row(fail(s)) overridden by s's explicit + ;; children. When state s is popped from the BFS queue, fail(s) has + ;; already been popped (strict depth order), so row(fail(s)) is + ;; finalized -- we copy it into row(s) (one fxvector-copy! of 256 + ;; cells), then override with s's explicit children's cells. + ;; + ;; This replaces the previous resolve-goto walk (per-cell hashtable + ;; lookups up the failure chain) with an equivalent count of fxvector + ;; reads/writes: 314M lookups -> ~600ms. + ;; + ;; State ids are assigned in BFS order (not build-trie order), which + ;; is what makes the depth-strict popping invariant hold. + (def (build-fail-and-dense! root num-states) + (let ([goto (make-fxvector (fx* num-states 256) 0)] + [out (make-vector num-states '())] + [queue (make-simple-queue)] + [next-id (box 1)]) + ;; Root: id 0, fail = 0 (default), out[0] = own term list. + (vector-set! root 3 0) + (vector-set! out 0 (vector-ref root 2)) + ;; Fill row(0) by walking root's alist: assign each child an id, + ;; set its fail = 0, set its out = own ++ out[root], enqueue. + ;; Row(child) gets initialized when the child is popped (= when + ;; row(0) is finalized). + (let pass1 ([al (vector-ref root 0)]) + (cond + [(null? al) (void)] + [else + (let* ([b (car (car al))] + [child (cdr (car al))] + [child-id (unbox next-id)]) + (set-box! next-id (fx+ child-id 1)) + (vector-set! child 3 child-id) + (vector-set! child 1 0) + (vector-set! out child-id + (append (vector-ref child 2) (vector-ref out 0))) + (fxvector-set! goto b child-id) + (queue-enqueue! queue child) + (pass1 (cdr al)))])) + (vector-set! root 0 #f) (let bfs () (unless (queue-empty? queue) - (let* ([s (queue-dequeue! queue)] - [node (vector-ref nodes s)] - [g (vector-ref node 0)]) - (let-values ([(keys vals) (hashtable-entries g)]) - (let ([k (vector-length keys)]) - (let children ([i 0]) - (when (fx< i k) - (let* ([b (vector-ref keys i)] - [t (vector-ref vals i)] - [t-fail (compute-fail s b t nodes)]) - (vector-set! (vector-ref nodes t) 1 t-fail) - (vector-set! (vector-ref nodes t) 2 - (append (vector-ref (vector-ref nodes t) 2) - (vector-ref (vector-ref nodes t-fail) 2))) - (queue-enqueue! queue t)) - (children (fx+ i 1))))))) - (bfs))))) - - ;; Walk the failure chain from parent s looking for a goto on byte b. - ;; Returns the resulting state, or 0 if we exhaust the chain at root. - ;; Guards against self-loop: never return t for fail(t). - (def (compute-fail s b t nodes) - (let walk ([f (vector-ref (vector-ref nodes s) 1)]) - (let* ([f-node (vector-ref nodes f)] - [f-goto (vector-ref f-node 0)] - [m (hashtable-ref f-goto b #f)]) - (cond - [(and m (not (fx= m t))) m] - [(fx= f 0) 0] - [else (walk (vector-ref f-node 1))])))) - - ;; Phase 3: collapse sparse trie + fail links into a dense (state, byte) - ;; -> next-state table. After this point the input trie isn't touched. - (def (build-dense-tables nodes num-states) - (let ([goto (make-fxvector (fx* num-states 256) 0)] - [out (make-vector num-states '())]) - (let states ([s 0]) - (when (fx< s num-states) - (let ([node (vector-ref nodes s)]) - (vector-set! out s (vector-ref node 2)) - (let bytes ([b 0]) - (when (fx< b 256) - (fxvector-set! goto (fx+ (fx* s 256) b) - (resolve-goto s b nodes)) - (bytes (fx+ b 1))))) - (states (fx+ s 1)))) + (let* ([s (queue-dequeue! queue)] + [s-id (vector-ref s 3)] + [s-fail (vector-ref s 1)] + [s-row (fx* s-id 256)]) + ;; row(s) := row(fail(s)). fail(s) was popped earlier + ;; (strict depth order) so its row is finalized. + (fxvector-copy! goto (fx* s-fail 256) goto s-row 256) + (let children ([al (vector-ref s 0)]) + (cond + [(null? al) (void)] + [else + (let* ([b (car (car al))] + [t (cdr (car al))] + ;; row(s)[b] is row(fail(s))[b] = goto*(fail(s),b) + ;; -- exactly fail(t). Read before overriding. + [t-fail (fxvector-ref goto (fx+ s-row b))] + [t-id (unbox next-id)]) + (set-box! next-id (fx+ t-id 1)) + (vector-set! t 3 t-id) + (vector-set! t 1 t-fail) + (vector-set! out t-id + (append (vector-ref t 2) + (vector-ref out t-fail))) + (fxvector-set! goto (fx+ s-row b) t-id) + (queue-enqueue! queue t) + (children (cdr al)))])) + ;; Drop alist now that we've consumed it. + (vector-set! s 0 #f)) + (bfs))) (values goto out))) - ;; For (state, byte): explicit goto if present, else recurse through - ;; failure link, terminating at root. - (def (resolve-goto s b nodes) - (let walk ([t s]) - (let* ([t-node (vector-ref nodes t)] - [t-goto (vector-ref t-node 0)] - [m (hashtable-ref t-goto b #f)]) - (cond - [m m] - [(fx= t 0) 0] - [else (walk (vector-ref t-node 1))])))) - ;; --- BFS queue (head-list . tail-list) ---------------------------------- (def (make-simple-queue) (cons '() '()))