sorted-map/set: red-black tree instead of O(n) sorted vector (jolt-0hbr) (#137)

Sorted collections were a sorted VECTOR — insert-at = (into (conj (subvec es
0 i) x) (subvec es i)) is O(n) per assoc with a large constant, so building was
O(n^2): 2000 entries took 55.6s.

Replace the rep with a red-black tree (assoc/dissoc/get/contains O(log n)),
ported from the ClojureScript PersistentTreeMap (cljs.core: tree-map-add /
balance-left / balance-right / tree-map-append / balance-*-del). This tier (25)
loads before 30-macros so deftype isn't available; a node is a plain vector
[color k v left right] and cljs's BlackNode/RedNode methods become functions —
the algorithm is unchanged. A sorted-set stores elements as keys with a nil
value; its ops project the key.

The seed read the old :entries vector directly for equality/printing; route
those through a new :entries op that materializes ascending from the tree
(core_types/sorted-entries-arr + main.janet's printer).

2000 sorted-map assocs: 55.6s -> 0.98s (57x); now O(log n) (per-op cost flat
from n=2000 to 10000). Correctness in test/integration/sorted-rbtree-test.janet
(shuffled insert ordering, delete rebalancing, custom comparator, comparator
lookup, subseq, count); sorted specs + full gate green. (key/val on sorted
entries stays a pre-existing gap — entries are pvecs not host tuples; jolt-jk23.)

Co-authored-by: Yogthos <yogthos@gmail.com>
This commit is contained in:
Dmitri Sotnikov 2026-06-16 06:04:46 +00:00 committed by GitHub
parent 91e246c682
commit 82525b6a81
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 319 additions and 97 deletions

View file

@ -2,29 +2,30 @@
;;
;; A sorted-map / sorted-set is a tagged host table
;; {:jolt/type :jolt/sorted-map|:jolt/sorted-set
;; :entries VECTOR ; comparator-ordered: [k v] pairs / elements
;; :tree RB-NODE | nil ; a red-black tree, comparator-ordered
;; :cnt N ; element count (O(1))
;; :cmp FN-or-nil ; 3-way comparator; nil = natural order (compare)
;; :ops {op-kw fn}} ; this tier's implementations, attached to the value
;;
;; ALL the semantics live here in Clojure. The Janet seed keeps only its
;; dispatch branches (conj/assoc/get/seq/count/…), each a one-line call through
;; the value's own :ops table — so the ops travel WITH the value (correct
;; across contexts, forks, and AOT images; no module-level hooks to re-wire).
;; The wrapper table itself is minted and read through the minimal host value
;; primitives: jolt.host/tagged-table + jolt.host/ref-put! + jolt.host/ref-get.
;; The tree is a left-leaning-free red-black tree — Rich Hickey's algorithm,
;; ported from the ClojureScript PersistentTreeMap (cljs.core: tree-map-add /
;; balance-left / balance-right / tree-map-append / balance-*-del). assoc / get /
;; dissoc / contains are O(log n); the old sorted-VECTOR rep was O(n) per assoc,
;; O(n^2) to build (jolt-684u's sibling, jolt-0hbr). cljs uses BlackNode/RedNode
;; deftypes, but this tier loads before 30-macros (no deftype), so a node is a
;; plain vector [color k v left right] (color :red/:black; left/right node|nil)
;; and the methods become functions — the algorithm is identical.
;;
;; Clojure semantics this port fixes vs the old Janet kernel: lookup and
;; membership go through the COMPARATOR ((contains? (sorted-set 1) 1.0) was a
;; deep= scan; assoc/conj of a comparator-equal key replaces/no-ops), equality
;; is representation-agnostic ((= (sorted-map :a 1) {:a 1})), empty?/empty see
;; the collection rather than the wrapper, (empty sc) keeps the comparator,
;; iteration (map/reduce/filter) works, and sorted colls canonicalize as map
;; keys. Entries keep the FIRST-inserted key on replace, as Clojure's
;; PersistentTreeMap does.
;; A sorted-SET stores its elements as keys with a nil value; its ops project the
;; key. ALL the semantics live here in Clojure; the Janet seed keeps only its
;; dispatch branches (conj/assoc/get/seq/count/…), each a one-line call through
;; the value's own :ops table, so the ops travel WITH the value (correct across
;; contexts, forks, and AOT images). The wrapper is minted/read through the host
;; value primitives jolt.host/tagged-table + ref-put! + ref-get.
;; Raw field read on the wrapper (host primitive). Plain `get` on a sorted coll
;; IS the comparator lookup — it dispatches back into these ops, so reading
;; :entries/:cmp/:ops with it would recurse forever.
;; :tree/:cmp/:ops with it would recurse forever.
(defn- sfield [sc k] (jolt.host/ref-get sc k))
;; Clojure's fn->comparator: a comparator fn may return a number (3-way) or a
@ -38,48 +39,201 @@
(defn- the-cmp [sc] (or (sfield sc :cmp) compare))
;; Lowest index in [0, n) whose key is >= k under cmp (n when none).
(defn- lower-bound [es keyf cmp k]
(loop [lo 0 hi (count es)]
(if (< lo hi)
(let [mid (quot (+ lo hi) 2)]
(if (neg? (cmp (keyf (nth es mid)) k))
(recur (inc mid) hi)
(recur lo mid)))
lo)))
;; --- red-black tree nodes: [color key val left right] -----------------------
(defn- nd-key [n] (nth n 1))
(defn- nd-val [n] (nth n 2))
(defn- nd-left [n] (nth n 3))
(defn- nd-right [n] (nth n 4))
(defn- red? [n] (and n (identical? :red (nth n 0))))
(defn- black? [n] (and n (identical? :black (nth n 0))))
(defn- mk-red [k v l r] [:red k v l r])
(defn- mk-black [k v l r] [:black k v l r])
;; BlackNode.blacken = self; RedNode.blacken = a black copy.
(defn- blacken [n] (if (red? n) [:black (nd-key n) (nd-val n) (nd-left n) (nd-right n)] n))
;; BlackNode.redden = a red copy; RedNode.redden = invariant violation (never hit
;; on the paths that call it: redden is only applied to a known-black node).
(defn- redden [n] [:red (nd-key n) (nd-val n) (nd-left n) (nd-right n)])
;; replace a node's key/val/children KEEPING its color.
(defn- replace-node [n k v l r] (if (red? n) (mk-red k v l r) (mk-black k v l r)))
;; Index of the comparator-equal entry, or -1.
(defn- find-idx [sc keyf k]
(let [es (sfield sc :entries)
cmp (the-cmp sc)
i (lower-bound es keyf cmp k)]
(if (and (< i (count es)) (zero? (cmp (keyf (nth es i)) k))) i -1)))
;; --- insert balancing (the RedNode/BlackNode .balance-left/.balance-right) ---
(defn- ins-balance-left [ins parent]
(if (red? ins)
(let [l (nd-left ins) r (nd-right ins)]
(cond
(red? l) (mk-red (nd-key ins) (nd-val ins)
(blacken l)
(mk-black (nd-key parent) (nd-val parent) r (nd-right parent)))
(red? r) (mk-red (nd-key r) (nd-val r)
(mk-black (nd-key ins) (nd-val ins) l (nd-left r))
(mk-black (nd-key parent) (nd-val parent) (nd-right r) (nd-right parent)))
:else (mk-black (nd-key parent) (nd-val parent) ins (nd-right parent))))
(mk-black (nd-key parent) (nd-val parent) ins (nd-right parent))))
(defn- make-sorted [tag es cmp ops]
(defn- ins-balance-right [ins parent]
(if (red? ins)
(let [l (nd-left ins) r (nd-right ins)]
(cond
(red? r) (mk-red (nd-key ins) (nd-val ins)
(mk-black (nd-key parent) (nd-val parent) (nd-left parent) l)
(blacken r))
(red? l) (mk-red (nd-key l) (nd-val l)
(mk-black (nd-key parent) (nd-val parent) (nd-left parent) (nd-left l))
(mk-black (nd-key ins) (nd-val ins) (nd-right l) r))
:else (mk-black (nd-key parent) (nd-val parent) (nd-left parent) ins)))
(mk-black (nd-key parent) (nd-val parent) (nd-left parent) ins)))
;; node .add-left / .add-right (parent gains a new left/right subtree `ins`)
(defn- add-left [parent ins]
(if (red? parent)
(mk-red (nd-key parent) (nd-val parent) ins (nd-right parent))
(ins-balance-left ins parent)))
(defn- add-right [parent ins]
(if (red? parent)
(mk-red (nd-key parent) (nd-val parent) (nd-left parent) ins)
(ins-balance-right ins parent)))
;; insert k/v into tree, assuming k is NOT already present (the caller checks).
(defn- tree-ins [cmp tree k v]
(if (nil? tree)
(mk-red k v nil nil)
(if (neg? (cmp k (nd-key tree)))
(add-left tree (tree-ins cmp (nd-left tree) k v))
(add-right tree (tree-ins cmp (nd-right tree) k v)))))
;; replace the value at an existing key, keeping the tree structure (and the
;; first-inserted key, like Clojure's PersistentTreeMap).
(defn- tree-replace [cmp tree k v]
(let [c (cmp k (nd-key tree))]
(cond
(zero? c) (replace-node tree (nd-key tree) v (nd-left tree) (nd-right tree))
(neg? c) (replace-node tree (nd-key tree) (nd-val tree) (tree-replace cmp (nd-left tree) k v) (nd-right tree))
:else (replace-node tree (nd-key tree) (nd-val tree) (nd-left tree) (tree-replace cmp (nd-right tree) k v)))))
(defn- tree-lookup [tree cmp k]
(loop [t tree]
(if (nil? t)
nil
(let [c (cmp k (nd-key t))]
(cond (zero? c) t
(neg? c) (recur (nd-left t))
:else (recur (nd-right t)))))))
;; --- delete balancing (cljs standalone balance-left / balance-right / *-del) -
(defn- balance-left [k v ins right]
(if (red? ins)
(let [il (nd-left ins) ir (nd-right ins)]
(cond
(red? il) (mk-red (nd-key ins) (nd-val ins) (blacken il) (mk-black k v ir right))
(red? ir) (mk-red (nd-key ir) (nd-val ir)
(mk-black (nd-key ins) (nd-val ins) il (nd-left ir))
(mk-black k v (nd-right ir) right))
:else (mk-black k v ins right)))
(mk-black k v ins right)))
(defn- balance-right [k v left ins]
(if (red? ins)
(let [il (nd-left ins) ir (nd-right ins)]
(cond
(red? ir) (mk-red (nd-key ins) (nd-val ins) (mk-black k v left il) (blacken ir))
(red? il) (mk-red (nd-key il) (nd-val il)
(mk-black k v left (nd-left il))
(mk-black (nd-key ins) (nd-val ins) (nd-right il) ir))
:else (mk-black k v left ins)))
(mk-black k v left ins)))
(defn- balance-left-del [k v del right]
(cond
(red? del) (mk-red k v (blacken del) right)
(black? right) (balance-right k v del (redden right))
(and (red? right) (black? (nd-left right)))
(mk-red (nd-key (nd-left right)) (nd-val (nd-left right))
(mk-black k v del (nd-left (nd-left right)))
(balance-right (nd-key right) (nd-val right) (nd-right (nd-left right)) (redden (nd-right right))))
:else (throw (ex-info "red-black tree invariant violation" {}))))
(defn- balance-right-del [k v left del]
(cond
(red? del) (mk-red k v left (blacken del))
(black? left) (balance-left k v (redden left) del)
(and (red? left) (black? (nd-right left)))
(mk-red (nd-key (nd-right left)) (nd-val (nd-right left))
(balance-left (nd-key left) (nd-val left) (redden (nd-left left)) (nd-left (nd-right left)))
(mk-black k v (nd-right (nd-right left)) del))
:else (throw (ex-info "red-black tree invariant violation" {}))))
;; merge two subtrees (the children of a removed node)
(defn- tree-append [left right]
(cond
(nil? left) right
(nil? right) left
(red? left)
(if (red? right)
(let [app (tree-append (nd-right left) (nd-left right))]
(if (red? app)
(mk-red (nd-key app) (nd-val app)
(mk-red (nd-key left) (nd-val left) (nd-left left) (nd-left app))
(mk-red (nd-key right) (nd-val right) (nd-right app) (nd-right right)))
(mk-red (nd-key left) (nd-val left) (nd-left left)
(mk-red (nd-key right) (nd-val right) app (nd-right right)))))
(mk-red (nd-key left) (nd-val left) (nd-left left) (tree-append (nd-right left) right)))
(red? right)
(mk-red (nd-key right) (nd-val right) (tree-append left (nd-left right)) (nd-right right))
:else
(let [app (tree-append (nd-right left) (nd-left right))]
(if (red? app)
(mk-red (nd-key app) (nd-val app)
(mk-black (nd-key left) (nd-val left) (nd-left left) (nd-left app))
(mk-black (nd-key right) (nd-val right) (nd-right app) (nd-right right)))
(balance-left-del (nd-key left) (nd-val left) (nd-left left)
(mk-black (nd-key right) (nd-val right) app (nd-right right)))))))
;; remove k from tree, assuming k IS present (the caller checks).
(defn- tree-del [cmp tree k]
(let [c (cmp k (nd-key tree))]
(cond
(zero? c) (tree-append (nd-left tree) (nd-right tree))
(neg? c) (let [del (tree-del cmp (nd-left tree) k)]
(if (black? (nd-left tree))
(balance-left-del (nd-key tree) (nd-val tree) del (nd-right tree))
(mk-red (nd-key tree) (nd-val tree) del (nd-right tree))))
:else (let [del (tree-del cmp (nd-right tree) k)]
(if (black? (nd-right tree))
(balance-right-del (nd-key tree) (nd-val tree) (nd-left tree) del)
(mk-red (nd-key tree) (nd-val tree) (nd-left tree) del))))))
;; in-order walk: conj (proj node) for each node, ascending.
(defn- tree-collect [t proj acc]
(if (nil? t)
acc
(tree-collect (nd-right t) proj
(conj (tree-collect (nd-left t) proj acc) (proj t)))))
(defn- make-sorted [tag tree cnt cmp ops]
(-> (jolt.host/tagged-table tag)
(jolt.host/ref-put! :entries es)
(jolt.host/ref-put! :tree tree)
(jolt.host/ref-put! :cnt cnt)
(jolt.host/ref-put! :cmp cmp)
(jolt.host/ref-put! :ops ops)))
(defn- insert-at [es i x] (into (conj (subvec es 0 i) x) (subvec es i)))
(defn- remove-at [es i] (into (subvec es 0 i) (subvec es (inc i))))
;; entries as a vector (ascending), the materialized form seq/rseq/subseq use.
(defn- sc-entries [sc proj]
(tree-collect (sfield sc :tree) proj []))
;; --- sorted-map ops ---------------------------------------------------------
(defn- map-entry [t] [(nd-key t) (nd-val t)])
(defn- sm-get [sm k not-found]
(let [i (find-idx sm first k)]
(if (neg? i) not-found (second (nth (sfield sm :entries) i)))))
(let [n (tree-lookup (sfield sm :tree) (the-cmp sm) k)]
(if (nil? n) not-found (nd-val n))))
(defn- sm-assoc-1 [sm k v]
(let [es (sfield sm :entries)
cmp (the-cmp sm)
i (lower-bound es first cmp k)
found (and (< i (count es)) (zero? (cmp (first (nth es i)) k)))]
(make-sorted :jolt/sorted-map
(if found
(assoc es i [(first (nth es i)) v])
(insert-at es i [k v]))
(sfield sm :cmp) (sfield sm :ops))))
(let [cmp (the-cmp sm) tree (sfield sm :tree)
node (tree-lookup tree cmp k)]
(cond
(and node (= v (nd-val node))) sm
node (make-sorted :jolt/sorted-map (tree-replace cmp tree k v) (sfield sm :cnt) (sfield sm :cmp) (sfield sm :ops))
:else (make-sorted :jolt/sorted-map (blacken (tree-ins cmp tree k v)) (inc (sfield sm :cnt)) (sfield sm :cmp) (sfield sm :ops)))))
(defn- sm-assoc-many [sm kvs]
(let [n (count kvs)]
@ -90,14 +244,14 @@
(recur (sm-assoc-1 m (nth kvs i) (nth kvs (inc i))) (+ i 2))
m))))
(defn- sm-dissoc-many [sm ks]
(reduce (fn [m k]
(let [i (find-idx m first k)]
(if (neg? i)
m
(make-sorted :jolt/sorted-map (remove-at (sfield m :entries) i)
(sfield m :cmp) (sfield m :ops)))))
sm ks))
(defn- sm-dissoc-1 [sm k]
(let [cmp (the-cmp sm) tree (sfield sm :tree)]
(if (nil? (tree-lookup tree cmp k))
sm
(let [t (tree-del cmp tree k)]
(make-sorted :jolt/sorted-map (when t (blacken t)) (dec (sfield sm :cnt)) (sfield sm :cmp) (sfield sm :ops))))))
(defn- sm-dissoc-many [sm ks] (reduce sm-dissoc-1 sm ks))
;; conj on a map: a [k v] pair (2-vector / map-entry) or a map to merge;
;; nil is a no-op, as in Clojure.
@ -110,69 +264,68 @@
(defn- sm-conj-many [sm xs] (reduce sm-conj-1 sm xs))
;; --- sorted-set ops ---------------------------------------------------------
;; --- sorted-set ops (elements stored as keys, nil value) --------------------
(defn- ss-get [ss x not-found]
(let [i (find-idx ss identity x)]
(if (neg? i) not-found (nth (sfield ss :entries) i))))
(let [n (tree-lookup (sfield ss :tree) (the-cmp ss) x)]
(if (nil? n) not-found (nd-key n))))
(defn- ss-conj-1 [ss x]
(let [es (sfield ss :entries)
cmp (the-cmp ss)
i (lower-bound es identity cmp x)]
(if (and (< i (count es)) (zero? (cmp (nth es i) x)))
(let [cmp (the-cmp ss) tree (sfield ss :tree)]
(if (tree-lookup tree cmp x)
ss
(make-sorted :jolt/sorted-set (insert-at es i x) (sfield ss :cmp) (sfield ss :ops)))))
(make-sorted :jolt/sorted-set (blacken (tree-ins cmp tree x nil)) (inc (sfield ss :cnt)) (sfield ss :cmp) (sfield ss :ops)))))
(defn- ss-conj-many [ss xs] (reduce ss-conj-1 ss xs))
(defn- ss-disj-many [ss xs]
(reduce (fn [s x]
(let [i (find-idx s identity x)]
(if (neg? i)
s
(make-sorted :jolt/sorted-set (remove-at (sfield s :entries) i)
(sfield s :cmp) (sfield s :ops)))))
ss xs))
(defn- ss-disj-1 [ss x]
(let [cmp (the-cmp ss) tree (sfield ss :tree)]
(if (nil? (tree-lookup tree cmp x))
ss
(let [t (tree-del cmp tree x)]
(make-sorted :jolt/sorted-set (when t (blacken t)) (dec (sfield ss :cnt)) (sfield ss :cmp) (sfield ss :ops))))))
(defn- ss-disj-many [ss xs] (reduce ss-disj-1 ss xs))
;; --- the ops tables the Janet seed dispatches through ------------------------
(def ^:private sm-ops
{:count (fn [sm] (count (sfield sm :entries)))
:seq (fn [sm] (seq (sfield sm :entries)))
:rseq (fn [sm] (seq (vec (reverse (sfield sm :entries)))))
:first (fn [sm] (first (sfield sm :entries)))
{:count (fn [sm] (sfield sm :cnt))
:entries (fn [sm] (sc-entries sm map-entry))
:seq (fn [sm] (seq (sc-entries sm map-entry)))
:rseq (fn [sm] (seq (vec (reverse (sc-entries sm map-entry)))))
:first (fn [sm] (first (sc-entries sm map-entry)))
:get sm-get
:contains (fn [sm k] (not (neg? (find-idx sm first k))))
:contains (fn [sm k] (not (nil? (tree-lookup (sfield sm :tree) (the-cmp sm) k))))
:assoc sm-assoc-many
:dissoc sm-dissoc-many
:conj sm-conj-many
:empty (fn [sm] (make-sorted :jolt/sorted-map [] (sfield sm :cmp) (sfield sm :ops)))})
:empty (fn [sm] (make-sorted :jolt/sorted-map nil 0 (sfield sm :cmp) (sfield sm :ops)))})
(def ^:private ss-ops
{:count (fn [ss] (count (sfield ss :entries)))
:seq (fn [ss] (seq (sfield ss :entries)))
:rseq (fn [ss] (seq (vec (reverse (sfield ss :entries)))))
:first (fn [ss] (first (sfield ss :entries)))
{:count (fn [ss] (sfield ss :cnt))
:entries (fn [ss] (sc-entries ss nd-key))
:seq (fn [ss] (seq (sc-entries ss nd-key)))
:rseq (fn [ss] (seq (vec (reverse (sc-entries ss nd-key)))))
:first (fn [ss] (first (sc-entries ss nd-key)))
:get ss-get
:contains (fn [ss x] (not (neg? (find-idx ss identity x))))
:contains (fn [ss x] (not (nil? (tree-lookup (sfield ss :tree) (the-cmp ss) x))))
:conj ss-conj-many
:disj ss-disj-many
:empty (fn [ss] (make-sorted :jolt/sorted-set [] (sfield ss :cmp) (sfield ss :ops)))})
:empty (fn [ss] (make-sorted :jolt/sorted-set nil 0 (sfield ss :cmp) (sfield ss :ops)))})
;; --- constructors + predicates -----------------------------------------------
(defn sorted-map [& kvs]
(sm-assoc-many (make-sorted :jolt/sorted-map [] nil sm-ops) (vec kvs)))
(sm-assoc-many (make-sorted :jolt/sorted-map nil 0 nil sm-ops) (vec kvs)))
(defn sorted-map-by [comparator & kvs]
(sm-assoc-many (make-sorted :jolt/sorted-map [] (fn->cmp comparator) sm-ops) (vec kvs)))
(sm-assoc-many (make-sorted :jolt/sorted-map nil 0 (fn->cmp comparator) sm-ops) (vec kvs)))
(defn sorted-set [& xs]
(ss-conj-many (make-sorted :jolt/sorted-set [] nil ss-ops) (vec xs)))
(ss-conj-many (make-sorted :jolt/sorted-set nil 0 nil ss-ops) (vec xs)))
(defn sorted-set-by [comparator & xs]
(ss-conj-many (make-sorted :jolt/sorted-set [] (fn->cmp comparator) ss-ops) (vec xs)))
(ss-conj-many (make-sorted :jolt/sorted-set nil 0 (fn->cmp comparator) ss-ops) (vec xs)))
(defn sorted-map? [x] (= :jolt/sorted-map (sfield x :jolt/type)))
(defn sorted-set? [x] (= :jolt/sorted-set (sfield x :jolt/type)))
@ -184,13 +337,14 @@
;; nil, like Clojure.
(defn- sc-keyf [sc] (if (sorted-map? sc) first identity))
(defn- sc-proj [sc] (if (sorted-map? sc) map-entry nd-key))
(defn- sub-filter [sc tests]
(let [cmp (the-cmp sc)
keyf (sc-keyf sc)]
(filterv (fn [e]
(every? (fn [[test k]] (test (cmp (keyf e) k) 0)) tests))
(sfield sc :entries))))
(sc-entries sc (sc-proj sc)))))
(defn subseq
([sc test k] (seq (sub-filter sc [[test k]])))

View file

@ -57,10 +57,14 @@
(defn core-sorted-map? [x] (and (table? x) (= :jolt/sorted-map (x :jolt/type))))
(defn core-sorted-set? [x] (and (table? x) (= :jolt/sorted-set (x :jolt/type))))
(defn core-sorted? [x] (or (core-sorted-map? x) (core-sorted-set? x)))
# The :entries vector as a Janet array (entries are jolt vectors: pvecs in
# immutable mode, arrays in mutable mode) — for the seed's printers/equality.
# The comparator-ordered entries as a Janet array (entries are jolt vectors:
# pvecs in immutable mode, arrays in mutable mode) — for the seed's printers/
# equality. Materialized from the red-black tree via the coll's own :entries op
# (jolt-0hbr); the old sorted-VECTOR rep is read directly as a fallback.
(defn sorted-entries-arr [coll]
(let [e (coll :entries)] (if (pvec? e) (pv->array e) e)))
(def ef (let [ops (coll :ops)] (and ops (ops :entries))))
(def e (if ef (ef coll) (coll :entries)))
(if (pvec? e) (pv->array e) e))
# Lazy cell chain over an indexed (tuple/array) collection, walking by INDEX —
# O(1) per step. Slicing the remainder per step (the old shape) made every

