diff --git a/src/jolt/aot.janet b/src/jolt/aot.janet new file mode 100644 index 0000000..0626a2e --- /dev/null +++ b/src/jolt/aot.janet @@ -0,0 +1,47 @@ +# Ahead-of-time images for compiled namespaces. +# +# Compile-by-default turns each form into Janet bytecode at load time. AOT skips +# that work on subsequent runs by serializing a namespace's compiled vars to a +# bytecode image and loading them back. +# +# The trick is the marshal dictionary. A compiled jolt function closes over core +# fns (core-map, +, …) and var cells; those core fns are Janet cfunctions/closures +# that can't be marshaled by value. But the runtime env that holds them is baked +# into the binary and is byte-for-byte identical at save and load time, so we +# marshal *against* it: core fns are referenced by name, and only the user's +# bytecode plus its var cells are actually serialized. + +(use ./compiler) # jolt-runtime-env +(use ./types) + +# Forward dict (key -> value) for unmarshal; reverse (value -> key) for marshal. +# Built from the runtime env, which chains to the Janet boot env, so both core fns +# and Janet builtins resolve by name. +(defn- fwd-dict [] (env-lookup jolt-runtime-env)) +(defn- rev-dict [] (invert (env-lookup jolt-runtime-env))) + +(defn marshal-ns + "Marshal namespace `ns-name`'s var mappings to a byte buffer. The whole mappings + table is marshaled in one call so var cells shared between defs stay shared." + [ctx ns-name] + (marshal ((ctx-find-ns ctx ns-name) :mappings) (rev-dict))) + +(defn unmarshal-ns! + "Install mappings produced by marshal-ns into `ns-name` in ctx, overwriting + same-named vars. Returns ns-name." + [ctx ns-name bytes] + (let [mappings (unmarshal bytes (fwd-dict)) + ns (ctx-find-ns ctx ns-name)] + (each [sym v] (pairs mappings) (put (ns :mappings) sym v)) + ns-name)) + +(defn save-ns + "Write an AOT image of compiled namespace `ns-name` to `path`." + [ctx ns-name path] + (spit path (marshal-ns ctx ns-name))) + +(defn load-ns-image + "Read an AOT image written by save-ns back into ctx under `ns-name`. Skips + parse/analyze/emit/compile entirely — the bytecode is already built." + [ctx ns-name path] + (unmarshal-ns! ctx ns-name (slurp path))) diff --git a/test/integration/aot-test.janet b/test/integration/aot-test.janet new file mode 100644 index 0000000..95c1588 --- /dev/null +++ b/test/integration/aot-test.janet @@ -0,0 +1,43 @@ +# AOT image round-trip: compile a namespace, marshal it to bytecode, load it into +# a FRESH context, and run the loaded functions without recompiling. +(use ../../src/jolt/api) +(use ../../src/jolt/aot) +(use ../../src/jolt/types) + +(print "AOT image round-trip...") + +(def img-path (string (or (os/getenv "TMPDIR") "/tmp") "/jolt-aot-test.jimg")) + +# 1. Compile a namespace into ctx1: a constant, a fn over it, a fn using core +# fns, and a recursive fn. +(def ctx1 (init {:compile? true})) +(ctx-set-current-ns ctx1 "demo") +(eval-string ctx1 "(def base 100)") +(eval-string ctx1 "(defn add-base [x] (+ x base))") +(eval-string ctx1 "(defn sum-sq [xs] (reduce + (map (fn [x] (* x x)) xs)))") +(eval-string ctx1 "(defn fact [n] (if (zero? n) 1 (* n (fact (dec n)))))") + +(assert (= 107 (eval-string ctx1 "(add-base 7)")) "ctx1 add-base") +(assert (= 14 (eval-string ctx1 "(sum-sq [1 2 3])")) "ctx1 sum-sq") +(assert (= 120 (eval-string ctx1 "(fact 5)")) "ctx1 fact") + +# 2. Save an AOT image of the compiled namespace. +(save-ns ctx1 "demo" img-path) +(assert (os/stat img-path) "image written") + +# 3. Load it into a brand-new context — no recompilation of demo. +(def ctx2 (init {:compile? true})) +(load-ns-image ctx2 "demo" img-path) +(ctx-set-current-ns ctx2 "demo") + +(assert (= 107 (eval-string ctx2 "(add-base 7)")) "ctx2 add-base from image") +(assert (= 14 (eval-string ctx2 "(sum-sq [1 2 3])")) "ctx2 sum-sq from image") +(assert (= 3628800 (eval-string ctx2 "(fact 10)")) "ctx2 fact from image (new arg)") + +# 4. The loaded vars are live: redefining one is visible to callers compiled in +# ctx2 that reference it. +(eval-string ctx2 "(def base 1000)") +(assert (= 1007 (eval-string ctx2 "(add-base 7)")) "loaded var still redefinable") + +(os/rm img-path) +(print "AOT round-trip passed!")