diff --git a/host/chez/collections.ss b/host/chez/collections.ss index 35936a5..e4b6622 100644 --- a/host/chez/collections.ss +++ b/host/chez/collections.ss @@ -32,45 +32,161 @@ out)) ;; ============================================================================ -;; persistent vector — copy-on-write over a Scheme vector +;; persistent vector — 32-way trie + tail (Clojure's PersistentVector) ;; ============================================================================ -;; A pvec carries an `ent` flag: #t marks a MAP ENTRY (the [k v] pair seq'd out -;; of a map). A map entry equals its [k v] vector and walks like one (nth/count/ -;; seq/=/hash/print all read only `v`), but is NOT `vector?` and IS `map-entry?` -;; — matching Clojure's MapEntry. The flag defaults #f, so every -;; existing `(make-pvec v)` builds a plain vector; modifying an entry (conj/assoc) -;; likewise yields a plain vector. -(define-record-type pvec - (fields v ent) - (protocol (lambda (new) (case-lambda ((v) (new v #f)) ((v e) (new v e))))) - (nongenerative chez-pvec-v1)) -(define empty-pvec (make-pvec (vector))) -(define (jolt-vector . xs) (make-pvec (list->vector xs))) -(define (make-map-entry k v) (make-pvec (vector k v) #t)) -(define (jolt-map-entry? x) (and (pvec? x) (pvec-ent x) #t)) -(define (pvec-count p) (vector-length (pvec-v p))) +;; cnt elements live in a trie of 32-wide nodes (root, height = shift bits) plus a +;; trailing `tail` chunk of 1..32. conj appends to the tail and, when it fills, +;; pushes it into the trie by path-copy — so conj is O(1) amortized and a linear +;; build is O(n), not the O(n^2) of a flat copy-on-write array. nth/assoc/pop are +;; O(log32 n). Trie nodes are Scheme vectors holding only their live children +;; (grown left-to-right), so a node's length is its child count. +;; +;; `ent` #t marks a MAP ENTRY (the [k v] pair seq'd out of a map). An entry has 2 +;; elements (all in the tail), equals its [k v] vector and walks like one, and is +;; both vector? (Clojure's MapEntry implements IPersistentVector) and map-entry?. +;; Modifying an entry (conj/assoc/pop) yields a plain vector (ent #f). +;; +;; make-pvec and pvec-v keep the old flat-vector API: make-pvec builds a trie from +;; a Scheme vector (every existing caller still passes one) and pvec-v materializes +;; it back, so only this file's internals change. +(define pv-bits 5) +(define pv-width 32) +(define pv-mask 31) +(define pv-empty-node (vector)) +(define-record-type (pvec mk-pvec pvec?) + (fields cnt shift root tail ent) (nongenerative chez-pvec-v2)) + +;; trailing helpers over Scheme vectors used by the trie +(define (vec-snoc v x) ; copy v with x appended + (let* ((n (vector-length v)) (out (make-vector (fx+ n 1)))) + (let loop ((i 0)) (when (fx=? i (pv-tailoff (pvec-cnt p))) + (pvec-tail p) + (let loop ((node (pvec-root p)) (level (pvec-shift p))) + (if (fx>? level 0) + (loop (vector-ref node (fxand (fxsra i level) pv-mask)) (fx- level pv-bits)) + node)))) + ;; jolt models every number as a double, so vector indices arrive as flonums — ;; coerce an integer-valued index to a Scheme fixnum before bounds math. (define (->idx i) (if (fixnum? i) i (if (flonum? i) (exact (floor i)) i))) +(define (pvec-count p) (pvec-cnt p)) (define (pvec-nth-d p i d) - (let ((v (pvec-v p)) (i (->idx i))) - (if (and (fixnum? i) (fx>=? i 0) (fxidx i))) + (if (and (fixnum? i) (fx>=? i 0) (fx? (fxsra cnt pv-bits) (fxsll 1 shift)) + ;; root overflow: grow the trie a level + (mk-pvec (fx+ cnt 1) (fx+ shift pv-bits) + (vector (pvec-root p) (pv-new-path shift tail-node)) + (vector x) #f) + (mk-pvec (fx+ cnt 1) shift + (pv-push-tail cnt shift (pvec-root p) tail-node) + (vector x) #f)))))) + +(define (pv-assoc-trie level node i x) + (if (fx=? level 0) + (vec-set node (fxand i pv-mask) x) + (let ((subidx (fxand (fxsra i level) pv-mask))) + (vec-set node subidx (pv-assoc-trie (fx- level pv-bits) (vector-ref node subidx) i x))))) (define (pvec-assoc p i x) ; i in [0,count]; =count appends - (let* ((v (pvec-v p)) (n (vector-length v)) (i (->idx i))) - (cond ((and (fx>=? i 0) (fxidx i)) (cnt (pvec-cnt p))) + (cond + ((fx=? i cnt) (pvec-conj p x)) + ((and (fx>=? i 0) (fx=? i (pv-tailoff cnt)) + (mk-pvec cnt (pvec-shift p) (pvec-root p) + (vec-set (pvec-tail p) (fxand i pv-mask) x) #f) + (mk-pvec cnt (pvec-shift p) + (pv-assoc-trie (pvec-shift p) (pvec-root p) i x) (pvec-tail p) #f))) + (else (error 'assoc "vector index out of bounds"))))) (define (pvec-peek p) - (let ((n (pvec-count p))) (if (fx=? n 0) jolt-nil (vector-ref (pvec-v p) (fx- n 1))))) + (let ((n (pvec-cnt p))) (if (fx=? n 0) jolt-nil (pvec-nth-d p (fx- n 1) jolt-nil)))) +;; pop the last trie chunk back into the tail; #f means the subtree emptied. +(define (pv-pop-tail cnt level node) + (let ((subidx (fxand (fxsra (fx- cnt 2) level) pv-mask))) + (cond + ((fx>? level pv-bits) + (let ((newchild (pv-pop-tail cnt (fx- level pv-bits) (vector-ref node subidx)))) + (cond ((and (not newchild) (fx=? subidx 0)) #f) + (newchild (vec-set node subidx newchild)) + (else (vec-take node subidx))))) + ((fx=? subidx 0) #f) + (else (vec-take node subidx))))) (define (pvec-pop p) - (let ((n (pvec-count p))) - (if (fx=? n 0) (error 'pop "can't pop empty vector") - (make-pvec (vec-copy-range (pvec-v p) 0 (fx- n 1)))))) + (let ((cnt (pvec-cnt p)) (shift (pvec-shift p))) + (cond + ((fx=? cnt 0) (error 'pop "can't pop empty vector")) + ((fx=? cnt 1) empty-pvec) + ((fx>? (fx- cnt (pv-tailoff cnt)) 1) + (mk-pvec (fx- cnt 1) shift (pvec-root p) (vec-drop-last (pvec-tail p)) #f)) + (else + (let* ((new-tail (pv-chunk-for p (fx- cnt 2))) + (popped (pv-pop-tail cnt shift (pvec-root p))) + (new-root (or popped pv-empty-node))) + (if (and (fx>? shift pv-bits) (fxvector xs))) +(define (make-map-entry k v) (make-pvec (vector k v) #t)) +(define (jolt-map-entry? x) (and (pvec? x) (pvec-ent x) #t)) ;; ============================================================================ ;; bitmap HAMT — keys hashed by jolt-hash, leaves compared by jolt= diff --git a/test/chez/corpus.edn b/test/chez/corpus.edn index ac39aba..424a8fd 100644 --- a/test/chez/corpus.edn +++ b/test/chez/corpus.edn @@ -3211,4 +3211,9 @@ {:suite "protocols / arity dispatch" :label "protocol method extended to a host type with an arg" :expected "\"S:hi:5\"" :actual "(do (defprotocol Q (mq [this a])) (extend-protocol Q java.lang.String (mq [this a] (str \"S:\" this \":\" a))) (mq \"hi\" 5))"} {:suite "protocols / arity dispatch" :label "protocol method used as a value" :expected "[1 2]" :actual "(do (defprotocol Z (z [this])) (defrecord ZR [v] Z (z [_] v)) (mapv z [(->ZR 1) (->ZR 2)]))"} {:suite "protocols / arity dispatch" :label "no-extra-arg method dispatches by type" :expected "[16 25]" :actual "(do (defprotocol Sh (ar [s])) (defrecord Sq [x] Sh (ar [_] (* x x))) (defrecord Sq2 [y] Sh (ar [_] (* y y))) [(ar (->Sq 4)) (ar (->Sq2 5))])"} + {:suite "vectors / large trie" :label "linear build equals range across paths" :expected "true" :actual "(= (vec (range 2000)) (reduce conj [] (range 2000)) (into [] (range 2000)))"} + {:suite "vectors / large trie" :label "nth across a multi-level vector" :expected "[0 1024 2000 2099]" :actual "(let [v (vec (range 2100))] [(nth v 0) (nth v 1024) (nth v 2000) (nth v 2099)])"} + {:suite "vectors / large trie" :label "assoc in a large vector" :expected "[1098 :x 1100]" :actual "(let [v (assoc (vec (range 1100)) 1099 :x)] [(nth v 1098) (nth v 1099) (count v)])"} + {:suite "vectors / large trie" :label "pop down across a level boundary" :expected "[1023 1022]" :actual "(let [p (pop (vec (range 1024)))] [(count p) (peek p)])"} + {:suite "vectors / large trie" :label "pop to empty then rebuild" :expected "[0 5]" :actual "(let [e (loop [v (vec (range 100))] (if (seq v) (recur (pop v)) v))] [(count e) (count (into e (range 5)))])"} ]