Merge pull request #2114 from se-apc/strip_fix

Allow list of chunks to be given to strip*()

OTP-15680

3 files changed, 147 insertions, 12 deletions

diff --git a/lib/stdlib/doc/src/beam_lib.xml b/lib/stdlib/doc/src/beam_lib.xml
index 8bb4cf9101..bb44ca3201 100644
--- a/lib/stdlib/doc/src/beam_lib.xml
+++ b/lib/stdlib/doc/src/beam_lib.xml
@@ -470,6 +470,18 @@ CryptoKeyFun(clear) -> term()</code>
     </func>
 
     <func>
+      <name name="strip" arity="2" since=""/>
+      <fsummary>Remove chunks not needed by the loader from a BEAM file.
+      </fsummary>
+      <desc>
+        <p>Removes all chunks from a BEAM
+          file except those needed by the loader or passed in. In particular,
+          the debug information (chunk <c>debug_info</c> and <c>abstract_code</c>)
+          is removed.</p>
+      </desc>
+    </func>
+
+    <func>
       <name name="strip_files" arity="1" since=""/>
       <fsummary>Removes chunks not needed by the loader from BEAM files.
       </fsummary>
@@ -483,6 +495,19 @@ CryptoKeyFun(clear) -> term()</code>
     </func>
 
     <func>
+      <name name="strip_files" arity="2" since=""/>
+      <fsummary>Removes chunks not needed by the loader from BEAM files.
+      </fsummary>
+      <desc>
+        <p>Removes all chunks except
+          those needed by the loader or passed in from BEAM files. In particular,
+          the debug information (chunk <c>debug_info</c> and <c>abstract_code</c>)
+          is removed. The returned list contains one element for each
+          specified filename, in the same order as in <c>Files</c>.</p>
+      </desc>
+    </func>
+
+    <func>
       <name name="strip_release" arity="1" since=""/>
       <fsummary>Remove chunks not needed by the loader from all BEAM files of
         a release.</fsummary>
@@ -497,6 +522,20 @@ CryptoKeyFun(clear) -> term()</code>
     </func>
 
     <func>
+      <name name="strip_release" arity="2" since=""/>
+      <fsummary>Remove chunks not needed by the loader from all BEAM files of
+        a release.</fsummary>
+      <desc>
+        <p>Removes all chunks
+          except those needed by the loader or passed in from the BEAM files of a
+          release. <c><anno>Dir</anno></c> is to be the installation root
+          directory. For example, the current OTP release can be
+          stripped with the call
+          <c>beam_lib:strip_release(code:root_dir())</c>.</p>
+      </desc>
+    </func>
+
+    <func>
       <name name="version" arity="1" since=""/>
       <fsummary>Read the module version of the BEAM file.</fsummary>
       <desc>
diff --git a/lib/stdlib/src/beam_lib.erl b/lib/stdlib/src/beam_lib.erl
index 3386cfcbe6..aa992f17ab 100644
--- a/lib/stdlib/src/beam_lib.erl
+++ b/lib/stdlib/src/beam_lib.erl
@@ -32,8 +32,12 @@
 	 all_chunks/1,
 	 diff_dirs/2,
 	 strip/1,
+	 strip/2,
 	 strip_files/1,
+	 strip_files/2,
 	 strip_release/1,
+	 strip_release/2,
+	 significant_chunks/0,
 	 build_module/1,
 	 version/1,
 	 md5/1,
@@ -188,7 +192,16 @@ diff_dirs(Dir1, Dir2) ->
       Beam2 :: beam().
 
 strip(FileName) ->
-    try strip_file(FileName)
+    strip(FileName, []).
+
+-spec strip(Beam1, AdditionalChunks) ->
+        {'ok', {module(), Beam2}} | {'error', 'beam_lib', info_rsn()} when
+      Beam1 :: beam(),
+      AdditionalChunks :: [chunkid()],
+      Beam2 :: beam().
+
+strip(FileName, AdditionalChunks) ->
+    try strip_file(FileName, AdditionalChunks)
     catch Error -> Error end.
     
 -spec strip_files(Files) ->
@@ -196,8 +209,17 @@ strip(FileName) ->
       Files :: [beam()],
       Beam :: beam().
 
