Fix prune of callback modules from AVM libs

src/packbeam_api.erl

@@ -747,8 +747,16 @@ parse_beam(Data, _Tmp, in_data, _Pos, Accum) ->
     case is_beam(Accum) orelse is_entrypoint(Accum) of
         true ->
             StrippedData = strip_padding(Data),
-            {ok, {Module, ChunkRefs}} = beam_lib:chunks(StrippedData, [imports, exports, atoms]),
-            [{module, Module}, {chunk_refs, ChunkRefs}, {data, StrippedData} | Accum];
+            {ok, {Module, ChunkRefs}} = beam_lib:chunks(
+                StrippedData, [imports, exports, atoms, "LitT", "LitU"], [allow_missing_chunks]
+            ),
+            [
+                {module, Module},
+                {chunk_refs, ChunkRefs},
+                {uncompressed_literals, get_uncompressed_literals(ChunkRefs)},
+                {data, StrippedData}
+                | Accum
+            ];
         _ ->
             [{data, Data} | Accum]
     end.
@@ -758,6 +766,20 @@ strip_padding(<<0:8, Rest/binary>>) ->
 strip_padding(Data) ->
     Data.
 
+%% @private
+get_uncompressed_literals(ChunkRefs) ->
+    case proplists:get_value("LitT", ChunkRefs) of
+        missing_chunk ->
+            case proplists:get_value("LitU", ChunkRefs) of
+                missing_chunk -> undefined;
+                LitU -> LitU
+            end;
+        <<0:32, Data/binary>> ->
+            Data;
+        <<_Size:4/binary, Data/binary>> ->
+            zlib:uncompress(Data)
+    end.
+
 %% @private
 maybe_uncompress_literals(Chunks) ->
     case proplists:get_value("LitT", Chunks) of

test/packbeam_api_tests.erl

@@ -553,6 +553,98 @@ packbeam_extract_test() ->
 
     ok.
 
+packbeam_create_prune_supervisor_callback_test() ->
+    %% Test that prune does not remove a module referenced only in a
+    %% supervisor child spec literal. In this scenario:
+    %%   start_mod (entrypoint) -> my_sup (via atom) -> my_worker (only in literal)
+    %% The my_worker atom only appears in the literals table of my_sup,
+    %% not in its atoms chunk or imports. Prune must still keep it.
+    AVMFile = dest_dir("packbeam_prune_sup_callback_test.avm"),
+    ?assertMatch(
+        ok,
+        packbeam_api:create(
+            AVMFile,
+            [
+                test_beam_path("start_mod.beam"),
+                test_beam_path("my_sup.beam"),
+                test_beam_path("my_worker.beam"),
+                test_beam_path("d.beam")
+            ],
+            #{prune => true}
+        )
+    ),
+
+    ParsedFiles = packbeam_api:list(AVMFile),
+
+    ?assert(is_list(ParsedFiles)),
+
+    %% start_mod calls my_sup, my_sup has my_worker only in literal child specs.
+    %% d is not referenced by anyone.
+    ?assert(parsed_file_contains_module(start_mod, ParsedFiles)),
+    ?assert(parsed_file_contains_module(my_sup, ParsedFiles)),
+    ?assert(parsed_file_contains_module(my_worker, ParsedFiles)),
+    ?assertNot(parsed_file_contains_module(d, ParsedFiles)),
+
+    %% Pruned file should be smaller than unpruned (d.beam was removed).
+    UnprunedAVMFile = dest_dir("packbeam_prune_sup_callback_unpruned.avm"),
+    ok = packbeam_api:create(
+        UnprunedAVMFile,
+        [
+            test_beam_path("start_mod.beam"),
+            test_beam_path("my_sup.beam"),
+            test_beam_path("my_worker.beam"),
+            test_beam_path("d.beam")
+        ],
+        #{}
+    ),
+    PrunedSize = filelib:file_size(AVMFile),
+    UnprunedSize = filelib:file_size(UnprunedAVMFile),
+    ?assert(PrunedSize < UnprunedSize),
+
+    ok.
+
+packbeam_create_prune_supervisor_callback_from_avm_test() ->
+    %% Same as above, but with my_worker coming from a lib AVM file.
+    %% This tests the AVM parsing path where uncompressed_literals
+    %% may not be available.
+    LibAVMFile = dest_dir("packbeam_prune_sup_callback_lib.avm"),
+    ?assertMatch(
+        ok,
+        packbeam_api:create(
+            LibAVMFile,
+            [
+                test_beam_path("my_sup.beam"),
+                test_beam_path("my_worker.beam"),
+                test_beam_path("d.beam")
+            ],
+            #{lib => true}
+        )
+    ),
+
+    AVMFile = dest_dir("packbeam_prune_sup_callback_from_avm_test.avm"),
+    ?assertMatch(
+        ok,
+        packbeam_api:create(
+            AVMFile,
+            [
+                test_beam_path("start_mod.beam"),
+                LibAVMFile
+            ],
+            #{prune => true}
+        )
+    ),
+
+    ParsedFiles = packbeam_api:list(AVMFile),
+
+    ?assert(is_list(ParsedFiles)),
+
+    ?assert(parsed_file_contains_module(start_mod, ParsedFiles)),
+    ?assert(parsed_file_contains_module(my_sup, ParsedFiles)),
+    ?assert(parsed_file_contains_module(my_worker, ParsedFiles)),
+    ?assertNot(parsed_file_contains_module(d, ParsedFiles)),
+
+    ok.
+
 file_exists(Path) ->
     filelib:is_file(Path).
 

test/prune_sup/my_sup.erl

@@ -0,0 +1,18 @@
+%%
+%% Copyright (c) 2026 Peter M <petermm@gmail.com>
+%% All rights reserved.
+%%
+%% SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
+
+%% @doc A supervisor-like module whose init returns child specs.
+%% The child module (my_worker) is only referenced as an atom.
+-module(my_sup).
+
+-export([start_link/0, init/1]).
+
+start_link() ->
+    init([]).
+
+init(_Args) ->
+    ChildSpecs = [#{id => my_worker, start => {my_worker, start_link, []}}],
+    {ok, {#{}, ChildSpecs}}.

test/prune_sup/my_worker.erl

@@ -0,0 +1,14 @@
+%%
+%% Copyright (c) 2026 Peter M <petermm@gmail.com>
+%% All rights reserved.
+%%
+%% SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
+
+%% @doc A worker module only referenced as a callback in a supervisor
+%% child spec. No module directly imports this module.
+-module(my_worker).
+
+-export([start_link/0]).
+
+start_link() ->
+    ok.

test/prune_sup/start_mod.erl

@@ -0,0 +1,15 @@
+%%
+%% Copyright (c) 2026 Peter M <petermm@gmail.com>
+%% All rights reserved.
+%%
+%% SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
+
+%% @doc A start module that starts a supervisor.
+%% The supervisor module (my_sup) is referenced via atoms and imports.
+-module(start_mod).
+
+-export([start/0]).
+
+start() ->
+    Sup = my_sup,
+    Sup:start_link().