Persistent hash map: HAMT instead of O(n) copy-on-write (jolt-684u) (#136)

* Add benchmark suite for alloc/dispatch/collection workloads (jolt-1r86)

The ray tracer is float-compute-bound (devirt, alloc removal, type-proving all
measured flat on it), so it can't validate the optimization passes. Add a small
cross-language suite (AWFY + CLBG style, portable Clojure) isolating the axes it
misses:

  binary-trees  allocation / GC pressure (escaping short-lived records)
  dispatch      megamorphic protocol dispatch (~1M dispatches/s; WP can't devirt)
  collections   persistent map/vector churn

bench/run.sh runs them; bench/README.md maps each to the pass it exercises.

collections immediately surfaced jolt-684u: the persistent hash map is O(n) per
assoc (flat copy-on-write bucket array, not a HAMT) — n=4000 assocs take 50s.
Invisible to the ray tracer (no maps).

* Persistent hash map: HAMT instead of O(n) copy-on-write (jolt-684u)

The map was a flat bucket array whose assoc copied the whole array every insert
(O(n)/assoc, O(n^2) to build). Compounding it, small maps are Janet structs that
only promoted to phm for collection keys — never for size — so a scalar-key map
stayed an O(n)-copy struct forever. Building a 4000-entry map took ~50s.

Two fixes, following ClojureScript's design:

- phm.janet is now a HAMT (hash array mapped trie): BitmapIndexedNode /
  ArrayNode / HashCollisionNode, 32-way, 5 hash bits per level, structural
  sharing — assoc/dissoc/get are O(log32 n). Translated from cljs.core, adapted
  to Janet's 32-bit bit-ops (the hash is carried unsigned, the level index is
  extracted with arithmetic, and bits are tested with band against 1<<i since
  brushift rejects negative bitmaps). The public phm-* API and the value shape
  (:jolt/type :jolt/phm, :cnt) are unchanged; transients are a separate rep and
  untouched.

- core_coll promotes a struct map to a phm past 8 entries (not only for
  collection keys), mirroring cljs PersistentArrayMap -> PersistentHashMap, so
  incremental building isn't O(n^2).

20000 raw assocs: 7.1s -> 0.105s. The collections benchmark: 16.7s -> 0.2s.
Correctness covered by test/unit/phm-hamt-test.janet (oracle vs a Janet table,
nil keys, dissoc, a real hash-collision pair, and a sub-linear-assoc guard);
full gate green.

---------

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

src/jolt/core_coll.janet

@@ -19,6 +19,12 @@
 # Collections
 # ============================================================
 
+# Small maps are Janet structs (native, O(1) get) but assoc copies them whole
+# (O(n)); past this many entries a map promotes to the phm HAMT (O(log n) assoc)
+# so incremental building isn't O(n^2). Mirrors cljs PersistentArrayMap's
+# HASHMAP-THRESHOLD (jolt-684u).
+(def- map-array-threshold 8)
+
 # Is x a map value (for conj/merge semantics: conj-ing a map merges its entries)?
 (defn- map-value? [x]
   (or (phm? x) (and (struct? x) (nil? (get x :jolt/type)))))
@@ -124,16 +130,21 @@
             (if (= idx (length result)) (array/push result v) (put result idx v)))
           (+= i 2))
         (if (tuple? m) (tuple/slice (tuple ;result)) result))
