Merge pull request #91 from jolt-lang/perf-maps-math

perf: inline keyword lookup + map literals, clojure.math, indexed reduce

src/jolt/api.janet

@@ -124,6 +124,35 @@
   # exists (50-io tier), so it rides the end of overlay load.
   (install-print-method-cb! ctx))
 
+# clojure.math (Clojure 1.11) backed directly by Janet's math natives
+# (jolt-h79). The vars hold plain Janet functions, so compiled calls
+# direct-link — unlike Math/sqrt interop forms, which are in the frozen
+# interpret-only punt set (~5us/call vs ~30ns here). Installed as a
+# populated namespace, so (require '[clojure.math :as m]) is a no-op pass.
+(defn- install-clojure-math! [ctx]
+  (def ns (ctx-find-ns ctx "clojure.math"))
+  (def fns
+    {"sqrt" math/sqrt "cbrt" math/cbrt "pow" math/pow
+     "exp" math/exp "expm1" math/expm1
+     "log" math/log "log10" math/log10 "log1p" math/log1p
+     "sin" math/sin "cos" math/cos "tan" math/tan
+     "asin" math/asin "acos" math/acos "atan" math/atan "atan2" math/atan2
+     "sinh" math/sinh "cosh" math/cosh "tanh" math/tanh
+     "floor" math/floor "ceil" math/ceil "rint" math/round
+     "round" (fn cm-round [x] (math/round x))
+     "signum" (fn cm-signum [x] (cond (< x 0) -1.0 (> x 0) 1.0 0.0))
+     "to-degrees" (fn cm-to-degrees [r] (/ (* r 180.0) math/pi))
+     "to-radians" (fn cm-to-radians [d] (/ (* d math/pi) 180.0))
+     "hypot" (fn cm-hypot [a b] (math/sqrt (+ (* a a) (* b b))))
+     "floor-div" (fn cm-floor-div [a b] (math/floor (/ a b)))
+     "floor-mod" (fn cm-floor-mod [a b] (mod a b))
+     "E" math/e "PI" math/pi})
+  (eachp [nm f] fns (ns-intern ns nm f))
+  # mark loaded so maybe-require-ns never goes looking for a source file
+  (def loaded (or (get (ctx :env) :loaded-namespaces)
+                  (let [t @{}] (put (ctx :env) :loaded-namespaces t) t)))
+  (put loaded "clojure.math" true))
+
 (defn init
   "Create a new Jolt evaluation context.
   opts may contain:
@@ -154,6 +183,7 @@
     (install-async! ctx)
     # Host contract (ns jolt.host): the seam the portable jolt-core compiler calls.
     (host/install! ctx)
+    (install-clojure-math! ctx)
     # require/maybe-require-ns route loaded namespaces through the loader's
     # compile-or-interpret eval-toplevel (the evaluator can't import the loader
     # — that would be circular — so it reads this hook). Without it, required

src/jolt/backend.janet

@@ -250,9 +250,16 @@
    # jolt's bit fns are 2-arg (unlike Clojure's variadic), so these emit native
    # only at exactly the arity the interpreted fn accepts; bit-not is unary.
    "bit-and" 'band "bit-or" 'bor "bit-xor" 'bxor
-   "bit-shift-left" 'blshift "bit-shift-right" 'brshift "bit-not" 'bnot})
+   "bit-shift-left" 'blshift "bit-shift-right" 'brshift "bit-not" 'bnot
+   # janet min/max are variadic with Clojure's numeric semantics; nil?/some?
+   # lower to janet's fastfun = / not= against nil (pure opcodes), and `not`
+   # to janet not — all hot in predicate-heavy loops (jolt-4vr). Same
+   # documented numbers-only relaxation as the arithmetic ops above.
+   "min" 'min "max" 'max
+   "nil?" 'jolt-nil? "some?" 'jolt-some? "not" 'not})
 
-(def- unary-ops {'++ true '-- true 'bnot true})
+(def- unary-ops {'++ true '-- true 'bnot true
+                 'jolt-nil? true 'jolt-some? true 'not true})
 (def- binary-ops {'mod true '% true 'band true 'bor true 'bxor true
                   'blshift true 'brshift true})
 
@@ -270,8 +277,15 @@
     (nil? op) nil
     (and (get unary-ops op) (not= nargs 1)) nil
     (and (get binary-ops op) (not= nargs 2)) nil
+    (and (or (= op 'min) (= op 'max)) (= nargs 0)) nil
     op))
 
+# Janet-level gensym for the inline fast paths: (use ./core) shadows janet's
+# gensym with jolt's (which returns a jolt symbol STRUCT — useless as a janet
+# binding target). _fp$ mirrors the reserved _r$ compiler prefix.
+(var- fp-counter 0)
+(defn- jsym [] (symbol "_fp$" (++ fp-counter)))
+
 (defn- emit-invoke [ctx node]
   (def fnode (norm-node (node :fn)))
   (def args (map |(emit ctx $) (vview (node :args))))
@@ -280,7 +294,41 @@
     nop (case nop
           '++ ['+ (in args 0) 1]
           '-- ['- (in args 0) 1]
+          'jolt-nil? ['= nil (in args 0)]
+          'jolt-some? ['not= nil (in args 0)]
           (tuple nop ;args))
