diff --git a/jolt-core/clojure/core/20-coll.clj b/jolt-core/clojure/core/20-coll.clj index 829b2aa..7dc44e6 100644 --- a/jolt-core/clojure/core/20-coll.clj +++ b/jolt-core/clojure/core/20-coll.clj @@ -73,9 +73,32 @@ (persistent! (reduce (fn [counts x] (assoc! counts x (inc (get counts x 0)))) (transient {}) coll))) +;; Buckets are transient vectors, not persistent ones: the JVM form rebuilds the +;; bucket's persistent vector per element (conj (get ret k []) x), an O(log n) +;; trie path-rebuild + alloc per element — so a coarse grouping (few large +;; buckets) is bound on that conj, not the map build. Push onto a per-bucket +;; native array (O(1)) instead, then bulk-build the persistent map ONCE. +;; Distinct keys are recorded in a side vector so the buckets can be frozen in +;; place (no second map rebuild). A bucket's FIRST element is stored as a cheap +;; persistent [x]; only the second element promotes it to a transient — so an +;; all-singletons grouping pays no transient alloc and matches the old cost, +;; while any bucket that actually grows rides the O(1) push. (defn group-by [f coll] - (persistent! - (reduce (fn [ret x] (let [k (f x)] (assoc! ret k (conj (get ret k []) x)))) (transient {}) coll))) + (let [tm (transient {}) + ks (reduce (fn [ks x] + (let [k (f x) + b (get tm k)] + (if (nil? b) + (do (assoc! tm k [x]) (conj! ks k)) + (if (vector? b) + (do (assoc! tm k (conj! (transient b) x)) ks) + (do (conj! b x) ks))))) + (transient []) coll)] + (reduce (fn [_ k] + (let [b (get tm k)] + (if (vector? b) nil (assoc! tm k (persistent! b))))) + nil (persistent! ks)) + (persistent! tm))) (defn not-empty [coll] (if (or (nil? coll) (zero? (count coll))) nil coll)) diff --git a/src/jolt/core_coll.janet b/src/jolt/core_coll.janet index 8d945a5..2c16041 100644 --- a/src/jolt/core_coll.janet +++ b/src/jolt/core_coll.janet @@ -182,7 +182,7 @@ (if (core-transient? m) (case (m :kind) :vector (if (and (number? k) (>= k 0) (< k (length (m :arr)))) (in (m :arr) k) default) - :map (let [p (get (m :tbl) (canon-key k))] (if p (in p 1) default)) + :map (let [p (get (m :tbl) (tbl-key k))] (if p (in p 1) default)) :set (if (nil? (get (m :tbl) (canon-key k))) default k)) (if (set? m) (phs-get m k default) (if (phm? m) (phm-get m k default) @@ -251,7 +251,7 @@ (if (core-transient? coll) (case (coll :kind) :vector (and (number? key) (>= key 0) (< key (length (coll :arr)))) - (not (nil? (get (coll :tbl) (canon-key key))))) + (not (nil? (get (coll :tbl) (if (= :map (coll :kind)) (tbl-key key) (canon-key key)))))) (if (set? coll) (phs-contains? coll key) (if (phm? coll) (phm-contains? coll key) (if (pvec? coll) (and (number? key) (>= key 0) (< key (pv-count coll))) diff --git a/src/jolt/core_extra.janet b/src/jolt/core_extra.janet index 6a44bb0..c73e1da 100644 --- a/src/jolt/core_extra.janet +++ b/src/jolt/core_extra.janet @@ -343,7 +343,7 @@ (or (phm? coll) (and (struct? coll) (nil? (get coll :jolt/type)))) (let [t @{}] (each pair (realize-for-iteration coll) - (put t (canon-key (in pair 0)) @[(in pair 0) (in pair 1)])) + (put t (tbl-key (in pair 0)) @[(in pair 0) (in pair 1)])) @{:jolt/type :jolt/transient :kind :map :tbl t}) # mutable-build arrays (vectors/lists) — copy into a transient vector (array? coll) @{:jolt/type :jolt/transient :kind :vector :arr (array/slice coll)} @@ -365,11 +365,11 @@ # a [k v] pair (map-entry / 2-vector) (and (or (pvec? x) (tuple? x) (array? x)) (= 2 (if (pvec? x) (pv-count x) (length x)))) - (put (t :tbl) (canon-key (vnth x 0)) @[(vnth x 0) (vnth x 1)]) + (put (t :tbl) (tbl-key (vnth x 0)) @[(vnth x 0) (vnth x 1)]) # a map: merge all its entries (or (phm? x) (and (struct? x) (nil? (get x :jolt/type)))) (each e (map-entries-of x) - (put (t :tbl) (canon-key (in e 0)) @[(in e 0) (in e 1)])) + (put (t :tbl) (tbl-key (in e 0)) @[(in e 0) (in e 1)])) (error "conj! on a transient map requires a [key value] pair or a map"))) t) @@ -380,7 +380,7 @@ (when (not (and (number? k) (= k (math/floor k)) (>= k 0) (<= k (length a)))) (error (string "Index " k " out of bounds for assoc! on a transient vector of length " (length a)))) (if (= k (length a)) (array/push a v) (put a k v))) - :map (put (t :tbl) (canon-key k) @[k v]) + :map (put (t :tbl) (tbl-key k) @[k v]) (error "assoc! expects a transient vector or map")) t) @@ -405,7 +405,7 @@ (defn core-dissoc! [t & ks] (if (and (core-transient? t) (= :map (t :kind))) - (do (tr-check-active! t) (each k ks (put (t :tbl) (canon-key k) nil)) t) + (do (tr-check-active! t) (each k ks (put (t :tbl) (tbl-key k) nil)) t) (error "dissoc! requires a transient map"))) (defn core-disj! [t & xs] diff --git a/src/jolt/core_types.janet b/src/jolt/core_types.janet index 229ccaa..48868ef 100644 --- a/src/jolt/core_types.janet +++ b/src/jolt/core_types.janet @@ -96,6 +96,14 @@ k))) (set-canonicalize-key! canon-key) +# Janet tables silently drop a nil key (put/get with nil is a no-op), but Clojure +# maps allow a nil key. The transient map keys its native table by canon-key, and +# canon-key returns nil only for nil input — so route nil to a unique sentinel. +# The sentinel is a fresh mutable table; canon-key never produces one, so it can't +# collide with the canon-key of any real key. (phm keeps its own has-nil slot.) +(def tbl-nil-key @{}) +(defn tbl-key [k] (if (nil? k) tbl-nil-key (canon-key k))) + # All [k v] entries of a map (struct or phm), nil-valued keys included. Use this # instead of (keys (phm-to-struct m)) — phm-to-struct drops keys whose value is # nil, which is exactly what Clojure maps must keep. diff --git a/test/spec/maps-spec.janet b/test/spec/maps-spec.janet index 2656453..706d75f 100644 --- a/test/spec/maps-spec.janet +++ b/test/spec/maps-spec.janet @@ -206,5 +206,21 @@ ["dissoc after bulk" "nil" "(get (dissoc (into {} (map (fn [i] [i i]) (range 100))) 50) 50)"] ["frequencies count" "3" "(get (frequencies [1 2 2 1 2 1]) 1)"] ["frequencies coll-key" "2" "(get (frequencies [[1 2] [1 2] [3 4]]) [1 2])"] + # nil is a legal map key, but a Janet table drops a nil key — the transient + # map (canon-keyed table) used to silently lose the nil-key entry, so + # group-by/frequencies/assoc dropped the whole nil bucket. tbl-key routes nil + # to a sentinel; phm keeps its own has-nil slot. + ["frequencies nil key" "2" "(get (frequencies [nil nil 1]) nil)"] + ["group-by nil key" "[nil nil]" "(get (group-by identity [nil nil 1]) nil)"] + ["group-by nil count" "2" "(count (group-by identity [nil nil 1]))"] + ["transient nil key" ":x" "(let [t (transient {})] (assoc! t nil :x) (get (persistent! t) nil))"] + ["transient nil get" "true" "(let [t (transient {})] (assoc! t nil :x) (contains? t nil))"] + ["transient nil dissoc" ":gone" "(let [t (transient {})] (assoc! t nil :x) (dissoc! t nil) (get (persistent! t) nil :gone))"] + # group-by buckets are built on transient vectors (O(1) push) and frozen once, + # rather than an O(log n) persistent conj per element. Result is identical: + # bucket contents and order match the persistent build across a coarse + # grouping (few large buckets — the case bound on the per-bucket conj). ["group-by bucket" "[1 3 5]" "(get (group-by odd? (range 1 6)) true)"] + ["group-by big bucket" "true" "(= (group-by even? (range 200)) {true (vec (filter even? (range 200))) false (vec (filter odd? (range 200)))})"] + ["group-by order" "[0 3 6 9]" "(get (group-by (fn [x] (mod x 3)) (range 10)) 0)"] ["hash-map bulk = incr" "true" "(= (apply hash-map (mapcat (fn [i] [i i]) (range 50))) (reduce (fn [m i] (assoc m i i)) {} (range 50)))"])