Introduce beam_utils:replace_labels/4

1 files changed, 66 insertions, 2 deletions

diff --git a/lib/compiler/src/beam_utils.erl b/lib/compiler/src/beam_utils.erl
index cc6e54ca16..df41e35c82 100644
--- a/lib/compiler/src/beam_utils.erl
+++ b/lib/compiler/src/beam_utils.erl
@@ -23,14 +23,14 @@
 -module(beam_utils).
 -export([is_killed_block/2,is_killed/3,is_killed_at/3,
 	 is_not_used/3,
-	 empty_label_index/0,index_label/3,index_labels/1,
+	 empty_label_index/0,index_label/3,index_labels/1,replace_labels/4,
 	 code_at/2,bif_to_test/3,is_pure_test/1,
 	 live_opt/1,delete_live_annos/1,combine_heap_needs/2,
 	 split_even/1]).
 
 -export_type([code_index/0,module_code/0,instruction/0]).
 
--import(lists, [member/2,sort/1,reverse/1,splitwith/2]).
+-import(lists, [map/2,member/2,sort/1,reverse/1,splitwith/2]).
 
 %% instruction() describes all instructions that are used during optimzation
 %% (from beam_a to beam_z).
@@ -160,6 +160,18 @@ index_label(Lbl, Is0, Acc) ->
 code_at(L, Ll) ->
     gb_trees:get(L, Ll).
 
+%% replace_labels(FunctionIs, Tail, ReplaceDb, Fallback) -> FunctionIs.
+%%  Replace all labels in instructions according to the ReplaceDb.
+%%  If label is not found the Fallback is called with the label to
+%%  produce a new one.
+
+-spec replace_labels([instruction()],
+                     [instruction()],
+                     #{beam_asm:label() => beam_asm:label()},
+                     fun((beam_asm:label()) -> term())) -> [instruction()].
+replace_labels(Is, Acc, D, Fb) ->
+    replace_labels_1(Is, Acc, D, Fb).
+
 %% bif_to_test(Bif, [Op], Fail) -> {test,Test,Fail,[Op]}
 %%  Convert a BIF to a test. Fail if not possible.
 
@@ -643,6 +655,58 @@ index_labels_1([], Acc) -> gb_trees:from_orddict(sort(Acc)).
 drop_labels([{label,_}|Is]) -> drop_labels(Is);
 drop_labels(Is) -> Is.
 
+
+replace_labels_1([{test,Test,{f,Lbl},Ops}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{test,Test,{f,label(Lbl, D, Fb)},Ops}|Acc], D, Fb);
+replace_labels_1([{test,Test,{f,Lbl},Live,Ops,Dst}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{test,Test,{f,label(Lbl, D, Fb)},Live,Ops,Dst}|Acc], D, Fb);
+replace_labels_1([{select,I,R,{f,Fail0},Vls0}|Is], Acc, D, Fb) ->
+    Vls = map(fun ({f,L}) -> {f,label(L, D, Fb)};
+		   (Other) -> Other
+	      end, Vls0),
+    Fail = label(Fail0, D, Fb),
+    replace_labels_1(Is, [{select,I,R,{f,Fail},Vls}|Acc], D, Fb);
+replace_labels_1([{'try',R,{f,Lbl}}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{'try',R,{f,label(Lbl, D, Fb)}}|Acc], D, Fb);
+replace_labels_1([{'catch',R,{f,Lbl}}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{'catch',R,{f,label(Lbl, D, Fb)}}|Acc], D, Fb);
+replace_labels_1([{jump,{f,Lbl}}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{jump,{f,label(Lbl, D, Fb)}}|Acc], D, Fb);
+replace_labels_1([{loop_rec,{f,Lbl},R}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{loop_rec,{f,label(Lbl, D, Fb)},R}|Acc], D, Fb);
+replace_labels_1([{loop_rec_end,{f,Lbl}}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{loop_rec_end,{f,label(Lbl, D, Fb)}}|Acc], D, Fb);
+replace_labels_1([{wait,{f,Lbl}}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{wait,{f,label(Lbl, D, Fb)}}|Acc], D, Fb);
+replace_labels_1([{wait_timeout,{f,Lbl},To}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{wait_timeout,{f,label(Lbl, D, Fb)},To}|Acc], D, Fb);
+replace_labels_1([{bif,Name,{f,Lbl},As,R}|Is], Acc, D, Fb) when Lbl =/= 0 ->
+    replace_labels_1(Is, [{bif,Name,{f,label(Lbl, D, Fb)},As,R}|Acc], D, Fb);
+replace_labels_1([{gc_bif,Name,{f,Lbl},Live,As,R}|Is], Acc, D, Fb) when Lbl =/= 0 ->
+    replace_labels_1(Is, [{gc_bif,Name,{f,label(Lbl, D, Fb)},Live,As,R}|Acc], D, Fb);
+replace_labels_1([{call,Ar,{f,Lbl}}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{call,Ar,{f,label(Lbl, D, Fb)}}|Acc], D, Fb);
+replace_labels_1([{make_fun2,{f,Lbl},U1,U2,U3}|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [{make_fun2,{f,label(Lbl, D, Fb)},U1,U2,U3}|Acc], D, Fb);
+replace_labels_1([{bs_init,{f,Lbl},Info,Live,Ss,Dst}|Is], Acc, D, Fb) when Lbl =/= 0 ->
+    replace_labels_1(Is, [{bs_init,{f,label(Lbl, D, Fb)},Info,Live,Ss,Dst}|Acc], D, Fb);
+replace_labels_1([{bs_put,{f,Lbl},Info,Ss}|Is], Acc, D, Fb) when Lbl =/= 0 ->
+    replace_labels_1(Is, [{bs_put,{f,label(Lbl, D, Fb)},Info,Ss}|Acc], D, Fb);
+replace_labels_1([{put_map=I,{f,Lbl},Op,Src,Dst,Live,List}|Is], Acc, D, Fb)
+  when Lbl =/= 0 ->
+    replace_labels_1(Is, [{I,{f,label(Lbl, D, Fb)},Op,Src,Dst,Live,List}|Acc], D, Fb);
+replace_labels_1([{get_map_elements=I,{f,Lbl},Src,List}|Is], Acc, D, Fb) when Lbl =/= 0 ->
+    replace_labels_1(Is, [{I,{f,label(Lbl, D, Fb)},Src,List}|Acc], D, Fb);
+replace_labels_1([I|Is], Acc, D, Fb) ->
+    replace_labels_1(Is, [I|Acc], D, Fb);
+replace_labels_1([], Acc, _, _) -> Acc.
+
+label(Old, D, Fb) ->
+    case D of
+        #{Old := New} -> New;
+        _ -> Fb(Old)
+    end.
+
 %% Help functions for combine_heap_needs.
 
 combine_alloc_lists(Al1, Al2) ->