-    # map (struct/table). Promote to a phm when any new key is a collection (a
+    # map (struct/table). Promote to a phm when (a) any new key is a collection
-    # Janet struct/table would key it by identity) or any new key/value is nil (a
+    # (a Janet struct/table would key it by identity) or any new key/value is nil
-    # struct drops nil; phm preserves it, matching Clojure). m itself is a struct
+    # (a struct drops nil; phm preserves it), or (b) the result would exceed the
-    # here (phm handled above), so only the new kvs can introduce these.
+    # small-map threshold — a Janet struct copies wholesale on assoc (O(n)), so a
+    # growing map must ride the phm HAMT (O(log n)) past ~8 entries. Mirrors cljs
+    # PersistentArrayMap -> PersistentHashMap (jolt-684u). m is a struct here
+    # (phm handled above), so only the current size + new kvs matter.
     (let [coll-key (do (var c false) (var i 0)
                      (while (< i (length kvs))
                        (let [k (in kvs i) v (in kvs (+ i 1))]
                          (when (or (table? k) (array? k) (nil? k) (nil? v)) (set c true)))
-                       (+= i 2)) c)]
+                       (+= i 2)) c)
-      (if coll-key
+          promote (or coll-key
+                      (> (+ (if m (length m) 0) (/ (length kvs) 2)) map-array-threshold))]
+      (if promote
         (do (var result (make-phm))
             (when m (each k (keys m) (set result (phm-assoc result k (get m k)))))
             (var i 0) (while (< i (length kvs)) (set result (phm-assoc result (in kvs i) (in kvs (+ i 1)))) (+= i 2))

src/jolt/phm.janet

@@ -1,29 +1,43 @@
-# PersistentHashMap implementation for Jolt
+# PersistentHashMap for Jolt — a HAMT (hash array mapped trie), the structure
-# Bucket-based hash map with copy-on-write semantics. The bucket array GROWS
+# Clojure/ClojureScript/jank use. 32-way branching, 5 hash bits per level, with
-# (doubling, rehash) when the load factor passes 2 entries/bucket, so lookups
+# structural sharing: assoc/dissoc/get are O(log32 n) ~ effectively constant.
-# stay O(1)-ish at any size — with a fixed 8 buckets, a 100-entry map was a
+# Replaces the old flat copy-on-write bucket array, which was O(n) per assoc
-# ~12-entry linear scan per get (the jolt-s3y map-read regression). The bucket
+# (O(n^2) to build — jolt-684u). Translated from the ClojureScript
-# count is derived from (length (m :buckets)), so marshaled images from before
+# PersistentHashMap (cljs.core: BitmapIndexedNode / ArrayNode / HashCollisionNode).
-# this change keep working.
 #
-# REP vs API: this file is ONLY the hash-map representation (phm-* primitives).
+# REP vs API: this file is ONLY the map representation (phm-* primitives). The
-# The Clojure-facing map ops (assoc/dissoc/get/conj/count/seq dispatch, nil-key
+# Clojure-facing map ops (assoc/dissoc/get/conj/count/seq dispatch, nil-key,
-# handling, merge) live in core_coll.janet / core_types.janet, which recognize
+# merge) live in core_coll.janet / core_types.janet, which recognize phm by its
-# phm by its `:jolt/deftype` string. PersistentHashSet is layered on top in
+# `:jolt/deftype` string and call these primitives. PersistentHashSet is layered
-# phs.janet; LazySeq (historically here) now lives in lazyseq.janet.
+# on top in phs.janet. Transients are a separate mutable-table rep (core_types),
+# so they don't touch this file.
+#
+# Node representation (Janet tuples, tagged at index 0; their arrays are built
+# fresh on every modify and never mutated in place, so sharing is safe):
+#   [:bin bitmap arr]   bitmap-indexed: arr is [k v k v ...]; a nil k means v is
+#                       a sub-node (the slot recurses one level deeper).
+#   [:an  cnt arr]      array-node: arr is 32 slots of sub-node-or-nil.
+#   [:hcn hash cnt arr] hash-collision: arr is [k v k v ...] of same-hash keys.
+# The map itself stays a table tagged :jolt/phm with :cnt (read directly by
+# core), plus :root (the trie, nil when empty) and :has-nil/:nil-val (Clojure
+# maps allow a nil key, which the trie can't store because a nil k marks a
+# sub-node).
+#
+# Note on Janet `=`: it is STRUCTURAL on tuples, not identity, so cljs's
+# `(identical? n node)` "nothing changed" early-outs are dropped here — we just
+# rebuild the O(log32 n) path (an extra alloc only when overwriting an identical
+# value). nil returns (a node that became empty) are kept; those are real nils.
 
-(def- initial-buckets 8)
+(def- DEFTYPE "jolt.lang.persistent-hash-map.PersistentHashMap")
 
 (defn phm? [x]
-  (and (table? x)
+  (and (table? x) (= DEFTYPE (x :jolt/deftype))))
-       (= "jolt.lang.persistent-hash-map.PersistentHashMap" (x :jolt/deftype))))
 
-# Keys are hashed and compared by VALUE. Scalars (keywords/strings/numbers) are
+# Keys are hashed and compared by VALUE. Scalars are value-hashable in Janet;
-# value-hashable in Janet already, but collection keys (a phm/pvec/plist map or
+# collection keys (a phm/pvec/plist/vector) are Janet tables hashed by identity,
-# vector) are Janet tables hashed by identity — so they're canonicalized to a
+# so they're canonicalized to a value-hashable struct/tuple first.
-# value-hashable struct/tuple first. `canonicalize-key` is injected by core (which
+# `canonicalize-key` is injected by core (which knows those types); phm stays
-# knows the pvec/plist/phm types); phm stays dependency-free. Keys are still
+# dependency-free. Keys are STORED as-is, so retrieval/iteration return originals.
-# *stored* as-is, so retrieval and iteration return the original key objects.
 (var canonicalize-key nil)
 (defn set-canonicalize-key!
   "Install the value-canonicalizer for collection keys (called by core)."
@@ -35,143 +49,258 @@
     k))
 (defn canon
   "Public canonicalizer: maps a key to its value-hashable form (identity for
-  scalars). Used by callers that index the same canonicalized tables phm uses
+  scalars). Used by callers that index the same canonicalized tables phm uses."
-  (e.g. transient maps/sets)."
   [k] (ck k))
