Scope whole-program optimization to app namespaces (jolt-87e) (#126)

Under JOLT_OPTIMIZE a -m program run inferred + specialized EVERY loaded
namespace, including every transitive dependency. On a dep-heavy app that's
prohibitive: malli-app cold-started in ~2m10s (hundreds of dep namespaces, each
run through the per-form inline + inference passes).

The closed world a whole-program pass reasons over is the APP, not its
libraries. jolt-deps now passes the project's own source roots (its deps.edn
:paths) to the runtime as JOLT_APP_PATHS. A namespace loaded from an app root
gets full optimization (and joins the one whole-program fixpoint); a dependency
namespace compiles at default cost — :inline? off for its load, so the per-form
optimize passes don't run over library code — staying direct-linked but
generically typed (the open-world default). With no app roots declared (a bare
program run, or jolt without jolt-deps) everything counts as app, so behavior is
unchanged.

malli-app JOLT_OPTIMIZE cold start: 2m10s -> 4.5s. Compute-heavy programs whose
hot code is their own namespaces (the typed ray tracer) are unaffected — their
code is app code and still fully optimized (9s/frame render). Applied at runtime
in main for the same baked-at-build-time reason as JOLT_PATH; added to the
ctx-image cache key. Help text corrected: optimization is opt-in, not default.

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

src/jolt/api.janet

@@ -177,6 +177,16 @@
       (each p (get opts :paths []) (array/push roots p))
       (when-let [jp (os/getenv "JOLT_PATH")]
         (each p (string/split ":" jp) (when (> (length p) 0) (array/push roots p)))))
+    # App source roots (jolt-87e): the project's own roots (opts :app-paths, then
+    # JOLT_APP_PATHS from jolt-deps). The whole-program inference fixpoint is
+    # scoped to namespaces loaded from these — deps load-infer per-ns instead.
+    # Empty => no app/dep distinction, so whole-program covers every namespace
+    # (a bare program run with no deps.edn keeps its old behavior).
+    (let [aps @[]]
+      (each p (get opts :app-paths []) (array/push aps p))
+      (when-let [jap (os/getenv "JOLT_APP_PATHS")]
+        (each p (string/split ":" jap) (when (> (length p) 0) (array/push aps p))))
+      (put (ctx :env) :app-source-paths aps))
     # Collection representation (persistent vs mutable) is selected at BUILD time
     # via JOLT_MUTABLE (see config.janet); init-core! registers vec/vector/conj/
     # etc. that produce the mode-appropriate values, so nothing extra to load.

src/jolt/config.janet

@@ -24,7 +24,7 @@
 # Every env var that shapes the built/optimized ctx. Both image caches key on
 # this exact list, so adding a knob here updates every cache key at once.
 (def ctx-shaping-env-vars
-  ["JOLT_PATH" "JOLT_MUTABLE" "JOLT_AOT_CORE" "JOLT_FEATURES"
+  ["JOLT_PATH" "JOLT_APP_PATHS" "JOLT_MUTABLE" "JOLT_AOT_CORE" "JOLT_FEATURES"
    "JOLT_INTERPRET" "JOLT_INTERPRET_MACROS" "JOLT_DIRECT_LINK"
    "JOLT_NO_DIRECT_LINK" "JOLT_OPTIMIZE" "JOLT_WHOLE_PROGRAM"
    "JOLT_NO_WHOLE_PROGRAM" "JOLT_SHAPE" "JOLT_NO_SHAPE"

src/jolt/deps.janet

@@ -256,6 +256,21 @@
         (spit cache-file (string/format "%j" {:key key :roots roots}))
         roots))))
 
+(defn project-source-roots
+  "The project's OWN source roots — its deps.edn :paths plus any alias
+  :extra-paths, joined to cwd — as opposed to dependency roots. These are the
+  'app' namespaces: the runtime scopes the whole-program inference fixpoint to
+  them (JOLT_APP_PATHS) so a dep-heavy app's startup doesn't re-infer every
+  transitive dependency namespace (jolt-87e)."
+  [deps-edn-path &opt aliases]
+  (def out @[])
+  (when (os/stat deps-edn-path)
+    (def edn (load-config deps-edn-path))
+    (def extra (combine-aliases edn aliases))
+    (each r (src-roots (os/cwd) edn) (array/push out r))
+    (each pp (extra :extra-paths) (array/push out (string (os/cwd) "/" pp))))
+  out)
+
 # --- :tasks (the honest subset of babashka's) ----------------------------------
 # A STRING task is a shell command. A MAP task carries :main-opts (jolt args —
 # `-e "(...)"` covers expression tasks) and an optional :doc. Babashka-style

