escriptize: AOT precompilation and trailer embedding

Add --aot <target> and --jit_beams options: every beam in the packed AVM
is compiled to native code with AtomVM's jit_precompile (beams that
already carry an avmN chunk, e.g. a precompiled atomvmlib, are kept as
is), and the AVM is repacked with the start module preserved.

Add --embed trailer|objcopy and default to trailer: the AVM is appended
to the AtomVM binary followed by a 16-byte trailer (8-byte little-endian
size + ATOMVMv1 magic). This needs no external tooling and, unlike
objcopy, survives macOS code signature checks.

Also fix packbeam_api:create options to use start_module: the start key
was silently ignored and the entrypoint flag only survived by being
carried over from the input AVM.

Author: pguyot

src/atomvm_escriptize_provider.erl

@@ -33,17 +33,28 @@
     {atomvmlib, $l, "atomvmlib", string, "Path to atomvmlib.avm (default: auto discovered)"},
     {output, $o, "output", string, "Output executable name (default: app name)"},
     {objcopy, $c, "objcopy", string, "Path to objcopy tool (auto-detected if not specified)"},
-    {start, $s, "start", atom, "Start module (default: app name)"}
+    {start, $s, "start", atom, "Start module (default: app name)"},
+    {aot, $a, "aot", string,
+        "AOT-compile beams to native code for the given JIT target (e.g. aarch64)"},
+    {jit_beams, $j, "jit_beams", string,
+        "Path to the AtomVM JIT compiler beams (required with --aot)"},
+    {embed, $e, "embed", string,
+        "AVM embedding mechanism: trailer (default) or objcopy"}
 ]).
 
 -define(DEFAULT_OPTS, #{
     atomvm_binary => undefined,
     atomvmlib => undefined,
     output => undefined,
     objcopy => undefined,
-    start => undefined
+    start => undefined,
+    aot => undefined,
+    jit_beams => undefined,
+    embed => "trailer"
 }).
 
+-define(ATOMVM_TRAILER_MAGIC, <<"ATOMVMv1">>).
+
 %%
 %% provider implementation
 %%
@@ -88,7 +99,6 @@ do(State) ->
         TargetAVM = filename:join(DirName, Name ++ ".avm"),
         AtomVMLib = get_atomvmlib_path(Opts),
         AtomVMBinary = get_atomvm_binary(Opts),
-        ObjCopyTool = get_objcopy_tool(Opts),
         OutputExe = get_output_path(Opts, DirName, Name),
 
         % Get start module (default to app name)
@@ -99,13 +109,24 @@ do(State) ->
             end,
 
         % Create packed AVM with atomvmlib
-        PackedAVM = create_packed_avm(TargetAVM, AtomVMLib, DirName, Name, StartModule),
+        PackedAVM0 = create_packed_avm(TargetAVM, AtomVMLib, DirName, Name, StartModule),
+
+        % Optionally AOT-compile the packed beams to native code
+        PackedAVM = maybe_aot_precompile(PackedAVM0, Opts, DirName, Name, StartModule),
 
         % Copy AtomVM binary
         ok = copy_atomvm_binary(AtomVMBinary, OutputExe),
 
         % Embed AVM into executable
-        ok = embed_avm(ObjCopyTool, OutputExe, PackedAVM),
+        case maps:get(embed, Opts) of
+            "objcopy" ->
+                ObjCopyTool = get_objcopy_tool(Opts),
+                ok = embed_avm(ObjCopyTool, OutputExe, PackedAVM);
+            "trailer" ->
+                ok = embed_avm_trailer(OutputExe, PackedAVM);
+            OtherEmbed ->
+                throw({unknown_embed_mechanism, OtherEmbed})
+        end,
 
         % Make executable
         ok = make_executable(OutputExe),
@@ -335,7 +356,7 @@ create_packed_avm(TargetAVM, AtomVMLib, DirName, Name, StartModule) ->
     % Read both AVM files
     case {filelib:is_file(AtomVMLib), filelib:is_file(TargetAVM)} of
         {true, true} ->