-# Identity/scalar equality first — the common case — before paying for
+# Identity/scalar equality first (the common case), then value equality.
-# canonicalization of collection keys.
 (defn- key= [a b] (or (= a b) (= (ck a) (ck b))))
+# Janet bit ops are 32-bit and split awkwardly: band/bor/bxor want SIGNED
+# operands, brushift wants UNSIGNED — and `hash` is a signed 32-bit int. So we
+# carry the hash as an UNSIGNED 32-bit value, extract the 5-bit level index with
+# arithmetic (mask), and test/count bits with band against (1<<i) only — never
+# brushift on a possibly-negative bitmap.
+(defn- khash [k] (let [h (hash (ck k))] (if (< h 0) (+ h 0x100000000) h)))
 
-(defn- hash-idx [m k]
+# --- HAMT node machinery (translated from cljs.core) ------------------------
-  (if (nil? k) 0 (mod (hash (ck k)) (length (m :buckets)))))
+(def- EMPTY-BIN [:bin 0 (tuple)])
 
-# Index of key k in a bucket (entries are stored stride-2: k v k v ...), or nil.
+(defn- mask [h shift] (% (math/floor (/ h (blshift 1 shift))) 32))
-# The single scan all the bucket ops share — keeping it one place stops the
+(defn- bitpos [h shift] (blshift 1 (mask h shift)))
-# stride logic drifting between them.
+# popcount of a 32-bit bitmap; bidx = popcount of the bits below level index m.
-(defn- bucket-index-of [bucket k]
+(defn- popcount [bm]
-  (var i 0) (def n (length bucket)) (var found nil)
+  (var c 0) (var i 0)
-  (while (< i n)
+  (while (< i 32) (when (not= 0 (band bm (blshift 1 i))) (++ c)) (++ i))
-    (if (key= k (in bucket i)) (do (set found i) (break)))
+  c)
-    (+= i 2))
+(defn- bidx [bitmap m]
-  found)
+  (var c 0) (var i 0)
+  (while (< i m) (when (not= 0 (band bitmap (blshift 1 i))) (++ c)) (++ i))
+  c)
 
-(defn- phm-bucket-find [bucket k]
+(defn- cset1 [arr i a] (def n (array ;arr)) (put n i a) n)
-  (let [i (bucket-index-of bucket k)]
+(defn- cset2 [arr i a j b] (def n (array ;arr)) (put n i a) (put n j b) n)
-    (if i (in bucket (+ i 1)) nil)))
+# remove the pair at pair-index p (elements 2p, 2p+1)
+(defn- remove-pair [arr p]
+  (def out (array)) (def a (* 2 p)) (def b (+ a 1)) (def L (length arr))
+  (var x 0) (while (< x L) (unless (or (= x a) (= x b)) (array/push out (in arr x))) (++ x))
+  out)
+# element-index of key k in a collision arr [k v k v ...], or -1
+(defn- hcn-find [arr cnt k]
+  (var i 0) (def L (* 2 cnt)) (var r -1)
+  (while (< i L) (when (key= k (in arr i)) (set r i) (break)) (+= i 2))
+  r)
+# rebuild a bin-node from an array-node's slots when it shrinks (<=8) — each
+# surviving sub-node becomes a nil-key slot.
+(defn- pack-array-node [arr removed-idx]
+  (def out (array)) (var bitmap 0) (var i 0)
+  (while (< i 32)
+    (when (and (not= i removed-idx) (not (nil? (in arr i))))
+      (array/push out nil) (array/push out (in arr i))
+      (set bitmap (bor bitmap (blshift 1 i))))
+    (++ i))
+  [:bin bitmap out])
 
