From 58f019bc14367e09193a625c65c07dff4c052df1 Mon Sep 17 00:00:00 2001
From: giles <giles.bradshaw@rose-ash.com>
Date: Mon, 11 May 2026 22:26:45 +0000
Subject: [PATCH] =?UTF-8?q?JIT:=20lib/jit.sx=20=E2=80=94=20SX-level=20conv?=
 =?UTF-8?q?enience=20layer?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three primitives + a wrapper, all portable across hosts:

  with-jit-threshold N body...  — temporarily set threshold, restore on exit
  with-jit-budget    N body...  — temporarily set LRU budget
  with-fresh-jit       body...  — clear cache before & after body

  jit-report                    — human-readable stats string for logging
  jit-disable!  / jit-enable!   — convenience around set-budget! 0

The host (OCaml here, will be JS/Python eventually) only needs to provide
the underlying primitives (jit-stats, jit-set-threshold!, jit-set-budget!,
jit-reset-cache!, jit-reset-counters!). The ergonomics live in shared SX.

Used together with Phase 1 (tiered compilation) and Phase 2 (LRU eviction)
to give application developers fine-grained control over the JIT cache:
isolated test runs use with-fresh-jit, hot benchmark sections use
with-jit-threshold 1, memory-constrained pages use jit-set-budget! to
cap the cache.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 lib/jit.sx | 89 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)
 create mode 100644 lib/jit.sx

diff --git a/lib/jit.sx b/lib/jit.sx
new file mode 100644
index 00000000..7a0b8fca
--- /dev/null
+++ b/lib/jit.sx
@@ -0,0 +1,89 @@
+;; lib/jit.sx — SX-level convenience wrappers over the JIT cache control
+;; primitives (jit-stats, jit-set-threshold!, jit-set-budget!, jit-reset-cache!,
+;; jit-reset-counters!). Host-specific implementations live in
+;; hosts/<host>/lib/sx_*.ml; the API surface is portable across hosts.
+
+;; with-jit-threshold — temporarily set the JIT call-count threshold for
+;; the duration of body, restoring the previous value on exit. Useful for
+;; sections that want eager compilation (threshold=1) or want to skip JIT
+;; entirely (threshold=999999) for diagnostic comparison.
+(defmacro
+  with-jit-threshold
+  (n &rest body)
+  `(let
+     ((__old (get (jit-stats) "threshold")))
+     (jit-set-threshold! ,n)
+     (let
+       ((__r (do ,@body)))
+       (jit-set-threshold! __old)
+       __r)))
+
+;; with-jit-budget — temporarily set the LRU cache budget. Setting to 0
+;; disables JIT entirely (everything falls through to the interpreter);
+;; large values are effectively unbounded.
+(defmacro
+  with-jit-budget
+  (n &rest body)
+  `(let
+     ((__old (get (jit-stats) "budget")))
+     (jit-set-budget! ,n)
+     (let
+       ((__r (do ,@body)))
+       (jit-set-budget! __old)
+       __r)))
+
+;; with-fresh-jit — clear the cache before body, run body, clear again
+;; after. Use between sessions / request batches / test suites where you
+;; want deterministic timing free of carryover.
+(defmacro
+  with-fresh-jit
+  (&rest body)
+  `(let
+     ((__r (do (jit-reset-cache!) ,@body)))
+     (jit-reset-cache!)
+     __r))
+
+;; jit-report — human-readable summary of current JIT state. Returns a
+;; string suitable for logging.
+(define
+  jit-report
+  (fn
+    ()
+    (let
+      ((s (jit-stats)))
+      (let
+        ((compiled (get s "compiled"))
+          (skipped (get s "below-threshold"))
+          (failed (get s "compile-failed"))
+          (evicted (get s "evicted"))
+          (cache-size (get s "cache-size"))
+          (budget (get s "budget"))
+          (threshold (get s "threshold")))
+        (let
+          ((total (+ compiled skipped failed)))
+          (str
+            "jit: " cache-size "/" budget " cached "
+            "(thr=" threshold ") · "
+            compiled " compiled, "
+            skipped " below-thr, "
+            failed " failed, "
+            evicted " evicted "
+            "(" (if (> total 0) (* 100 (/ compiled total)) 0) "% compile rate)"))))))
+
+;; jit-disable! / jit-enable! — convenience helpers. Disabling sets budget
+;; to 0 which causes the VM to skip JIT entirely on the next call. Enable
+;; restores the budget to its previous value (or 5000 if no previous).
+(define _jit-saved-budget (list 5000))
+
+(define
+  jit-disable!
+  (fn
+    ()
+    (set! _jit-saved-budget (list (get (jit-stats) "budget")))
+    (jit-set-budget! 0)))
+
+(define
+  jit-enable!
+  (fn
+    ()
+    (jit-set-budget! (first _jit-saved-budget))))