chunk-first: pull the trie leaf instead of flattening the whole vector

A pvec is a 32-way trie, but na-chunk-first built each block by calling
pvec-v on the full backing vector — materializing all n elements to a
flat Scheme vector — then copying the 32-wide window out of it. That made
chunk-first O(n), so walking a vector chunk-by-chunk (Clojure's real
chunked map/filter fast path) was O(n^2): a ported chunked map over 500K
elements took 39s, superlinear to ~700s at 2M.

na-chunk-size equals pv-width and blocks are 32-aligned, so a block is
exactly one trie leaf — pv-chunk-for hands it back in O(log n). Copy that
leaf directly; fall back to per-index reads for the rare window that
crosses a leaf boundary. Chunked map is now linear, ~133x faster at 500K
(293ms) and within ~2.3x of the native seq loop, which makes a
clojure-in-clojure seq tier viable.

Corpus rows pin chunk-first window contents + chunk-rest boundaries
against JVM; fixed a stale 'always false' chunked-seq? label.
This commit is contained in:
Yogthos 2026-06-28 01:56:26 -04:00
parent 6940b2c7f5
commit 6d441e2d00
2 changed files with 22 additions and 2 deletions

View file

@ -109,9 +109,24 @@
(let* ((v (cseq-cvec s)) (i (cseq-ci s)))
(cons v (fxmin (fx+ i na-chunk-size) (pvec-count v))))))
(define (na-chunked-seq? x) (and (na-vblock x) #t))
;; Copy the block [i, end) straight out of the pvec trie's 32-element leaf node
;; (pv-chunk-for is O(log n)). na-chunk-size == pv-width and blocks are 32-aligned,
;; so a block is exactly one leaf; the rare non-aligned window crossing a leaf
;; boundary falls back to per-index reads. Flattening the whole backing vector
;; per block (pvec-v) made chunk-first O(n), so walking a vector chunk-by-chunk
;; was O(n^2).
(define (na-chunk-first s)
(let ((vb (na-vblock s)))
(if vb (make-pvec (vec-copy-range (pvec-v (car vb)) (cseq-ci s) (cdr vb)))
(if vb
(let* ((pv (car vb)) (i (cseq-ci s)) (end (cdr vb)) (len (fx- end i))
(node (pv-chunk-for pv i)) (off (fxand i pv-mask)))
(if (fx<=? (fx+ off len) (vector-length node))
(make-pvec (vec-copy-range node off (fx+ off len)))
(let ((out (make-vector len)))
(let loop ((j 0))
(if (fx<? j len)
(begin (vector-set! out j (pvec-nth-d pv (fx+ i j) jolt-nil)) (loop (fx+ j 1)))
(make-pvec out))))))
(jolt-first s)))) ; eager-buffer fallback
(define (na-chunk-rest s)
(let ((vb (na-vblock s)))