-(defn phm-bucket-contains? [bucket k]
+# mutual recursion across node types
-  (not (nil? (bucket-index-of bucket k))))
+(var node-assoc nil) (var node-lookup nil) (var node-without nil) (var create-node nil)
 
-(defn- phm-bucket-assoc [bucket k v]
+(defn- bin-assoc [node shift h key val added]
-  (def n (length bucket))
+  (def bitmap (in node 1)) (def arr (in node 2))
-  (def found-i (bucket-index-of bucket k))
+  (def m (mask h shift)) (def bit (blshift 1 m))
-  (if (not (nil? found-i))
+  (def idx (bidx bitmap m))
-    (let [nb @[]] (var j 0)
+  (if (zero? (band bitmap bit))
-      (while (< j n) (array/push nb (if (= j (+ found-i 1)) v (in bucket j))) (++ j)) nb)
+    (let [nkeys (popcount bitmap)]
-    (let [nb @[]] (var j 0)
+      (if (>= nkeys 16)
-      (while (< j n) (array/push nb (in bucket j)) (++ j))
+        # expand to an array-node
-      (array/push nb k) (array/push nb v) nb)))
+        (let [nodes (array/new-filled 32 nil)
+              jdx (mask h shift)]
+          (put nodes jdx (bin-assoc EMPTY-BIN (+ shift 5) h key val added))
+          (var i 0) (var j 0)
+          (while (< i 32)
+            (unless (zero? (band bitmap (blshift 1 i)))
+              (let [ek (in arr (* 2 j)) ev (in arr (+ 1 (* 2 j)))]
+                (put nodes i (if (nil? ek) ev
+                               (bin-assoc EMPTY-BIN (+ shift 5) (khash ek) ek ev added)))
+                (++ j)))
+            (++ i))
+          [:an (+ nkeys 1) nodes])
+        # insert (key val) into the bin arr at pair-index idx
+        (let [out (array)]
+          (array/concat out (array/slice arr 0 (* 2 idx)))
+          (array/push out key) (array/push out val)
+          (array/concat out (array/slice arr (* 2 idx)))
+          (put added 0 true)
+          [:bin (bor bitmap bit) out])))
+    # slot occupied
+    (let [kon (in arr (* 2 idx)) von (in arr (+ 1 (* 2 idx)))]
+      (cond
+        (nil? kon)
+          [:bin bitmap (cset1 arr (+ 1 (* 2 idx)) (node-assoc von (+ shift 5) h key val added))]
+        (key= key kon)
+          [:bin bitmap (cset1 arr (+ 1 (* 2 idx)) val)]
+        (do (put added 0 true)
+            [:bin bitmap (cset2 arr (* 2 idx) nil (+ 1 (* 2 idx))
+                               (create-node (+ shift 5) kon von h key val))])))))
 