-strip_files(Files) when is_list(Files) ->
-    try strip_fils(Files)
+strip_files(Files) ->
+    strip_files(Files, []).
+
+-spec strip_files(Files, AdditionalChunks) ->
+        {'ok', [{module(), Beam}]} | {'error', 'beam_lib', info_rsn()} when
+      Files :: [beam()],
+      AdditionalChunks :: [chunkid()],
+      Beam :: beam().
+
+strip_files(Files, AdditionalChunks) when is_list(Files) ->
+    try strip_fils(Files, AdditionalChunks)
     catch Error -> Error end.
 
 -spec strip_release(Dir) ->
@@ -207,7 +229,17 @@ strip_files(Files) when is_list(Files) ->
       Reason :: {'not_a_directory', term()} | info_rsn().
 
 strip_release(Root) ->
-    catch strip_rel(Root).
+    strip_release(Root, []).
+
+-spec strip_release(Dir, AdditionalChunks) ->
+        {'ok', [{module(), file:filename()}]}
+      | {'error', 'beam_lib', Reason} when
+      Dir :: atom() | file:filename(),
+      AdditionalChunks :: [chunkid()],
+      Reason :: {'not_a_directory', term()} | info_rsn().
+
+strip_release(Root, AdditionalChunks) ->
+    catch strip_rel(Root, AdditionalChunks).
 
 -spec version(Beam) ->
                      {'ok', {module(), [Version :: term()]}} |
@@ -401,17 +433,17 @@ cmp_lists([{Id, C1} | R1], [{Id, C2} | R2]) ->
 cmp_lists(_, _) ->
     error(different_chunks).
     
-strip_rel(Root) ->
+strip_rel(Root, AdditionalChunks) ->
     ok = assert_directory(Root),
-    strip_fils(filelib:wildcard(filename:join(Root, "lib/*/ebin/*.beam"))).
+    strip_fils(filelib:wildcard(filename:join(Root, "lib/*/ebin/*.beam")), AdditionalChunks).
 
 %% -> {ok, [{Mod, BinaryOrFileName}]} | throw(Error)
-strip_fils(Files) ->
-    {ok, [begin {ok, Reply} = strip_file(F), Reply end || F <- Files]}.
+strip_fils(Files, AdditionalChunks) ->
+    {ok, [begin {ok, Reply} = strip_file(F, AdditionalChunks), Reply end || F <- Files]}.
 
 %% -> {ok, {Mod, FileName}} | {ok, {Mod, binary()}} | throw(Error)
-strip_file(File) ->
-    {ok, {Mod, Chunks}} = read_significant_chunks(File, significant_chunks()),
+strip_file(File, AdditionalChunks) ->
+    {ok, {Mod, Chunks}} = read_significant_chunks(File, AdditionalChunks ++ significant_chunks()),
     {ok, Stripped0} = build_module(Chunks),
     Stripped = compress(Stripped0),
     case File of
diff --git a/lib/stdlib/test/beam_lib_SUITE.erl b/lib/stdlib/test/beam_lib_SUITE.erl
index 6418dc7eb6..4b2694320e 100644
--- a/lib/stdlib/test/beam_lib_SUITE.erl
+++ b/lib/stdlib/test/beam_lib_SUITE.erl
@@ -35,7 +35,7 @@
 
 -export([all/0, suite/0,groups/0,init_per_suite/1, end_per_suite/1, 
 	 init_per_group/2,end_per_group/2, 
-	 normal/1, error/1, cmp/1, cmp_literals/1, strip/1, otp_6711/1,
+	 normal/1, error/1, cmp/1, cmp_literals/1, strip/1, strip_add_chunks/1, otp_6711/1,
          building/1, md5/1, encrypted_abstr/1, encrypted_abstr_file/1]).
 
 -export([init_per_testcase/2, end_per_testcase/2]).
@@ -45,7 +45,7 @@ suite() ->
      {timetrap,{minutes,2}}].
 
 all() -> 
-    [error, normal, cmp, cmp_literals, strip, otp_6711,
+    [error, normal, cmp, cmp_literals, strip, strip_add_chunks, otp_6711,
      building, md5, encrypted_abstr, encrypted_abstr_file].
 
 groups() -> 
@@ -401,6 +401,69 @@ strip(Conf) when is_list(Conf) ->
 		  Source5D1, BeamFile5D1]),
     ok.
 
