Dead-code elimination in jolt uberscript (jolt-atg) (#130)

A bundle is closed-world — everything it needs is inlined and nothing is
required afterward — so a user defn unreachable from the entry's reference
graph can be dropped. The bundler now computes reachability from main-ns/-main
plus every non-prunable form and drops dead defn/defn- by exact source span
(formatting and reader macros in the surviving code are untouched).

Conservative and sound: only plain defn/defn- are prunable; a defn is kept if
its bare or ns-qualified name appears in any kept form, the closure runs to a
fixpoint, and any use of resolve/ns-resolve/requiring-resolve/find-var/intern/
eval/load-string disables pruning entirely. A parse failure on any file also
falls back to verbatim bundling, so the command stays as robust as a plain
concatenation. defmethod/defrecord/extend bodies are non-prunable and scanned,
so a fn reached only via dynamic dispatch stays live.

New reader/parse-all-spans returns [form start end] byte offsets so the drop
is a verbatim slice, not a re-print.

30-fn library used by a 3-fn entry: bundle 1114 -> 437 bytes (61% smaller, 27
dead fns dropped), output byte-identical.

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

src/jolt/main.janet

@@ -496,6 +496,45 @@
       (load-string ctx (string "(apply " ns-name "/-main *command-line-args*)")))
     ([err fib] (report-error err fib) (os/exit 1))))
 
+# --- uberscript dead-code elimination (jolt-atg) ---------------------------
+# A bundle is closed-world: everything it needs is inlined and nothing is
+# required afterward, so a user `defn` unreachable from the entry's reference
+# graph can be dropped. Conservative + sound: only plain defn/defn- are
+# prunable; a defn is kept if its (bare or ns-qualified) name appears in any
+# kept form, the closure runs to a fixpoint, and any use of dynamic resolution
+# disables pruning entirely. The drop is by exact source span, so formatting
+# and reader macros in the surviving code are untouched.
+(defn- dce-sym? [x] (and (struct? x) (= :symbol (get x :jolt/type))))
+(def- dce-bailout-syms
+  {"resolve" true "ns-resolve" true "requiring-resolve" true "find-var" true
+   "intern" true "eval" true "load-string" true})
+(defn- dce-collect-syms [form acc]
+  (cond
+    (dce-sym? form) (do (put acc (get form :name) true)
+                        (when (get form :ns)
+                          (put acc (string (get form :ns) "/" (get form :name)) true)))
+    (indexed? form) (each x form (dce-collect-syms x acc))
+    (or (struct? form) (table? form)) (each k (keys form)
+                                        (dce-collect-syms k acc)
+                                        (dce-collect-syms (get form k) acc)))
+  acc)
+(defn- dce-defn-name [form]
+  (when (and (indexed? form) (>= (length form) 2) (dce-sym? (get form 0))
+             (let [nm (get (in form 0) :name)] (or (= nm "defn") (= nm "defn-")))
+             (dce-sym? (get form 1)))
+    (get (in form 1) :name)))
+(defn- dce-strip-spans [src dead]
+  (if (empty? dead) src
+    (do
+      (sort dead (fn [a b] (< (a 0) (b 0))))
+      (def buf @"")
+      (var cur 0)
+      (each [s e] dead
+        (when (> s cur) (buffer/push-string buf (string/slice src cur s)))
+        (set cur (max cur e)))
+      (buffer/push-string buf (string/slice src cur))
+      (string buf))))
+
 (defn- run-uberscript [out main-ns]
   # Bundle main-ns and everything it requires (from JOLT_PATH roots) into one
   # .clj that runs on a plain jolt — no deps, no jpm. We require the entry and
@@ -510,15 +549,64 @@
   (def files @[])
   (each f (get (ctx :env) :loaded-files)
     (unless (get seen f) (put seen f true) (array/push files f)))
+  # read every file's source and parse it into top-level forms with byte spans;
+  # if ANY file fails to parse, fall back to verbatim bundling (DCE off) so the
+  # uberscript stays exactly as robust as a plain concatenation.
+  (def srcs @{})
+  (def file-forms @{})
+  (def all-forms @[])
+  (var dce-ok true)
+  (each f files
+    (def src (slurp f))
+    (put srcs f src)
+    (def lst @[])
+    (try
+      (each [form s e] (parse-all-spans src f)
+        (def entry @{:start s :end e :dname (dce-defn-name form) :form form})
+        (array/push lst entry)
+        (array/push all-forms entry))
+      ([_err _fib] (set dce-ok false)))
+    (put file-forms f lst))
+  # disable DCE if the bundle resolves names dynamically (a defn could be
+  # reached by a string/symbol the reference scan can't see).
+  (when dce-ok
+    (def allsyms @{})
+    (each e all-forms (dce-collect-syms (e :form) allsyms))
+    (each nm (keys allsyms) (when (get dce-bailout-syms nm) (set dce-ok false))))
+  # reachability: seed from the entry (-main) and every non-prunable form, then
+  # close over the bodies of defns that become live.
+  (def live @{})
+  (when dce-ok
+    (def referenced @{"-main" true})
+    (each e all-forms (unless (e :dname) (dce-collect-syms (e :form) referenced)))
+    (var changed true)
+    (while changed
+      (set changed false)
+      (each e all-forms
+        (when (and (e :dname) (not (get live e)) (get referenced (e :dname)))
+          (put live e true)
+          (dce-collect-syms (e :form) referenced)
+          (set changed true)))))
+  (var dropped 0)
   (def buf @"")
   (buffer/push-string buf (string ";; Generated by `jolt uberscript` — " (length files) " namespace(s)\n\n"))
   (each f files
     (buffer/push-string buf (string ";; --- " f " ---\n"))
-    (buffer/push-string buf (slurp f))
+    (def src (get srcs f))
+    (if-not dce-ok
+      (buffer/push-string buf src)
+      (do
+        (def dead @[])
+        (each e (get file-forms f)
+          (when (and (e :dname) (not (get live e)))
+            (++ dropped)
+            (array/push dead [(e :start) (e :end)])))
+        (buffer/push-string buf (dce-strip-spans src dead))))
     (buffer/push-string buf "\n"))
   (buffer/push-string buf (string "\n(apply " main-ns "/-main *command-line-args*)\n"))
   (spit out (string buf))
-  (print "Wrote " out " (" (length files) " namespace(s))"))
+  (print "Wrote " out " (" (length files) " namespace(s)"
+         (if (and dce-ok (> dropped 0)) (string ", " dropped " dead fn(s) dropped") "") ")"))
 
 (defn- print-help []
   (print "Jolt — a Clojure interpreter on Janet\n")

src/jolt/reader.janet

@@ -792,3 +792,38 @@
     (set s rest*)
     (when (not (nil? form)) (array/push out [form form-line])))
   out)
