From 82525b6a81ca50d14f4a277dfa0d04db66c72e2e Mon Sep 17 00:00:00 2001 From: Dmitri Sotnikov Date: Tue, 16 Jun 2026 06:04:46 +0000 Subject: [PATCH] sorted-map/set: red-black tree instead of O(n) sorted vector (jolt-0hbr) (#137) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- jolt-core/clojure/core/25-sorted.clj | 334 ++++++++++++++++------ src/jolt/core_types.janet | 10 +- src/jolt/main.janet | 9 +- test/integration/sorted-rbtree-test.janet | 63 ++++ 4 files changed, 319 insertions(+), 97 deletions(-) create mode 100644 test/integration/sorted-rbtree-test.janet diff --git a/jolt-core/clojure/core/25-sorted.clj b/jolt-core/clojure/core/25-sorted.clj index e7b1483..689590f 100644 --- a/jolt-core/clojure/core/25-sorted.clj +++ b/jolt-core/clojure/core/25-sorted.clj @@ -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 -;; :cmp FN-or-nil ; 3-way comparator; nil = natural order (compare) -;; :ops {op-kw fn}} ; this tier's implementations, attached to the value +;; :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 +;; 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. +;; +;; 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; 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. -;; -;; 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. +;; 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]]))) diff --git a/src/jolt/core_types.janet b/src/jolt/core_types.janet index c23d2d5..229ccaa 100644 --- a/src/jolt/core_types.janet +++ b/src/jolt/core_types.janet @@ -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 diff --git a/src/jolt/main.janet b/src/jolt/main.janet index e326cbc..6e93e63 100644 --- a/src/jolt/main.janet +++ b/src/jolt/main.janet @@ -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 "}")) diff --git a/test/integration/sorted-rbtree-test.janet b/test/integration/sorted-rbtree-test.janet new file mode 100644 index 0000000..da81715 --- /dev/null +++ b/test/integration/sorted-rbtree-test.janet @@ -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!"))