group-by: transient-vector buckets + fix nil keys in transient maps (#155)

The map build already used a transient map, but each bucket was rebuilt with
a persistent (conj (get ret k []) x) per element — an O(log n) trie path
rebuild + alloc each. A coarse grouping (few large buckets) was bound on that
conj, not the map build. Buckets are now native arrays (transient vectors,
O(1) push) frozen once; distinct keys are tracked in a side vector so the
buckets freeze in place with no second map rebuild. A bucket's first element
stays a cheap persistent [x] and only promotes to a transient on the second,
so an all-singletons grouping pays no transient alloc.

  coarse (10/100 buckets, 50k): ~313ms -> ~125ms (~2.5x)
  2 buckets (50k):              ~322ms -> ~129ms (~2.5x)
  all-unique (50k):             ~949ms -> ~892ms (no regression)

Surfaced a latent bug: canon-key returns nil for a nil key and Janet tables
drop a nil key, so the canon-keyed transient map silently lost a nil-key
entry — group-by/frequencies/assoc!/into{} dropped the whole nil bucket
((group-by identity [nil nil 1]) gave {1 [1]}, not {nil [nil nil], 1 [1]}).
Route nil through a sentinel (tbl-key) at the transient-map keying sites;
persistent!/count/dissoc! work unchanged since the real [nil v] pair is kept
as the stored value, and phm already has its own has-nil slot. The transient
set has the analogous bug (needs phs nil support) — filed separately.

Co-authored-by: Yogthos <yogthos@gmail.com>

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))

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)))

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]

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.

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)))"])