+    # (:kw m) / (:kw m default) — inline the lookup (jank-style, jolt-4vr).
+    # The guard is (get m :jolt/type): janet compiles `get` to an opcode
+    # (~17ns) where a struct?-style cfunction predicate costs ~85ns/lookup.
+    # :jolt/type is a reserved key — user map literals can't contain it (the
+    # reader treats such maps as tagged forms) — and every table-backed rep
+    # that must NOT be raw-indexed carries it (phm — tagged for this guard —
+    # sorted, transient, pvec, atoms, lazy-seqs), so a non-nil tag routes to
+    # core-get's full semantics. Everything else (structs = literal maps,
+    # records with direct field keys, nil, janet arrays, scalars) gets janet
+    # `get` semantics, which match core-get for keyword keys. Structs never
+    # store nil values (nil values force the phm rep), so present-but-nil
+    # can't be confused with missing on the fast arm.
+    (and (= :const (fnode :op)) (keyword? (fnode :val))
+         (>= 2 (length args) 1))
+    (let [k (fnode :val)
+          m-expr (in args 0)
+          # when the subject is already a janet symbol (a local), read it
+          # directly — the guard + lookup both reference it, and locals are
+          # immutable reads, so no rebinding let is needed (saves a binding
+          # per lookup in exactly the hottest shape, (:k local))
+          m (if (symbol? m-expr) m-expr (jsym))
+          wrap (fn [body] (if (symbol? m-expr) body ['let [m m-expr] body]))]
+      (if (= 1 (length args))
+        (wrap ['if ['get m :jolt/type] (tuple core-get m k) ['get m k]])
+        (let [d-expr (in args 1)
+              d (if (symbol? d-expr) d-expr (jsym))
+              v (jsym)
+              body ['if ['get m :jolt/type]
+                    (tuple core-get m k d)
+                    ['let [v ['get m k]] ['if ['nil? v] d v]]]
+              body (if (symbol? d-expr) body ['let [d d-expr] body])]
+          (wrap body))))
     (direct-call? ctx fnode) (tuple (emit ctx fnode) ;args)
     (tuple jolt-call (emit ctx fnode) ;args)))
 
@@ -289,12 +337,50 @@
   (tuple make-vec (tuple/slice (array/concat @['tuple] items))))
 
 (defn- emit-map [ctx node]
-  (def args @[build-map-literal])
-  (each pair (vview (node :pairs))
-    (def p (vview pair))
-    (array/push args (emit ctx (in p 0)))
-    (array/push args (emit ctx (in p 1))))
-  (tuple/slice args))
+  (def pairs (vview (node :pairs)))
+  # Fast path (jolt-4vr): when every key is a scalar const (keyword/string/
+  # number/bool — never a collection, so value-hashing can't be needed from
+  # the keys), construct the Janet struct inline with one nil-check per
+  # value instead of calling variadic build-map-literal and re-scanning the
+  # kvs at runtime. A nil value still falls back to the phm rep (Clojure
+  # keeps nil entries; structs drop them).
+  (var fast (> (length pairs) 0))
+  (each pair pairs
+    (def k (norm-node (in (vview pair) 0)))
+    (def kv (get k :val))
+    (unless (and (= :const (k :op))
+                 (or (keyword? kv) (string? kv) (number? kv) (boolean? kv)))
+      (set fast false)))
+  (if fast
+    (do
+      (def binds @[])
+      (def skvs @['struct])
+      (def phm-args @[build-map-literal])
+      (def truthy @['and])
+      (each pair pairs
+        (def p (vview pair))
+        (def kk ((norm-node (in p 0)) :val))
+        (def vs (jsym))
+        (array/push binds vs)
+        (array/push binds (emit ctx (in p 1)))
+        (array/push truthy vs)
+        (array/push skvs kk) (array/push skvs vs)
+        (array/push phm-args kk) (array/push phm-args vs))
+      # `and` is pure branch opcodes, so the all-truthy common case pays no
+      # predicate calls at all. nil OR false values (rare) drop to
+      # build-map-literal, which re-checks nil properly (false values come
+      # back out on the struct arm there; nil values get the phm rep).
+      ['let (tuple/slice binds)
+       ['if (tuple/slice truthy)
+        (tuple/slice skvs)
+        (tuple/slice phm-args)]])
+    (do
+      (def args @[build-map-literal])
+      (each pair pairs
+        (def p (vview pair))
+        (array/push args (emit ctx (in p 0)))
+        (array/push args (emit ctx (in p 1))))
+      (tuple/slice args))))
 
 # A set literal: build (make-phs e1 e2 …) so each element is evaluated at runtime
 # then the persistent set is constructed — mirrors compiler.janet's emit-set-expr.

src/jolt/core.janet

@@ -880,10 +880,27 @@
                   (if (nil? rt) (set go false) (set cur (ls-rest-cached cur rt))))))))))
     (do
       (var stop false)
-      (each x (if (set? coll) (phs-seq coll) (realize-for-iteration coll))
-        (when (not stop)
-          (set acc (rf acc x))
-          (when (core-reduced? acc) (set acc (acc :val)) (set stop true))))))
+      (cond
+        # Indexed colls iterate in place — realize-for-iteration would copy a
+        # pvec into a fresh array (alloc + pv-nth per element) on EVERY
+        # reduce call, which dominates tight reduce-over-vector loops
+        # (jolt-4vr). Also breaks at `reduced` instead of scanning the tail.
+        (pvec? coll)
+        (do (def n (pv-count coll)) (var i 0)
+            (while (and (< i n) (not stop))
+              (set acc (rf acc (pv-nth coll i)))
+              (when (core-reduced? acc) (set acc (acc :val)) (set stop true))
+              (++ i)))
+        (or (tuple? coll) (array? coll))
+        (do (def n (length coll)) (var i 0)
+            (while (and (< i n) (not stop))
+              (set acc (rf acc (in coll i)))
+              (when (core-reduced? acc) (set acc (acc :val)) (set stop true))
+              (++ i)))
+        (each x (if (set? coll) (phs-seq coll) (realize-for-iteration coll))
+          (when (not stop)
+            (set acc (rf acc x))
+            (when (core-reduced? acc) (set acc (acc :val)) (set stop true)))))))
   acc)
 
 (defn- transduce-reduce

src/jolt/evaluator.janet

@@ -106,6 +106,11 @@
   [coll k default]
   (cond
     (jolt-transient? coll) (transient-lookup coll k default)
+    # sorted colls are tables — without this arm they fell into the raw
+    # table-get branch and (:k (sorted-map ...)) was always nil (jolt-4vr spec)
+    (and (table? coll) (or (= :jolt/sorted-map (coll :jolt/type))
+                           (= :jolt/sorted-set (coll :jolt/type))))
+      ((get (coll :ops) :get) coll k default)
     (phm? coll) (phm-get coll k default)
     (set? coll) (if (phs-contains? coll k) k default)
     (pvec? coll)

src/jolt/phm.janet

@@ -121,7 +121,7 @@
     (def grown (if (> new-cnt (* 2 nbuckets))
                  (rehash new-buckets (* 2 nbuckets))
                  new-buckets))
-    @{:jolt/deftype "jolt.lang.persistent-hash-map.PersistentHashMap"
+    @{: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]
@@ -135,7 +135,7 @@
             (var bi 0)
             (while (< bi nbuckets)
               (put new-buckets bi (if (= bi idx) new-bucket (get (m :buckets) bi))) (++ bi))
-            @{:jolt/deftype "jolt.lang.persistent-hash-map.PersistentHashMap"
+            @{:jolt/type :jolt/phm :jolt/deftype "jolt.lang.persistent-hash-map.PersistentHashMap"
               :cnt new-cnt :buckets new-buckets :_meta (m :_meta)})))
       m)))
 
