From 554019b9edd13ce220b84f1a895fd517593c30cf Mon Sep 17 00:00:00 2001
From: Paul Guyot <pguyot@kallisys.net>
Date: Sun, 5 Oct 2025 10:17:09 +0200
Subject: [PATCH] JIT: introduce filtering of modules with --arch and
 --platform arguments

Signed-off-by: Paul Guyot <pguyot@kallisys.net>
---
 src/packbeam.erl     |  7 +++-
 src/packbeam_api.erl | 94 ++++++++++++++++++++++++++++++++++++++------
 2 files changed, 89 insertions(+), 12 deletions(-)

diff --git a/src/packbeam.erl b/src/packbeam.erl
index be56573..d9f4d09 100644
--- a/src/packbeam.erl
+++ b/src/packbeam.erl
@@ -187,7 +187,8 @@ do_create(Opts, Args) ->
         #{
             prune => maps:get(prune, Opts, false),
             start_module => maps:get(start_module, Opts, undefined),
-            include_lines => not maps:get(remove_lines, Opts, false)
+            include_lines => not maps:get(remove_lines, Opts, false),
+            arch => maps:get(arch, Opts, undefined)
         }
     ),
     0.
@@ -346,5 +347,9 @@ parse_args(["-f", Format | T], {Opts, Args}) ->
     parse_args(["--format", Format | T], {Opts, Args});
 parse_args(["--format", Format | T], {Opts, Args}) ->
     parse_args(T, {Opts#{format => Format}, Args});
+parse_args(["--arch", Arch | T], {Opts, Args}) ->
+    parse_args(T, {Opts#{arch => list_to_atom(Arch)}, Args});
+parse_args(["--platform", Platform | T], {Opts, Args}) ->
+    parse_args(T, {Opts#{platform => list_to_atom(Platform)}, Args});
 parse_args([H | T], {Opts, Args}) ->
     parse_args(T, {Opts, [H | Args]}).
diff --git a/src/packbeam_api.erl b/src/packbeam_api.erl
index 6b97ac9..4ab1c22 100644
--- a/src/packbeam_api.erl
+++ b/src/packbeam_api.erl
@@ -54,7 +54,8 @@
     prune => boolean(),
     start_module => module() | undefined,
     application_module => module() | undefined,
-    include_lines => boolean()
+    include_lines => boolean(),
+    arch => atom()
 }.
 
 -export_type([
@@ -68,7 +69,9 @@
     prune => false,
     start_module => undefined,
     application_module => undefined,
-    include_lines => true
+    include_lines => true,
+    arch => undefined,
+    platform => generic_unix
 }).
 
 %%
@@ -88,7 +91,9 @@
 %%              prune => false,
 %%              start_module => undefined,
 %%              application_module => undefined,
-%%              include_lines => false
+%%              include_lines => false,
+%%              arch => undefined,
+%%              platform => generic_unix
 %%          }'
 %%
 %% @end
@@ -122,15 +127,20 @@ create(OutputPath, InputPaths, Options) ->
         prune := Prune,
         start_module := StartModule,
         application_module := ApplicationModule,
-        include_lines := IncludeLines
+        include_lines := IncludeLines,
+        arch := Arch,
+        platform := Platform
     } = maps:merge(?DEFAULT_OPTIONS, Options),
-    ParsedFiles = parse_files(InputPaths, StartModule, IncludeLines),
+    Files0 = parse_files(InputPaths, StartModule, IncludeLines),
+    Files1 = filter_jit(Files0, Arch, Platform),
+    Files2 =
+        if
+            Prune -> prune(Files1, ApplicationModule);
+            true -> Files1
+        end,
     write_packbeam(
         OutputPath,
-        case Prune of
-            true -> prune(ParsedFiles, ApplicationModule);
-            _ -> ParsedFiles
-        end
+        Files2
     ).
 
 %%-----------------------------------------------------------------------------
@@ -782,9 +792,14 @@ ends_with(String, Suffix) ->
 
 %% @private
 filter_chunks(Chunks, IncludeLines) ->
-    AllowedChunks = allowed_chunks(IncludeLines),
+    AllowedChunks0 = allowed_chunks(IncludeLines),
+    AllowedChunks1 =
+        case lists:keymember("avmN", 1, Chunks) of
+            true -> AllowedChunks0 -- ["Code", "Type"];
+            false -> AllowedChunks0 -- ["avmN"]
+        end,
     lists:filter(
-        fun({Tag, _Data}) -> lists:member(Tag, AllowedChunks) end,
+        fun({Tag, _Data}) -> lists:member(Tag, AllowedChunks1) end,
         Chunks
     ).
 
@@ -853,3 +868,60 @@ deduplicate([]) ->
     [];
 deduplicate([H | T]) ->
     [H | [X || X <- deduplicate(T), X /= H]].
+
+filter_jit(Files0, undefined, generic_unix) ->
+    Files0;
+filter_jit(Files0, aarch64, generic_unix) ->
+    lists:filter(
+        fun(File) ->
+            filter_jit0(File, [
+                "jit.beam", "jit_stream_mmap.beam", "jit_aarch64.beam", "jit_aarch64_asm.beam"
+            ])
+        end,
+        Files0
+    );
+filter_jit(Files0, x86_64, generic_unix) ->
+    lists:filter(
+        fun(File) ->
+            filter_jit0(File, [
+                "jit.beam", "jit_stream_mmap.beam", "jit_x86_64.beam", "jit_x86_64_asm.beam"
+            ])
+        end,
+        Files0
+    );
+filter_jit(Files0, armv6m, generic_unix) ->
+    lists:filter(
+        fun(File) ->
+            filter_jit0(File, [
+                "jit.beam", "jit_stream_mmap.beam", "jit_armv6m.beam", "jit_armv6m_asm.beam"
+            ])
+        end,
+        Files0
+    );
+filter_jit(Files0, armv6m, rp2) ->
+    lists:filter(
+        fun(File) ->
+            filter_jit0(File, [
+                "jit.beam", "jit_stream_flash.beam", "jit_armv6m.beam", "jit_armv6m_asm.beam"
+            ])
+        end,
+        Files0
+    );
+filter_jit(Files0, armv6m, stm32) ->
+    lists:filter(
+        fun(File) ->
+            filter_jit0(File, [
+                "jit.beam", "jit_stream_flash.beam", "jit_armv6m.beam", "jit_armv6m_asm.beam"
+            ])
+        end,
+        Files0
+    ).
+
+filter_jit0(File, JITModules) ->
+    ModuleName = lists:flatten(get_element_name(File)),
+    case ModuleName of
+        "jit" ++ _ ->
+            lists:member(ModuleName, JITModules);
+        _ ->
+            true
+    end.