-            packbeam_api:create(PackedAVM, [AtomVMLib, TargetAVM], #{start => StartModule}),
+            packbeam_api:create(PackedAVM, [AtomVMLib, TargetAVM], #{start_module => StartModule}),
             rebar_api:info("Created packed AVM: ~s with start module ~p", [PackedAVM, StartModule]),
             PackedAVM;
         {false, _} ->
@@ -344,6 +365,92 @@ create_packed_avm(TargetAVM, AtomVMLib, DirName, Name, StartModule) ->
             throw({file_not_found, TargetAVM})
     end.
 
+%% @private
+%% AOT-compile every beam in the packed AVM to native code for the given
+%% JIT target, using the jit_precompile compiler from the AtomVM sources.
+maybe_aot_precompile(PackedAVM, Opts, DirName, Name, StartModule) ->
+    case maps:get(aot, Opts) of
+        undefined ->
+            PackedAVM;
+        Target ->
+            JitBeams =
+                case maps:get(jit_beams, Opts) of
+                    undefined -> throw(jit_beams_required_with_aot);
+                    Path -> Path
+                end,
+            case filelib:is_file(filename:join(JitBeams, "jit_precompile.beam")) of
+                true -> ok;
+                false -> throw({jit_precompile_not_found, JitBeams})
+            end,
+            true = code:add_patha(JitBeams),
+            aot_precompile_avm(PackedAVM, Target, DirName, Name, StartModule)
+    end.
+
+%% @private
+aot_precompile_avm(PackedAVM, Target, DirName, Name, StartModule) ->
+    Elements = packbeam_api:list(PackedAVM),
+    BeamNames = [
+        packbeam_api:get_element_name(E)
+     || E <- Elements, packbeam_api:is_beam(E)
+    ],
+    WorkDir = filename:join(DirName, Name ++ "_aot"),
+    InDir = filename:join(WorkDir, "in"),
+    OutDir = filename:join(WorkDir, "out"),
+    ok = filelib:ensure_dir(filename:join(InDir, "dummy")),
+    ok = filelib:ensure_dir(filename:join(OutDir, "dummy")),
+    rebar_api:info("AOT-compiling ~p beams for ~s", [length(BeamNames), Target]),
+    lists:foreach(
+        fun(E) ->
+            case packbeam_api:is_beam(E) of
+                true ->
+                    ElementName = packbeam_api:get_element_name(E),
+                    Data = packbeam_api:get_element_data(E),
+                    case has_native_code(Data) of
+                        true ->
+                            % already AOT-compiled (e.g. a precompiled
+                            % atomvmlib): keep as is
+                            ok = file:write_file(filename:join(OutDir, ElementName), Data);
+                        false ->
+                            BeamPath = filename:join(InDir, ElementName),
+                            ok = file:write_file(BeamPath, Data),
+                            ok = jit_precompile:compile(Target, OutDir, false, BeamPath)
+                    end;
+                false ->
+                    ok
+            end
+        end,
+        Elements
+    ),
+    % Non-beam elements (e.g. priv files) keep their names by going through
+    % an AVM stripped of its beams
+    NonBeamsAVM = filename:join(WorkDir, "non_beams.avm"),
+    packbeam_api:delete(NonBeamsAVM, PackedAVM, BeamNames),
+    AotAVM = filename:join(DirName, Name ++ "_aot.avm"),
+    AotBeams = [filename:join(OutDir, BeamName) || BeamName <- BeamNames],
+    packbeam_api:create(AotAVM, AotBeams ++ [NonBeamsAVM], #{start_module => StartModule}),
+    rebar_api:info("Created AOT packed AVM: ~s", [AotAVM]),
+    AotAVM.
+
+%% @private
+has_native_code(BeamData) ->
+    case beam_lib:chunks(BeamData, ["avmN"], [allow_missing_chunks]) of
+        {ok, {_Module, [{"avmN", Chunk}]}} -> Chunk =/= missing_chunk;
+        _ -> false
+    end.
+
+%% @private
+%% Append the AVM and a 16-byte trailer (8-byte little-endian size +
+%% "ATOMVMv1" magic) to the executable. Unlike objcopy, this requires no
+%% external tooling and is tolerated by the macOS code signature checks.
+embed_avm_trailer(Executable, AVMFile) ->
+    {ok, AVMData} = file:read_file(AVMFile),
+    Size = byte_size(AVMData),
+    Trailer = <<Size:64/little, (?ATOMVM_TRAILER_MAGIC)/binary>>,
+    rebar_api:debug("Appending ~p byte AVM with trailer to ~s", [Size, Executable]),
+    {ok, Fd} = file:open(Executable, [append, binary, raw]),
+    ok = file:write(Fd, [AVMData, Trailer]),
+    ok = file:close(Fd).
+
 %% @private
 copy_atomvm_binary(Source, Dest) ->
     rebar_api:debug("Copying AtomVM binary from ~s to ~s", [Source, Dest]),

test/driver/src/escriptize_tests.erl

@@ -23,6 +23,7 @@
 
 run(Opts) ->
     ok = test_defaults(Opts),
+    ok = test_aot(Opts),
     ok.
 
 %% @private
@@ -44,7 +45,74 @@ test_defaults(Opts) ->
     ExecPath = test:make_path([AppDir, "_build/default/bin/myscript"]),
     ok = test:file_exists(ExecPath),
 
-    [] = test:execute_cmd(ExecPath, Opts),
+    % stderr carries VM diagnostics (e.g. JIT compilation timing); the
+    % escript contract is about stdout
+    [] = test:execute_cmd(ExecPath ++ " 2>/dev/null", Opts),
+
+    % The default (trailer) mechanism appends the AVM followed by a
+    % 16-byte trailer: 8-byte little-endian size + "ATOMVMv1" magic
+    {ok, Exec} = file:read_file(ExecPath),
+    TrailerOffset = byte_size(Exec) - 16,
+    <<_:TrailerOffset/binary, AVMSize:64/little, "ATOMVMv1">> = Exec,
+    true = AVMSize > 0 andalso AVMSize < TrailerOffset,
+
+    test:tick().
+
+%% @private
+%% AOT precompilation needs the jit compiler beams from an AtomVM source
+%% tree and an AtomVM binary built with JIT support; the test is skipped
+%% unless ATOMVM_REBAR3_PLUGIN_TEST_JIT_BEAMS, _JIT_TARGET (and optionally
+%% ATOMVM_REBAR3_PLUGIN_TEST_ATOMVM/_ATOMVMLIB) are set.
+test_aot(Opts) ->
+    case
+        {
+            os:getenv("ATOMVM_REBAR3_PLUGIN_TEST_JIT_BEAMS"),
+            os:getenv("ATOMVM_REBAR3_PLUGIN_TEST_JIT_TARGET")
+        }
+    of
+        {JitBeams, Target} when JitBeams =/= false, Target =/= false ->
+            test_aot(Opts, JitBeams, Target);
+        _ ->
+            io:format("(aot test skipped: ATOMVM_REBAR3_PLUGIN_TEST_JIT_* unset) "),
+            ok
+    end.
+
+%% @private
+test_aot(Opts, JitBeams, Target) ->
+    AppsDir = maps:get(apps_dir, Opts),
+    AppDir = test:make_path([AppsDir, "myscript"]),
+
+    CmdOpts0 = [
+        {"--aot", Target},
+        {"--jit_beams", JitBeams},
+        {"--output", "myscript_aot"}
+    ],
+    CmdOpts1 =
+        case os:getenv("ATOMVM_REBAR3_PLUGIN_TEST_ATOMVM") of
+            false -> CmdOpts0;
+            AtomVM -> [{"--atomvm_binary", AtomVM} | CmdOpts0]
+        end,
+    CmdOpts =
+        case os:getenv("ATOMVM_REBAR3_PLUGIN_TEST_ATOMVMLIB") of
+            false -> CmdOpts1;
+            AtomVMLib -> [{"--atomvmlib", AtomVMLib} | CmdOpts1]
+        end,
+
+    Cmd = create_escriptize_cmd(AppDir, CmdOpts, []),
+    Output = test:execute_cmd(Cmd, Opts),
+    test:debug(Output, Opts),
+
+    ok = test:expect_contains("AOT-compiling", Output),
+    ok = test:expect_contains("Created AOT packed AVM:", Output),
+    ok = test:expect_contains("Created standalone executable:", Output),
+
+    ExecPath = test:make_path([AppDir, "_build/default/bin/myscript_aot"]),
+    ok = test:file_exists(ExecPath),
+
+    Run = test:execute_cmd(ExecPath ++ " 2>&1 >/dev/null", Opts),
+    % AOT-compiled modules must not be compiled again at runtime
+    false = string:find(Run, "Compilation of") =/= nomatch,
+    [] = test:execute_cmd(ExecPath ++ " 2>/dev/null", Opts),
 
     test:tick().