View file

@ -133,13 +133,14 @@
(and (table? v) (= :jolt/regex (v :jolt/type)))
(do (push-str buf "#\"") (push-str buf (v :source)) (push-str buf "\""))
# sorted colls: their comparator-ordered :entries vector (a pvec in
# immutable mode, an array in mutable mode) is all the printer reads.
# sorted colls: their comparator-ordered entries, materialized from the
# red-black tree via the value's own :entries op (jolt-0hbr), is all the
# printer reads.
(and (table? v) (= :jolt/sorted-map (v :jolt/type)))
(do
(push-str buf "{")
(var first? true)
(each e (let [es (v :entries)] (if (pvec? es) (pv->array es) es))
(each e (let [ef (let [o (v :ops)] (and o (o :entries))) es (if ef (ef v) (v :entries))] (if (pvec? es) (pv->array es) es))
(if first? (set first? false) (push-str buf ", "))
(write-value (if (pvec? e) (pv-nth e 0) (in e 0)) buf)
(push-str buf " ")
@ -150,7 +151,7 @@
(do
(push-str buf "#{")
(var first? true)
(each x (let [es (v :entries)] (if (pvec? es) (pv->array es) es))
(each x (let [ef (let [o (v :ops)] (and o (o :entries))) es (if ef (ef v) (v :entries))] (if (pvec? es) (pv->array es) es))
(if first? (set first? false) (push-str buf " "))
(write-value x buf))
(push-str buf "}"))

View file

@ -0,0 +1,63 @@
# sorted-map / sorted-set are a red-black tree (jolt-0hbr), ported from the
# ClojureScript PersistentTreeMap. assoc/dissoc are O(log n); the old sorted-
# vector rep was O(n) per op (O(n^2) to build — 55s for 2000 entries). This
# drives big shuffled insert/delete sequences (which stress rebalancing) through
# the built binary and checks ordering + a sub-quadratic build time.
(import ../../src/jolt/api :as api)
(print "sorted-map/set red-black tree (jolt-0hbr)...")
(os/setenv "JOLT_DIRECT_LINK" "1")
(def ctx (api/init {:compile? true}))
(defn ev [s] (api/eval-string ctx s))
(var fails 0)
(defn check [label got expected]
(if (= got expected) (print " ok " label)
(do (++ fails) (printf " FAIL %s: want %p got %p" label expected got))))
# --- ordering is maintained through heavy rebalancing -----------------------
(check "shuffled 500 keys come back sorted"
(ev "(= (vec (keys (apply sorted-map (interleave (shuffle (vec (range 500))) (range 500))))) (vec (range 500)))")
true)
(check "sorted-set of shuffled 500 is sorted"
(ev "(= (vec (apply sorted-set (shuffle (vec (range 500))))) (vec (range 500)))")
true)
# --- delete maintains order + correctness (stresses delete rebalancing) ------
(check "dissoc every even key, odds remain in order"
(ev "(let [m (apply sorted-map (interleave (shuffle (vec (range 200))) (range 200)))
m2 (reduce dissoc m (range 0 200 2))]
(= (vec (keys m2)) (vec (range 1 200 2))))")
true)
(check "disj down to empty then rebuild"
(ev "(let [s (apply sorted-set (range 100))
s2 (reduce disj s (shuffle (vec (range 100))))]
(and (= 0 (count s2)) (= [1 2 3] (vec (conj s2 3 1 2 1)))))")
true)
# --- comparator + lookup correctness ----------------------------------------
(check "custom comparator (descending)"
(ev "(= (vec (keys (sorted-map-by > 1 :a 3 :c 2 :b))) [3 2 1])") true)
(check "get/contains go through comparator (1 vs 1.0)"
(ev "(and (contains? (sorted-set 1 2 3) 1.0) (= :a (get (sorted-map 1 :a) 1.0)))") true)
(check "first-inserted key kept on value replace"
(ev "(first (first (assoc (sorted-map 1 :a) 1.0 :b)))") 1)
# --- count + subseq ----------------------------------------------------------
(check "count after mixed ops"
(ev "(count (-> (apply sorted-map (interleave (range 50) (range 50))) (dissoc 10 20 30) (assoc 100 1 101 2)))") 49)
(check "subseq range"
(ev "(= (vec (map first (subseq (apply sorted-map (interleave (range 20) (range 20))) >= 5 < 9))) [5 6 7 8])") true)
# --- complexity: building a big map must be sub-quadratic --------------------
(def t0 (os/clock))
(ev "(count (loop [i 0 m (sorted-map)] (if (< i 5000) (recur (inc i) (assoc m (mod (* i 7919) 10007) i)) m)))")
(def elapsed (- (os/clock) t0))
(printf " 5000 assocs: %.2fs" elapsed)
(if (< elapsed 8.0)
(print " ok sorted assoc is sub-quadratic (< 8s)")
(do (++ fails) (printf " FAIL sorted build too slow (%.1fs) — O(n^2)?" elapsed)))
(if (> fails 0) (do (printf "sorted-rbtree: %d FAILED" fails) (os/exit 1))
(print "sorted-rbtree (jolt-0hbr) passed!"))