src/jolt/deps_cli.janet

@@ -55,6 +55,12 @@
   (def rs (string/join (roots aliases) ":"))
   (def existing (os/getenv "JOLT_PATH"))
   (os/setenv "JOLT_PATH" (if (and existing (> (length existing) 0)) (string rs ":" existing) rs))
+  # The project's OWN source roots (jolt-87e): the runtime scopes whole-program
+  # inference to namespaces under these, inferring deps per-ns at load instead of
+  # re-inferring every transitive dep in one fixpoint. Absent => the runtime
+  # treats all namespaces as app (whole-program over everything, as before).
+  (when (os/stat "deps.edn")
+    (os/setenv "JOLT_APP_PATHS" (string/join (deps/project-source-roots "deps.edn" aliases) ":")))
   (os/execute [(jolt-bin) ;extra-args] :p))
 
 (defn- usage []

src/jolt/eval_base.janet

@@ -421,6 +421,21 @@
         (when (not= (last chain) file) (array/push chain file))
         (propagate err fib))))))
 
+# jolt-87e: is a namespace loaded from `path` part of the APP (vs a dependency)?
+# True when its file sits under one of the declared app source roots
+# (:app-source-paths, from JOLT_APP_PATHS / jolt-deps). When NO app roots are
+# declared (a bare program run, or jolt invoked without jolt-deps), everything
+# counts as app so whole-program covers the whole program exactly as before.
+# Only app namespaces defer into the one whole-program fixpoint; dependency
+# namespaces infer per-ns at load, so a dep-heavy app's startup doesn't re-infer
+# hundreds of transitive dependency namespaces in a single closed-world pass.
+(defn- app-source-ns?
+  [ctx path]
+  (def roots (get (ctx :env) :app-source-paths))
+  (if (or (nil? roots) (empty? roots))
+    true
+    (and path (truthy? (some |(string/has-prefix? $ path) roots)))))
+
 (defn maybe-require-ns
   "If namespace ns-name isn't populated yet, load its source — from a file on the
   context's source roots, else from the stdlib baked into the image. Restores the
@@ -445,29 +460,52 @@
           (error (string "Could not locate " ns-name
                          " on the context's source paths (JOLT_PATH / :paths)")))
         (when (or path embedded)
-          (let [saved (ctx-current-ns ctx)]
+          (let [saved (ctx-current-ns ctx)
+                # jolt-87e: is this an app namespace, or a dependency/library? Only
+                # the app is the closed world the whole-program optimizer reasons
+                # over; dependencies are open-world libraries.
+                app? (app-source-ns? ctx path)
+                # Whole-program optimize is active for this load.
+                wp-active? (and (get (ctx :env) :inline?)
+                                (get (ctx :env) :whole-program?)
+                                (not (get (ctx :env) :infer-program-done?)))
+                # A dependency under whole-program optimize compiles at DEFAULT
+                # cost: :inline? off for its load, so the per-form inline +
+                # inference passes — the bulk of optimize-mode startup — don't run
+                # over hundreds of library forms. (Direct-linking + shape-recs stay
+                # on, exactly like a non-optimized direct-link build.) This is what
+                # makes JOLT_OPTIMIZE viable on dep-heavy apps; the app's own nses
+                # below keep full optimization. (jolt-87e)
+                dep-cheap? (and wp-active? (not app?))
+                saved-inline (get (ctx :env) :inline?)]
             # Stdlib files have no `ns` form, so switch into the target ns first
             # (their defs intern there); a library's own `ns` form overrides this.
             (ctx-set-current-ns ctx ns-name)
+            (when dep-cheap? (put (ctx :env) :inline? false))
             (if path
               (load-ns-source ctx (slurp path) path)
               (load-ns-source ctx embedded (string ns-name " (stdlib)")))
+            (when dep-cheap? (put (ctx :env) :inline? saved-inline))
             # Inter-procedural collection-type inference (jolt-767): once the whole
             # unit is loaded, run the closed-world fixpoint + recompile so param-
             # dependent lookups specialize. Only in optimization mode; best-effort
             # (a failure here must not break loading). Hook installed by the api to
             # avoid an evaluator->backend circular import.
             (when (get (ctx :env) :inline?)
-              (if (and (get (ctx :env) :whole-program?)
-                       (not (get (ctx :env) :infer-program-done?)))
-                # whole-program (jolt-t34): defer — record the ns and run ONE
-                # fixpoint over all units later (the closed-world pass sees every
-                # caller, so cross-ns param types propagate). Once that batch pass
-                # has run (infer-program-done?), a ns loaded later — a lazy require
-                # inside -main — can't join it, so fall back to per-ns inference.
+              (cond
+                # whole-program (jolt-t34), APP namespace: defer — record the ns and
+                # run ONE fixpoint over all app units later (the closed-world pass
+                # sees every caller, so cross-ns param types propagate).
+                (and wp-active? app?)
                 (let [lst (or (get (ctx :env) :inferred-nses)
                               (let [a @[]] (put (ctx :env) :inferred-nses a) a))]
                   (array/push lst ns-name))
+                # whole-program, DEPENDENCY namespace (jolt-87e): nothing to do —
+                # it compiled cheaply above (no inference to run or defer).
+                wp-active?
+                nil
+                # per-ns mode (whole-program off), or a lazy require AFTER the batch
+                # ran: infer this unit on its own.
                 (when-let [iu (get (ctx :env) :infer-unit!)]
                   (protect (iu ctx ns-name)))))
             # Record load order for tooling (uberscript): a dependency finishes

src/jolt/main.janet

@@ -534,12 +534,14 @@
   (print "  uberscript OUT -m NS  Bundle NS + its required namespaces into one .clj")
   (print "  --version, version    Print the Jolt version")
   (print "  -h, --help, help      Show this help\n")
-  (print "Running a program (a file, -m/-M) direct-links and optimizes by default;")
-  (print "the repl, -e, and nrepl-server stay open so you can redefine vars.")
+  (print "Running a program (a file, -m/-M) direct-links by default; type inference")
+  (print "and specialization are opt-in via JOLT_OPTIMIZE (the cost is paid once, then")
+  (print "cached). The repl, -e, and nrepl-server stay open so you can redefine vars.")
   (print "Environment:")
+  (print "  JOLT_OPTIMIZE=1         type-infer + specialize (whole-program over app nses)")
   (print "  JOLT_NO_DIRECT_LINK=1   keep a program run open/redefinable (no optimization)")
   (print "  JOLT_NO_WHOLE_PROGRAM=1 direct-link but skip the cross-namespace pass")
-  (print "  JOLT_DIRECT_LINK=1      force direct-linking on (e.g. for -e)")
+  (print "  JOLT_DIRECT_LINK=1      force direct-linking + optimization on (e.g. for -e)")
   (print "  JOLT_WHOLE_PROGRAM=1    force the whole-program pass on")
   (print "  JOLT_INTERPRET=1        run the tree-walking interpreter\n")
   (print "Dependencies (deps.edn) are handled by the separate jolt-deps tool."))
@@ -561,6 +563,15 @@
   (when-let [jp (os/getenv "JOLT_PATH")]
     (each p (string/split ":" jp)
       (when (> (length p) 0) (array/push (get (ctx :env) :source-paths) p))))
+  # JOLT_APP_PATHS (jolt-87e), likewise applied at runtime: the project's own
+  # source roots, set by jolt-deps. Whole-program inference is scoped to
+  # namespaces under these (deps compile at default cost), so a dep-heavy app's
+  # optimize-mode startup doesn't re-infer every transitive dependency. Read here
+  # for the same baked-at-build-time reason as JOLT_PATH above.
+  (let [aps @[]]
+    (when-let [jap (os/getenv "JOLT_APP_PATHS")]
+      (each p (string/split ":" jap) (when (> (length p) 0) (array/push aps p))))
+    (put (ctx :env) :app-source-paths aps))
   # JOLT_FEATURES, likewise, must be applied at runtime: reader-features-set!
   # runs at module load, which for a baked binary is BUILD time — so a process
   # that sets JOLT_FEATURES (e.g. to read a clj-targeted lib's :clj branches)

test/integration/app-scope-wholeprogram-test.janet

@@ -0,0 +1,67 @@
+# Whole-program inference scoped to app namespaces (jolt-87e).
+#
+# Auto-whole-program (a -m program run under direct-link) used to defer EVERY
+# loaded namespace — including every transitive dependency — into one closed-
+# world fixpoint, which is prohibitive on dep-heavy apps (hundreds of dep nses;
+# a ~2-minute cold start on malli). With app source roots declared (JOLT_APP_PATHS
+# / jolt-deps, here :app-paths), only the app's OWN namespaces join the whole-
+# program batch; dependency namespaces skip inference (they stay direct-linked
+# but generically typed — the open-world default). With NO app roots declared,
+# every namespace is treated as app (whole-program over everything, pre-87e).
+
+(use ../../src/jolt/api)
+
+(var failures 0)
+(defn- check [label got want]
+  (unless (= got want)
+    (++ failures)
+    (printf "FAIL [%s] got %q want %q" label got want)))
+
+# Lay down an app root and a dep root under a tmp dir.
+(def tmp (string (os/getenv "TMPDIR") "jolt-87e-" (os/time)))
+(def app-root (string tmp "/app"))
+(def dep-root (string tmp "/dep"))
+(os/mkdir tmp) (os/mkdir app-root) (os/mkdir dep-root)
+(spit (string dep-root "/mydep.clj")
+      "(ns mydep)\n(defn helper [x] (+ x 1))\n")
+(spit (string app-root "/myapp.clj")
+      "(ns myapp (:require [mydep]))\n(defn run [x] (mydep/helper x))\n")
+
+(defn- with-wp [f]
+  (def saved (os/getenv "JOLT_WHOLE_PROGRAM"))
+  (os/setenv "JOLT_WHOLE_PROGRAM" "1")
+  (defer (os/setenv "JOLT_WHOLE_PROGRAM" saved) (f)))
+
+# --- app roots declared: only the app ns defers into the batch ---------------
+(with-wp
+  (fn []
+    (let [ctx (init {:compile? true :direct-linking? true
+                     :paths [app-root dep-root] :app-paths [app-root]})]
+      (check "whole-program is on" (truthy? (get (ctx :env) :whole-program?)) true)
+      (eval-string ctx "(require '[myapp])")
+      (let [deferred (or (get (ctx :env) :inferred-nses) @[])]
+        (check "app ns deferred to batch" (truthy? (index-of "myapp" deferred)) true)
+        (check "dep ns NOT in batch (per-ns inferred)"
+               (truthy? (index-of "mydep" deferred)) false))
+      # still runs correctly after the (scoped) whole-program pass
+      (when-let [ip (get (ctx :env) :infer-program!)] (protect (ip ctx)))
+      (check "scoped whole-program program still correct"
+             (eval-string ctx "(myapp/run 41)") 42))))
+
+# --- no app roots declared: every ns defers (pre-87e whole-program) ----------
+(with-wp
+  (fn []
+    (let [ctx (init {:compile? true :direct-linking? true
+                     :paths [app-root dep-root]})]
+      (eval-string ctx "(require '[myapp])")
+      (let [deferred (or (get (ctx :env) :inferred-nses) @[])]
+        (check "no app roots: app ns deferred" (truthy? (index-of "myapp" deferred)) true)
+        (check "no app roots: dep ns ALSO deferred"
+               (truthy? (index-of "mydep" deferred)) true))
+      (when-let [ip (get (ctx :env) :infer-program!)] (protect (ip ctx)))
+      (check "unscoped whole-program still correct"
+             (eval-string ctx "(myapp/run 9)") 10))))
+
+(if (pos? failures)
+  (do (printf "app-scope-wholeprogram: %d failure(s)" failures) (os/exit 1))
+  (print "app-scope-wholeprogram: all cases passed"))