Move allocation combining from beam_flatten to beam_ssa_codegen

Continuing the simplification of beam_flatten, move the optimization
that eliminates a test_heap instruction following a binary construction
by incorporating the allocation of the heap space into the bs_init*
instruction itself.

This change does not change the generated code in any way.

Also remove beam_utils:combine_heap_needs/2, because beam_flatten
was the last user of it.

3 files changed, 39 insertions, 70 deletions

diff --git a/lib/compiler/src/beam_flatten.erl b/lib/compiler/src/beam_flatten.erl
index c8bc1b7b84..07c24abf7b 100644
--- a/lib/compiler/src/beam_flatten.erl
+++ b/lib/compiler/src/beam_flatten.erl
@@ -42,13 +42,6 @@ block([{block,Is0}|Is1], Acc) -> block(Is1, norm_block(Is0, Acc));
 block([I|Is], Acc) -> block(Is, [I|Acc]);
 block([], Acc) -> reverse(Acc).
 
-norm_block([{set,[],[],{alloc,R,{_,nostack,_,_}=Alloc}}|Is], Acc0) ->
-    case insert_alloc_in_bs_init(Acc0, Alloc) of
-	impossible ->
-	    norm_block(Is, reverse(norm_allocate(Alloc, R), Acc0));
-	Acc ->
-	    norm_block(Is, Acc)
-    end;
 norm_block([{set,[],[],{alloc,R,Alloc}}|Is], Acc0) ->
     norm_block(Is, reverse(norm_allocate(Alloc, R), Acc0));
 norm_block([{set,[D1],[S],get_hd},{set,[D2],[S],get_tl}|Is], Acc) ->
@@ -56,7 +49,7 @@ norm_block([{set,[D1],[S],get_hd},{set,[D2],[S],get_tl}|Is], Acc) ->
     norm_block(Is, [I|Acc]);
 norm_block([I|Is], Acc) -> norm_block(Is, [norm(I)|Acc]);
 norm_block([], Acc) -> Acc.
-    
+
 norm({set,[D],As,{bif,N,F}})      -> {bif,N,F,As,D};
 norm({set,[D],As,{alloc,R,{gc_bif,N,F}}}) -> {gc_bif,N,F,R,As,D};
 norm({set,[D],[],init})           -> {init,D};
@@ -90,27 +83,3 @@ norm_allocate({nozero,Ns,0,Inits}, Regs) ->
     [{allocate,Ns,Regs}|Inits];
 norm_allocate({nozero,Ns,Nh,Inits}, Regs) ->
     [{allocate_heap,Ns,Nh,Regs}|Inits].