-(defn- phm-bucket-dissoc [bucket k]
+(defn- an-assoc [node shift h key val added]
-  (def n (length bucket))
+  (def cnt (in node 1)) (def arr (in node 2))
-  (def found-i (bucket-index-of bucket k))
+  (def idx (mask h shift)) (def sub (in arr idx))
-  (if (nil? found-i) bucket
+  (if (nil? sub)
-    (if (= n 2) nil
+    [:an (+ cnt 1) (cset1 arr idx (bin-assoc EMPTY-BIN (+ shift 5) h key val added))]
-      (let [nb @[]] (var j 0)
+    [:an cnt (cset1 arr idx (node-assoc sub (+ shift 5) h key val added))]))
-        (while (< j found-i) (array/push nb (in bucket j)) (++ j))
+
-        (while (< j (- n 2)) (array/push nb (in bucket (+ j 2))) (++ j)) nb))))
+(defn- hcn-assoc [node shift h key val added]
+  (def chash (in node 1)) (def cnt (in node 2)) (def arr (in node 3))
+  (if (= h chash)
+    (let [idx (hcn-find arr cnt key)]
+      (if (= idx -1)
+        (let [out (array ;arr)] (array/push out key) (array/push out val)
+             (put added 0 true) [:hcn chash (+ cnt 1) out])
+        (if (= (in arr (+ 1 idx)) val) node
+          [:hcn chash cnt (cset1 arr (+ 1 idx) val)])))
+    # different hash at this level: wrap in a bin node, then assoc
+    (bin-assoc [:bin (bitpos chash shift) (array nil node)] shift h key val added)))
+
+(set create-node (fn [shift k1 v1 k2h k2 v2]
+  (def k1h (khash k1))
+  (if (= k1h k2h)
+    [:hcn k1h 2 (array k1 v1 k2 v2)]
+    (let [added @[false]
+          n1 (bin-assoc EMPTY-BIN shift k1h k1 v1 added)]
+      (bin-assoc n1 shift k2h k2 v2 added)))))
+
+(set node-assoc (fn [node shift h key val added]
+  (case (in node 0)
+    :bin (bin-assoc node shift h key val added)
+    :an (an-assoc node shift h key val added)
+    :hcn (hcn-assoc node shift h key val added))))
+
+(defn- bin-lookup [node shift h key nf]
+  (def bitmap (in node 1)) (def arr (in node 2))
+  (def m (mask h shift)) (def bit (blshift 1 m))
+  (if (zero? (band bitmap bit)) nf
+    (let [idx (bidx bitmap m) kon (in arr (* 2 idx)) von (in arr (+ 1 (* 2 idx)))]
+      (cond
+        (nil? kon) (node-lookup von (+ shift 5) h key nf)
+        (key= key kon) von
+        nf))))
+
+(defn- an-lookup [node shift h key nf]
+  (def sub (in (in node 2) (mask h shift)))
+  (if (nil? sub) nf (node-lookup sub (+ shift 5) h key nf)))
+
+(defn- hcn-lookup [node shift h key nf]
+  (def idx (hcn-find (in node 3) (in node 2) key))
+  (if (< idx 0) nf (in (in node 3) (+ 1 idx))))
+
+(set node-lookup (fn [node shift h key nf]
+  (case (in node 0)
+    :bin (bin-lookup node shift h key nf)
+    :an (an-lookup node shift h key nf)
+    :hcn (hcn-lookup node shift h key nf))))
+
+(defn- bin-without [node shift h key]
+  (def bitmap (in node 1)) (def arr (in node 2))
+  (def m (mask h shift)) (def bit (blshift 1 m))
+  (if (zero? (band bitmap bit)) node
+    (let [idx (bidx bitmap m) kon (in arr (* 2 idx)) von (in arr (+ 1 (* 2 idx)))]
+      (cond
+        (nil? kon)
+          (let [nn (node-without von (+ shift 5) h key)]
+            (cond (not (nil? nn)) [:bin bitmap (cset1 arr (+ 1 (* 2 idx)) nn)]
+                  (= bitmap bit) nil
+                  [:bin (bxor bitmap bit) (remove-pair arr idx)]))
+        (key= key kon)
+          (if (= bitmap bit) nil [:bin (bxor bitmap bit) (remove-pair arr idx)])
+        node))))
+
+(defn- an-without [node shift h key]
+  (def cnt (in node 1)) (def arr (in node 2))
+  (def idx (mask h shift)) (def sub (in arr idx))
+  (if (nil? sub) node
+    (let [nn (node-without sub (+ shift 5) h key)]
+      (cond
+        (nil? nn) (if (<= cnt 8) (pack-array-node arr idx) [:an (- cnt 1) (cset1 arr idx nil)])
+        [:an cnt (cset1 arr idx nn)]))))
+
+(defn- hcn-without [node shift h key]
+  (def chash (in node 1)) (def cnt (in node 2)) (def arr (in node 3))
+  (def idx (hcn-find arr cnt key))
+  (cond (= idx -1) node
+        (= cnt 1) nil
+        [:hcn chash (- cnt 1) (remove-pair arr (brshift idx 1))]))
+
+(set node-without (fn [node shift h key]
+  (case (in node 0)
+    :bin (bin-without node shift h key)
+    :an (an-without node shift h key)
+    :hcn (hcn-without node shift h key))))
+
+# depth-first walk: call (f k v) for every entry (the nil key is handled at the
+# map level, not in the trie).
+(defn- node-each [node f]
+  (case (in node 0)
+    :bin (let [arr (in node 2) L (length arr)]
+           (var i 0)
+           (while (< i L)
+             (let [k (in arr i) v (in arr (+ i 1))]
+               (if (nil? k) (when v (node-each v f)) (f k v)))
+             (+= i 2)))
+    :an (let [arr (in node 2)]
+          (var i 0) (while (< i 32) (when (in arr i) (node-each (in arr i) f)) (++ i)))
+    :hcn (let [arr (in node 3) L (* 2 (in node 2))]
+           (var i 0) (while (< i L) (f (in arr i) (in arr (+ i 1))) (+= i 2)))))
+
+# --- map value + public API -------------------------------------------------
+(defn- mk [cnt root has-nil nil-val meta]
+  @{:jolt/type :jolt/phm :jolt/deftype DEFTYPE
+    :cnt cnt :root root :has-nil has-nil :nil-val nil-val :_meta meta})
 
 (defn phm-get [m k &opt default]
   (default default nil)
-  (let [bucket (get (m :buckets) (hash-idx m k))]
+  (if (nil? k)
-    # Single pass with a presence flag (not nil-of-value): a key mapped to nil
+    (if (m :has-nil) (m :nil-val) default)
-    # is still present, so return nil (not the default) when it exists.
+    (let [root (m :root)]
-    (if bucket
+      (if root (node-lookup root 0 (khash k) k default) default))))
-      (let [i (bucket-index-of bucket k)]
-        (if i (in bucket (+ i 1)) default))
-      default)))
 
