diff --git a/src/packbeam_api.erl b/src/packbeam_api.erl index 3c0d11b..47f45a3 100644 --- a/src/packbeam_api.erl +++ b/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 diff --git a/test/packbeam_api_tests.erl b/test/packbeam_api_tests.erl index 0d9c0a3..6ce64a3 100644 --- a/test/packbeam_api_tests.erl +++ b/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). diff --git a/test/prune_sup/my_sup.erl b/test/prune_sup/my_sup.erl new file mode 100644 index 0000000..e3fc411 --- /dev/null +++ b/test/prune_sup/my_sup.erl @@ -0,0 +1,18 @@ +%% +%% Copyright (c) 2026 Peter M +%% 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}}. diff --git a/test/prune_sup/my_worker.erl b/test/prune_sup/my_worker.erl new file mode 100644 index 0000000..797296d --- /dev/null +++ b/test/prune_sup/my_worker.erl @@ -0,0 +1,14 @@ +%% +%% Copyright (c) 2026 Peter M +%% 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. diff --git a/test/prune_sup/start_mod.erl b/test/prune_sup/start_mod.erl new file mode 100644 index 0000000..7dd1a85 --- /dev/null +++ b/test/prune_sup/start_mod.erl @@ -0,0 +1,15 @@ +%% +%% Copyright (c) 2026 Peter M +%% 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().