-
-%% insert_alloc_in_bs_init(ReverseInstructionStream, AllocationInfo) ->
-%%                                  impossible | ReverseInstructionStream'
-%%   A bs_init/6 instruction should not be followed by a test heap instruction.
-%%   Given the AllocationInfo from a test heap instruction, merge the
-%%   allocation amounts into the previous bs_init/6 instruction (if any).
-%%
-insert_alloc_in_bs_init([{bs_put,_,_,_}=I|Is], Alloc) ->
-    %% The instruction sequence ends with an bs_put/4 instruction.
-    %% We'll need to search backwards for the bs_init/6 instruction.
-    insert_alloc_1(Is, Alloc, [I]);
-insert_alloc_in_bs_init(_, _) -> impossible.
-
-insert_alloc_1([{bs_init=Op,Fail,Info0,Live,Ss,Dst}|Is],
-	       {_,nostack,Ws2,[]}, Acc) when is_integer(Live) ->
-    %% The number of extra heap words is always in the second position
-    %% in the Info tuple.
-    Ws1 = element(2, Info0),
-    Al = beam_utils:combine_heap_needs(Ws1, Ws2),
-    Info = setelement(2, Info0, Al),
-    I = {Op,Fail,Info,Live,Ss,Dst},
-    reverse(Acc, [I|Is]);
-insert_alloc_1([{bs_put,_,_,_}=I|Is], Alloc, Acc) ->
-    insert_alloc_1(Is, Alloc, [I|Acc]).
diff --git a/lib/compiler/src/beam_ssa_codegen.erl b/lib/compiler/src/beam_ssa_codegen.erl
index cae9920f40..142df5192f 100644
--- a/lib/compiler/src/beam_ssa_codegen.erl
+++ b/lib/compiler/src/beam_ssa_codegen.erl
@@ -218,7 +218,7 @@ need_heap_never(_) -> false.
 
 need_heap_blks([{L,#cg_blk{is=Is0}=Blk0}|Bs], H0, Acc) ->
     {Is1,H1} = need_heap_is(reverse(Is0), H0, []),
-    {Ns,H} = need_heap_terminator(Bs, H1),
+    {Ns,H} = need_heap_terminator(Bs, L, H1),
     Is = Ns ++ Is1,
     Blk = Blk0#cg_blk{is=Is},
     need_heap_blks(Bs, H, [{L,Blk}|Acc]);
@@ -228,6 +228,13 @@ need_heap_blks([], H, Acc) ->
 need_heap_is([#cg_alloc{words=Words}=Alloc0|Is], N, Acc) ->
     Alloc = Alloc0#cg_alloc{words=add_heap_words(N, Words)},
     need_heap_is(Is, #need{}, [Alloc|Acc]);
+need_heap_is([#cg_set{anno=Anno,op=bs_init}=I0|Is], N, Acc) ->
+    Alloc = case need_heap_need(N) of
+                [#cg_alloc{words=Need}] -> alloc(Need);
+                [] -> 0
+            end,
+    I = I0#cg_set{anno=Anno#{alloc=>Alloc}},
+    need_heap_is(Is, #need{}, [I|Acc]);
 need_heap_is([#cg_set{op=Op,args=Args}=I|Is], N, Acc) ->
     case classify_heap_need(Op, Args) of
         {put,Words} ->
@@ -243,11 +250,31 @@ need_heap_is([#cg_set{op=Op,args=Args}=I|Is], N, Acc) ->
 need_heap_is([], N, Acc) ->
     {Acc,N}.
 
-need_heap_terminator([{_,#cg_blk{last=#cg_br{succ=Same,fail=Same}}}|_], N) ->
+need_heap_terminator([{_,#cg_blk{last=#cg_br{succ=L,fail=L}}}|_], L, N) ->
+    %% Fallthrough.
     {[],N};
-need_heap_terminator([{_,#cg_blk{}}|_], N) ->
+need_heap_terminator([{_,#cg_blk{is=Is,last=#cg_br{succ=L}}}|_], L, N) ->
+    case need_heap_need(N) of
+        [] ->
+            {[],#need{}};
+        [_|_]=Alloc ->
+            %% If the preceding instructions are a binary construction,
+            %% hoist the allocation and incorporate into the bs_init
+            %% instruction.
+            case reverse(Is) of
+                [#cg_set{op=succeeded},#cg_set{op=bs_init}|_] ->
+                    {[],N};
+                [#cg_set{op=bs_put}|_] ->
+                    {[],N};
+                _ ->
+                    %% Not binary construction. Must emit an allocation
+                    %% instruction in this block.
+                    {Alloc,#need{}}
+            end
+    end;
+need_heap_terminator([{_,#cg_blk{}}|_], _, N) ->
     {need_heap_need(N),#need{}};
-need_heap_terminator([], H) ->
+need_heap_terminator([], _, H) ->
     {need_heap_need(H),#need{}}.
 
 need_heap_need(#need{h=0,f=0}) -> [];
@@ -1022,12 +1049,13 @@ cg_block([#cg_set{op=bs_init,dst=Dst0,args=Args0,anno=Anno}=I,
           #cg_set{op=succeeded,dst=Bool}], {Bool,Fail0}, St) ->
     Fail = bif_fail(Fail0),
     Line = line(Anno),
+    Alloc = map_get(alloc, Anno),
     [#b_literal{val=Kind}|Args1] = Args0,
     case Kind of
         new ->
             [Dst,Size,{integer,Unit}] = beam_args([Dst0|Args1], St),
             Live = get_live(I),
-            {[Line|cg_bs_init(Dst, Size, Unit, Live, Fail)],St};
+            {[Line|cg_bs_init(Dst, Size, Alloc, Unit, Live, Fail)],St};
         private_append ->
             [Dst,Src,Bits,{integer,Unit}] = beam_args([Dst0|Args1], St),
             Flags = {field_flags,[]},
@@ -1037,7 +1065,7 @@ cg_block([#cg_set{op=bs_init,dst=Dst0,args=Args0,anno=Anno}=I,
             [Dst,Src,Bits,{integer,Unit}] = beam_args([Dst0|Args1], St),
             Flags = {field_flags,[]},
             Live = get_live(I),
-            Is = [Line,{bs_append,Fail,Bits,0,Live,Unit,Src,Flags,Dst}],
+            Is = [Line,{bs_append,Fail,Bits,Alloc,Live,Unit,Src,Flags,Dst}],
             {Is,St}
     end;
 cg_block([#cg_set{anno=Anno,op=bs_start_match,dst=Ctx0,args=[Bin0]}=I,
@@ -1531,13 +1559,13 @@ cg_bs_put(Fail, [{atom,Type},{literal,Flags}|Args]) ->
             [{Op,Fail,{field_flags,Flags},Src}]
     end.
 
-cg_bs_init(Dst, Size0, Unit, Live, Fail) ->
+cg_bs_init(Dst, Size0, Alloc, Unit, Live, Fail) ->
     Op = case Unit of
              1 -> bs_init_bits;
              8 -> bs_init2
          end,
     Size = cg_bs_init_size(Size0),
-    [{Op,Fail,Size,0,Live,{field_flags,[]},Dst}].
+    [{Op,Fail,Size,Alloc,Live,{field_flags,[]},Dst}].
 
 cg_bs_init_size({x,_}=R) -> R;
 cg_bs_init_size({y,_}=R) -> R;
diff --git a/lib/compiler/src/beam_utils.erl b/lib/compiler/src/beam_utils.erl
index 686d314c2d..39916a2af8 100644
--- a/lib/compiler/src/beam_utils.erl
+++ b/lib/compiler/src/beam_utils.erl
@@ -24,13 +24,11 @@
 -export([is_killed/3,is_killed_at/3,is_not_used/3,
 	 empty_label_index/0,index_label/3,index_labels/1,replace_labels/4,
 	 code_at/2,bif_to_test/3,is_pure_test/1,
-	 combine_heap_needs/2,
-	 split_even/1
-        ]).
+	 split_even/1]).
 
 -export_type([code_index/0,module_code/0,instruction/0]).
 
--import(lists, [flatmap/2,map/2,member/2,sort/1,reverse/1]).
+-import(lists, [map/2,member/2,sort/1,reverse/1]).
 
 -define(is_const(Val), (Val =:= nil orelse
                         element(1, Val) =:= integer orelse
@@ -218,19 +216,6 @@ is_pure_test({test,is_function2,_,[_,_]}) -> true;
 is_pure_test({test,Op,_,Ops}) -> 
     erl_internal:new_type_test(Op, length(Ops)).
 
-%% combine_heap_needs(HeapNeed1, HeapNeed2) -> HeapNeed
-%%  Combine the heap need for two allocation instructions.
-
--type heap_need_tag() :: 'floats' | 'words'.
--type heap_need() :: non_neg_integer() |
-                     {'alloc',[{heap_need_tag(),non_neg_integer()}]}.
--spec combine_heap_needs(heap_need(), heap_need()) -> heap_need().
-
-combine_heap_needs(H1, H2) when is_integer(H1), is_integer(H2) ->
-    H1 + H2;
-combine_heap_needs(H1, H2) ->
-    {alloc,combine_alloc_lists([H1,H2])}.
-
 %% split_even/1
 %% [1,2,3,4,5,6] -> {[1,3,5],[2,4,6]}
 
@@ -734,19 +719,6 @@ label(Old, D, Fb) ->
         _ -> Fb(Old)
     end.
 
-%% Help function for combine_heap_needs.
-
-combine_alloc_lists(Al0) ->
-    Al1 = flatmap(fun(Words) when is_integer(Words) ->
-                         [{words,Words}];
-                    ({alloc,List}) ->
-                         List
-                 end, Al0),
-    Al2 = sofs:relation(Al1),
-    Al3 = sofs:relation_to_family(Al2),
-    Al4 = sofs:to_external(Al3),
-    [{Tag,lists:sum(L)} || {Tag,L} <- Al4].
-
 %% live_opt/4.
 
 split_even([], Ss, Ds) ->