-# Rehash every entry of `buckets` into a fresh array of `nb` buckets.
+(def- NF (gensym))
-(defn- rehash [buckets nb]
+(defn phm-contains? [m k]
-  (def out (array/new-filled nb nil))
+  (if (nil? k) (truthy? (m :has-nil))
-  (each bucket buckets
+    (let [root (m :root)]
-    (when bucket
+      (if root (not= (node-lookup root 0 (khash k) k NF) NF) false))))
-      (var i 0) (var n (length bucket))
-      (while (< i n)
-        (let [k (in bucket i)
-              idx (if (nil? k) 0 (mod (hash (ck k)) nb))]
-          (when (nil? (in out idx)) (put out idx @[]))
-          (array/push (in out idx) k)
-          (array/push (in out idx) (in bucket (+ i 1))))
-        (+= i 2))))
-  out)
 
 (defn phm-assoc [m k v]
-  (let [cnt (m :cnt) idx (hash-idx m k)
+  (if (nil? k)
-        old-bucket (get (m :buckets) idx)
+    (mk (if (m :has-nil) (m :cnt) (+ (m :cnt) 1)) (m :root) true v (m :_meta))
-        had-key (if old-bucket (phm-bucket-contains? old-bucket k) false)
+    (let [added @[false]
-        new-bucket (phm-bucket-assoc (if old-bucket old-bucket @[]) k v)
+          root (or (m :root) EMPTY-BIN)
-        new-cnt (if had-key cnt (+ cnt 1))
+          nroot (node-assoc root 0 (khash k) k v added)]
-        nbuckets (length (m :buckets))
+      (mk (if (in added 0) (+ (m :cnt) 1) (m :cnt)) nroot (m :has-nil) (m :nil-val) (m :_meta)))))
-        new-buckets (array/new nbuckets)]
-    (var bi 0)
-    (while (< bi nbuckets)
-      (put new-buckets bi (if (= bi idx) new-bucket (get (m :buckets) bi))) (++ bi))
-    # Grow past load factor 2 (doubling) so buckets stay short. Done on the
-    # copy, so persistence is untouched.
-    (def grown (if (> new-cnt (* 2 nbuckets))
-                 (rehash new-buckets (* 2 nbuckets))
-                 new-buckets))
-    @{:jolt/type :jolt/phm :jolt/deftype "jolt.lang.persistent-hash-map.PersistentHashMap"
-      :cnt new-cnt :buckets grown :_meta (m :_meta)}))
 
 (defn phm-dissoc [m k]
-  (let [idx (hash-idx m k) old-bucket (get (m :buckets) idx)]
+  (if (nil? k)
-    (if old-bucket
+    (if (m :has-nil) (mk (- (m :cnt) 1) (m :root) false nil (m :_meta)) m)
-      (let [new-bucket (phm-bucket-dissoc old-bucket k)]
+    (let [root (m :root)]
-        (if (= new-bucket old-bucket) m
+      (if (and root (phm-contains? m k))
-          (let [new-cnt (- (m :cnt) 1)
+        (mk (- (m :cnt) 1) (node-without root 0 (khash k) k) (m :has-nil) (m :nil-val) (m :_meta))
-                nbuckets (length (m :buckets))
+        m))))
-                new-buckets (array/new nbuckets)]
-            (var bi 0)
-            (while (< bi nbuckets)
-              (put new-buckets bi (if (= bi idx) new-bucket (get (m :buckets) bi))) (++ bi))
-            @{:jolt/type :jolt/phm :jolt/deftype "jolt.lang.persistent-hash-map.PersistentHashMap"
-              :cnt new-cnt :buckets new-buckets :_meta (m :_meta)})))
-      m)))
 
 (defn phm-entries [m]
-  (var result @[]) (var bi 0) (def nb (length (m :buckets)))
+  (def out @[])
-  (while (< bi nb)
+  (when (m :has-nil) (array/push out [nil (m :nil-val)]))
-    (let [bucket (get (m :buckets) bi)]
+  (when (m :root) (node-each (m :root) (fn [k v] (array/push out [k v]))))
-      (when bucket
+  out)
-        (var i 0) (var n (length bucket))
-        (while (< i n) (array/push result [(in bucket i) (in bucket (+ i 1))]) (+= i 2))))
-    (++ bi))
-  result)
 
 (defn phm-to-struct [m]
-  (var result @{}) (var bi 0) (def nb (length (m :buckets)))
+  # a Janet struct can't hold a nil key (matches Clojure struct/keys behavior);
-  (while (< bi nb)
+  # every other entry carries over.
-    (let [bucket (get (m :buckets) bi)]
+  (def t @{})
-      (when bucket
+  (when (m :root) (node-each (m :root) (fn [k v] (put t k v))))
-        (var i 0) (var n (length bucket))
+  (table/to-struct t))
-        (while (< i n) (put result (in bucket i) (in bucket (+ i 1))) (+= i 2))))
-    (++ bi))
-  (table/to-struct result))
 
 (defn phm-count [m] (m :cnt))
 