@@ -167,7 +167,7 @@
 
 (defn make-phm [&opt kvs]
   (default kvs nil)
-  (var m @{:jolt/deftype "jolt.lang.persistent-hash-map.PersistentHashMap"
+  (var m @{:jolt/type :jolt/phm :jolt/deftype "jolt.lang.persistent-hash-map.PersistentHashMap"
            :cnt 0 :buckets (array/new-filled initial-buckets nil) :_meta nil})
   (when kvs
     (var i 0) (var n (length kvs))

test/spec/map-fastpath-spec.janet

@@ -0,0 +1,57 @@
+# Specification: keyword-invoke and map-literal semantics that the compiled
+# fast paths (jolt-4vr: inlined keyword lookup, inlined struct construction)
+# must preserve. Written BEFORE the optimization; every row passes on the
+# generic jolt-call/build-map-literal paths and must keep passing on the
+# inlined ones.
+(use ../support/harness)
+
+(defspec "maps / keyword invoke"
+  ["hit"                 "1"     "(:a {:a 1 :b 2})"]
+  ["miss"                "nil"   "(:z {:a 1})"]
+  ["miss with default"   ":d"    "(:z {:a 1} :d)"]
+  ["hit with default"    "1"     "(:a {:a 1} :d)"]
+  ["on nil"              "nil"   "(:a nil)"]
+  ["on nil with default" ":d"    "(:a nil :d)"]
+  ["nil value is present" "nil"  "(:a {:a nil} :d)"]
+  ["false value is present" "false" "(:a {:a false} :d)"]
+  ["on a vector"         "nil"   "(:a [1 2 3])"]
+  ["on a number"         "nil"   "(:a 42)"]
+  ["on a sorted map"     "2"     "(:b (sorted-map :a 1 :b 2))"]
+  ["on assoc result"     "3"     "(:c (assoc {:a 1} :c 3))"]
+  ["on a record field"   "5"     "(do (defrecord KFP [x]) (:x (->KFP 5)))"]
+  ["qualified keyword"   "1"     "(:n/a {:n/a 1})"]
+  ["nested in expr"      "6"     "(+ (:a {:a 1}) (:b {:b 2}) (:c {:c 3}))"]
+  ["evaluates map expr once" "[2 1]"
+   "(do (def cnt (atom 0)) (let [v (:a (do (swap! cnt inc) {:a 2}))] [v @cnt]))"])
+
+(defspec "maps / literal construction"
+  ["basic"            "{:a 1, :b 2}"  "{:a 1 :b 2}"]
+  ["empty"            "{}"            "{}"]
+  ["computed values"  "{:a 3}"        "{:a (+ 1 2)}"]
+  ["nil value kept"   "true"          "(contains? {:a nil} :a)"]
+  ["nil value lookup" "nil"           "(get {:a nil} :a :d)"]
+  ["string key"       "1"             "(get {\"k\" 1} \"k\")"]
+  ["number key"       ":one"          "(get {1 :one} 1)"]
+  ["collection key"   ":v"            "(get {[1 2] :v} [1 2])"]
+  ["collection value-equal key" ":v"  "(get {[1 2] :v} (vector 1 2))"]
+  ["computed key"     "1"             "(get {(keyword \"a\") 1} :a)"]
+  # NOTE: entry ORDER is reader-hash order today, not source order (Clojure
+  # evaluates array-map literals left-to-right) — pinned loosely; see jolt-p3c
+  ["values evaluate exactly once each" "[1 2 3]"
+   "(do (def log (atom [])) {:a (swap! log conj 1) :b (swap! log conj 2) :c (swap! log conj 3)} (vec (sort (deref log))))"]
+  ["count"            "3"             "(count {:a 1 :b 2 :c 3})"]
+  ["equality with phm" "true"         "(= {:a 1 :b 2} (assoc {:a 1} :b 2))"]
+  ["keys work after assoc" "2"        "(:b (assoc {:a 1 :b 2} :c 3))"]
+  ["literal in fn body"  "12"         "(do (defn mfp-mk [x] {:v (* x 2)}) (:v (mfp-mk 6)))"])
+
+(defspec "clojure.math (jolt-h79)"
+  ["sqrt"        "true"  "(< 1.4142 (clojure.math/sqrt 2) 1.4143)"]
+  ["pow"         "1024"  "(long (clojure.math/pow 2 10))"]
+  ["tan of 0"    "0"     "(long (clojure.math/tan 0))"]
+  ["round"       "3"                "(clojure.math/round 2.6)"]
+  ["floor"       "2"                "(clojure.math/floor 2.9)"]
+  ["signum"      "-1"               "(clojure.math/signum -7.2)"]
+  ["to-radians"  "true"             "(< 3.14 (clojure.math/to-radians 180) 3.15)"]
+  ["PI"          "true"             "(< 3.14 clojure.math/PI 3.15)"]
+  ["require + alias" "5"            "(do (require '[clojure.math :as m]) (long (m/hypot 3 4)))"]
+  ["as a value"  "[1 2]"            "(mapv (comp long clojure.math/sqrt) [1 4])"])