+strip_add_chunks(Conf) when is_list(Conf) ->
+    PrivDir = ?privdir,
+    {SourceD1, BeamFileD1} = make_beam(PrivDir, simple, member),
+    {Source2D1, BeamFile2D1} = make_beam(PrivDir, simple2, concat),
+    {Source3D1, BeamFile3D1} = make_beam(PrivDir, make_fun, make_fun),
+    {Source4D1, BeamFile4D1} = make_beam(PrivDir, constant, constant),
+    {Source5D1, BeamFile5D1} = make_beam(PrivDir, lines, lines),
+
+    NoOfTables = erlang:system_info(ets_count),
+    P0 = pps(),
+
+    %% strip binary
+    verify(not_a_beam_file, beam_lib:strip(<<>>)),
+    {ok, B1} = file:read_file(BeamFileD1),
+    {ok, {simple, NB1}} = beam_lib:strip(B1),
+
+    BId1 = chunk_ids(B1),
+    NBId1 = chunk_ids(NB1),
+    true = length(BId1) > length(NBId1),
+    compare_chunks(B1, NB1, NBId1),
+
+    %% Keep all the extra chunks
+    ExtraChunks = ["Abst" , "Dbgi" , "Attr" , "CInf" , "LocT" , "Atom" ],
+    {ok, {simple, AB1}} = beam_lib:strip(B1, ExtraChunks),
+    ABId1 = chunk_ids(AB1),
+    true = length(BId1) == length(ABId1),
+    compare_chunks(B1, AB1, ABId1),
+
+    %% strip file - Keep extra chunks
+    verify(file_error, beam_lib:strip(foo)),
+    {ok, {simple, _}} = beam_lib:strip(BeamFileD1, ExtraChunks),
+    compare_chunks(B1, BeamFileD1, ABId1),
+
+    %% strip_files
+    {ok, B2} = file:read_file(BeamFile2D1),
+    {ok, [{simple,_},{simple2,_}]} = beam_lib:strip_files([B1, B2], ExtraChunks),
+    {ok, [{simple,_},{simple2,_},{make_fun,_},{constant,_}]} =
+	beam_lib:strip_files([BeamFileD1, BeamFile2D1, BeamFile3D1, BeamFile4D1], ExtraChunks),
+
+    %% check that each module can be loaded.
+    {module, simple} = code:load_abs(filename:rootname(BeamFileD1)),
+    {module, simple2} = code:load_abs(filename:rootname(BeamFile2D1)),
+    {module, make_fun} = code:load_abs(filename:rootname(BeamFile3D1)),
+    {module, constant} = code:load_abs(filename:rootname(BeamFile4D1)),
+
+    %% check that line number information is still present after stripping
+    {module, lines} = code:load_abs(filename:rootname(BeamFile5D1)),
+    {'EXIT',{badarith,[{lines,t,1,Info}|_]}} = (catch lines:t(atom)),
+    false = code:purge(lines),
+    true = code:delete(lines),
+    {ok, {lines,BeamFile5D1}} = beam_lib:strip(BeamFile5D1),
+    {module, lines} = code:load_abs(filename:rootname(BeamFile5D1)),
+    {'EXIT',{badarith,[{lines,t,1,Info}|_]}} = (catch lines:t(atom)),
+
+    true = (P0 == pps()),
+    NoOfTables = erlang:system_info(ets_count),
+
+    delete_files([SourceD1, BeamFileD1,
+		  Source2D1, BeamFile2D1,
+		  Source3D1, BeamFile3D1,
+		  Source4D1, BeamFile4D1,
+		  Source5D1, BeamFile5D1]),
+    ok.
 
 otp_6711(Conf) when is_list(Conf) ->
     {'EXIT',{function_clause,_}} = (catch {a, beam_lib:info(3)}),
@@ -729,6 +792,7 @@ make_beam(Dir, Module, F) ->
     FileBase = filename:join(Dir, atom_to_list(Module)),
     Source = FileBase ++ ".erl",
     BeamFile = FileBase ++ ".beam",
+    file:delete(BeamFile),
     simple_file(Source, Module, F),
     {ok, _} = compile:file(Source, [{outdir,Dir}, debug_info, report]),
     {Source, BeamFile}.