-(defn phm-contains? [m k]
-  (let [bucket (get (m :buckets) (hash-idx m k))]
-    (if bucket (phm-bucket-contains? bucket k) false)))
-
 (defn make-phm [&opt kvs]
   (default kvs nil)
-  (var m @{:jolt/type :jolt/phm :jolt/deftype "jolt.lang.persistent-hash-map.PersistentHashMap"
+  (var m (mk 0 nil false nil nil))
-           :cnt 0 :buckets (array/new-filled initial-buckets nil) :_meta nil})
   (when kvs
-    (var i 0) (var n (length kvs))
+    (var i 0) (def n (length kvs))
     (while (< i n) (set m (phm-assoc m (kvs i) (kvs (+ i 1)))) (+= i 2)))
   m)
-

test/unit/phm-hamt-test.janet

@@ -0,0 +1,107 @@
+# Persistent hash map correctness + complexity (jolt-684u). Exercises the phm-*
+# primitives directly (scalar keys — the perf-critical path) against a plain
+# Janet table oracle, then asserts assoc is sub-linear per op (a HAMT, not the
+# old O(n) copy-on-write bucket array). Collection-key value-equality is covered
+# by the full clojure suite (needs core's canonicalize-key); here we test the
+# representation in isolation.
+(import ../../src/jolt/phm :as phm)
+
+(var fails 0)
+(defn check [label got expected]
+  (if (deep= got expected) (print "  ok   " label)
+    (do (++ fails) (printf "  FAIL %s: want %p got %p" label expected got))))
+
+# --- basic assoc / get / overwrite / dissoc / count -------------------------
+(var m (phm/make-phm))
+(check "empty count" (phm/phm-count m) 0)
+(check "get missing -> default" (phm/phm-get m :x :none) :none)
+(set m (phm/phm-assoc m :a 1))
+(set m (phm/phm-assoc m :b 2))
+(set m (phm/phm-assoc m :c 3))
+(check "count after 3" (phm/phm-count m) 3)
+(check "get a" (phm/phm-get m :a) 1)
+(check "get c" (phm/phm-get m :c) 3)
+(check "contains b" (phm/phm-contains? m :b) true)
+(check "contains missing" (phm/phm-contains? m :z) false)
+(set m (phm/phm-assoc m :b 22))           # overwrite
+(check "overwrite value" (phm/phm-get m :b) 22)
+(check "overwrite keeps count" (phm/phm-count m) 3)
+(set m (phm/phm-dissoc m :a))
+(check "dissoc removes" (phm/phm-get m :a :gone) :gone)
+(check "dissoc decrements" (phm/phm-count m) 2)
+(check "dissoc missing is noop" (phm/phm-count (phm/phm-dissoc m :zzz)) 2)
+
+# --- nil key (Clojure maps allow nil keys; struct drops them) ---------------
+(var n (phm/phm-assoc (phm/make-phm) nil :nilval))
+(check "nil key get" (phm/phm-get n nil) :nilval)
+(check "nil key count" (phm/phm-count n) 1)
+(check "nil key contains" (phm/phm-contains? n nil) true)
+(set n (phm/phm-assoc n :k :v))
+(check "nil + other count" (phm/phm-count n) 2)
+(check "nil after add other" (phm/phm-get n nil) :nilval)
+(set n (phm/phm-dissoc n nil))
+(check "dissoc nil key" (phm/phm-get n nil :gone) :gone)
+(check "dissoc nil count" (phm/phm-count n) 1)
+
+# --- many keys vs a Janet-table oracle (string + int + keyword keys) --------
+(def oracle @{})
+(var big (phm/make-phm))
+(def N 5000)
+(loop [i :range [0 N]]
+  (def k (cond (= 0 (% i 3)) (string "s" i)
+               (= 1 (% i 3)) i
+               (keyword "k" i)))
+  (put oracle k i)
+  (set big (phm/phm-assoc big k i)))
+(check "big count == oracle" (phm/phm-count big) (length oracle))
+(var mism 0)
+(eachp [k v] oracle (unless (= (phm/phm-get big k :MISS) v) (++ mism)))
+(check "all keys read back correctly" mism 0)
+# entries round-trip
+(check "entries count" (length (phm/phm-entries big)) (length oracle))
+(def back @{})
+(each e (phm/phm-entries big) (put back (in e 0) (in e 1)))
+(check "entries round-trip == oracle" back oracle)
+# dissoc half, recheck
+(var half big)
+(loop [i :range [0 N 2]]
+  (def k (cond (= 0 (% i 3)) (string "s" i) (= 1 (% i 3)) i (keyword "k" i)))
+  (set half (phm/phm-dissoc half k)))
+(var hmism 0)
+(loop [i :range [0 N]]
+  (def k (cond (= 0 (% i 3)) (string "s" i) (= 1 (% i 3)) i (keyword "k" i)))
+  (def want (if (even? i) :gone i))
+  (unless (= (phm/phm-get half k :gone) want) (++ hmism)))
+(check "dissoc half reads correctly" hmism 0)
+
+# --- hash collisions (HashCollisionNode path) -------------------------------
+# "k6595" and "k144747" both hash to 690120568 — distinct keys, same 32-bit hash.
+(when (= (hash "k6595") (hash "k144747"))   # guard: only if Janet's hash still collides them
+  (var c (phm/make-phm))
+  (set c (phm/phm-assoc c "k6595" :A))
+  (set c (phm/phm-assoc c "k144747" :B))
+  (set c (phm/phm-assoc c :other 99))
+  (check "collision: both keys present" (phm/phm-count c) 3)
+  (check "collision: get first" (phm/phm-get c "k6595") :A)
+  (check "collision: get second" (phm/phm-get c "k144747") :B)
+  (check "collision: overwrite one" (phm/phm-get (phm/phm-assoc c "k6595" :A2) "k6595") :A2)
+  (set c (phm/phm-dissoc c "k6595"))
+  (check "collision: dissoc one keeps other" (phm/phm-get c "k144747") :B)
+  (check "collision: dissoc removes" (phm/phm-get c "k6595" :gone) :gone)
+  (check "collision: count after dissoc" (phm/phm-count c) 2))
+
+# --- complexity: assoc must be sub-linear per op (HAMT, not O(n) copy) -------
+# Build 20000 entries; on the old O(n)-copy map this is O(n^2) (~minutes). A HAMT
+# does it in well under a second. Guard generously (5s) to avoid flakiness.
+(def t0 (os/clock))
+(var perf (phm/make-phm))
+(loop [i :range [0 20000]] (set perf (phm/phm-assoc perf i i)))
+(def elapsed (- (os/clock) t0))
+(printf "  20000 assocs: %.3fs" elapsed)
+(check "20000 assocs complete" (phm/phm-count perf) 20000)
+(if (< elapsed 5.0)
+  (print "  ok    assoc is sub-linear (< 5s)")
+  (do (++ fails) (printf "  FAIL  assoc too slow (%.2fs) — O(n) per op?" elapsed)))
+
+(if (> fails 0) (do (printf "phm-hamt: %d FAILED" fails) (os/exit 1))
+  (print "phm-hamt: all passed"))