+
+(defn parse-all-spans
+  "Parse every top-level form of source, returning an array of
+  [form abs-start abs-end] where abs-start is the byte offset of the form's
+  first non-trivia char and abs-end is just past it (leading whitespace/comment
+  trivia is excluded). Lets a caller reconstruct source by slicing verbatim —
+  used by `jolt uberscript`'s dead-code elimination to drop a defn's exact byte
+  span without re-printing (so formatting and reader macros survive)."
+  [source &opt file]
+  (def out @[])
+  (var s source)
+  (while (> (length (string/trim s)) 0)
+    (def base (- (length source) (length s)))
+    (var i 0)
+    (def n (length s))
+    (var scanning true)
+    (while (and scanning (< i n))
+      (def c (in s i))
+      (cond
+        (or (= c (chr "\n")) (= c (chr " ")) (= c (chr "\t")) (= c (chr "\r")) (= c (chr ","))) (++ i)
+        (= c (chr ";")) (while (and (< i n) (not= (in s i) (chr "\n"))) (++ i))
+        (set scanning false)))
+    (def [form rest*]
+      (try (parse-next s)
+        ([err fib]
+          (if (reader-error? err)
+            (error (format-reader-error
+                     {:jolt/type :jolt/reader-error :msg (err :msg)
+                      :pos (+ base (err :pos))}
+                     source file))
+            (propagate err fib)))))
+    (def consumed (- (length s) (length rest*)))
+    (set s rest*)
+    (when (not (nil? form)) (array/push out [form (+ base i) (+ base consumed)])))
+  out)