beam_ssa_type: Infer types on switch failure

If we know that the checked value is a singleton at the fail
block, it's safe to infer types as if we've made a direct
comparison against that value.

2 files changed, 84 insertions, 56 deletions

diff --git a/lib/compiler/src/beam_ssa_type.erl b/lib/compiler/src/beam_ssa_type.erl
index 8b7a31849c..0912ecb2c2 100644
--- a/lib/compiler/src/beam_ssa_type.erl
+++ b/lib/compiler/src/beam_ssa_type.erl
@@ -814,18 +814,9 @@ prune_switch_list([], _, _, _) -> [].
 update_successors(#b_br{bool=#b_literal{val=true},succ=Succ}=Last, Ts, D0) ->
     {Last, update_successor(Succ, Ts, D0)};
 update_successors(#b_br{bool=#b_var{}=Bool,succ=Succ,fail=Fail}=Last0,
-                  Ts0, D0) ->
-    Ts = case cerl_sets:is_element(Bool, D0#d.once) of
-             true ->
-                 %% This variable is defined in this block and is only
-                 %% referenced by this br terminator. Therefore, there is
-                 %% no need to include it in the type database passed on to
-                 %% the successors of this block.
-                 maps:remove(Bool, Ts0);
-             false ->
-                 Ts0
-         end,
-    case infer_types_br(Bool, Ts, D0) of
+                  Ts, D0) ->
+    UsedOnce = cerl_sets:is_element(Bool, D0#d.once),
+    case infer_types_br(Bool, Ts, UsedOnce, D0) of
         {#{}=SuccTs, #{}=FailTs} ->
             D1 = update_successor(Succ, SuccTs, D0),
             D = update_successor(Fail, FailTs, D1),
@@ -841,15 +832,12 @@ update_successors(#b_switch{arg=#b_var{}=V,fail=Fail0,list=List0}=Last0,
                   Ts, D0) ->
     UsedOnce = cerl_sets:is_element(V, D0#d.once),
 
-    {List1, D1} = update_switch(List0, V, UsedOnce, Ts, [], D0),
-
-    FailTs = case UsedOnce of
-                 true -> maps:remove(V, Ts);
-                 false -> subtract_sw_list(V, List0, Ts)
-             end,
+    {List1, D1} = update_switch(List0, V, Ts, UsedOnce, [], D0),
+    FailTs = update_switch_failure(V, List0, Ts, UsedOnce, D1),
 
     case FailTs of
         none ->
+            %% The fail block is unreachable; swap it with one of the choices.
             [{_, Fail} | List] = List1,
             Last = Last0#b_switch{fail=Fail,list=List},
             {Last, D1};
@@ -871,25 +859,33 @@ update_successors(#b_ret{arg=Arg}=Last, Ts, D0) ->
         end,
     {Last, D}.
 
-update_switch([{Val, Lbl}=Sw | List], V, UsedOnce, Ts, Acc, D0) ->
-    case infer_types_switch(V, Val, Ts, D0) of
+update_switch([{Val, Lbl}=Sw | List], V, Ts, UsedOnce, Acc, D0) ->
+    case infer_types_switch(V, Val, Ts, UsedOnce, D0) of
         none ->
-            update_switch(List, V, UsedOnce, Ts, Acc, D0);
-        SwTs0 ->
-            SwTs = case UsedOnce of
-                       true -> maps:remove(V, SwTs0);
-                       false -> SwTs0
-                   end,
+            update_switch(List, V, Ts, UsedOnce, Acc, D0);
+        SwTs ->
             D = update_successor(Lbl, SwTs, D0),
-            update_switch(List, V, UsedOnce, Ts, [Sw | Acc], D)
+            update_switch(List, V, Ts, UsedOnce, [Sw | Acc], D)
     end;
-update_switch([], _V, _UsedOnce, _Ts, Acc, D) ->
-    {Acc, D}.
+update_switch([], _V, _Ts, _UsedOnce, Acc, D) ->
+    {reverse(Acc), D}.
 
-subtract_sw_list(V, List, Ts) ->
+update_switch_failure(V, List, Ts, UsedOnce, D) ->
     case sub_sw_list_1(raw_type(V, Ts), List, Ts) of
-        none -> none;
-        Type -> Ts#{ V := Type }
+        none ->
+            none;
+        FailType ->
+            case beam_types:get_singleton_value(FailType) of
+                {ok, Value} ->
+                    %% This is the only possible value at the fail label, so we
+                    %% can infer types as if we matched it directly.
+                    Lit = #b_literal{val=Value},
+                    infer_types_switch(V, Lit, Ts, UsedOnce, D);
+                error when UsedOnce ->
+                    ts_remove_var(V, Ts);
+                error ->
+                    Ts
+            end
     end.
 
 sub_sw_list_1(Type, [{Val,_}|T], Ts) ->
@@ -1291,39 +1287,51 @@ raw_type(V, Ts) ->
 %%  'cons' would give 'nil' as the only possible type. The result of the
 %%  subtraction for L will be added to FailTypes.
 
-infer_types_br(#b_var{}=V, Ts, #d{ds=Ds}) ->
+infer_types_br(#b_var{}=V, Ts, UsedOnce, #d{ds=Ds}) ->
     #{V:=#b_set{op=Op,args=Args}} = Ds,
 
     {PosTypes, NegTypes} = infer_type(Op, Args, Ts, Ds),
 
-    SuccTs1 = meet_types(PosTypes, Ts),
-    FailTs1 = subtract_types(NegTypes, Ts),
-
-    SuccTs = infer_br_value(V, Ts, true, SuccTs1),
-    FailTs = infer_br_value(V, Ts, false, FailTs1),
+    SuccTs0 = meet_types(PosTypes, Ts),
+    FailTs0 = subtract_types(NegTypes, Ts),
 
-    {SuccTs, FailTs}.
+    case UsedOnce of
+        true ->
+            %% The branch variable is defined in this block and is only
+            %% referenced by this terminator. Therefore, there is no need to
+            %% include it in the type database passed on to the successors of
+            %% of this block.
+            SuccTs = ts_remove_var(V, SuccTs0),
+            FailTs = ts_remove_var(V, FailTs0),
+            {SuccTs, FailTs};
+        false ->
+            SuccTs = infer_br_value(V, true, SuccTs0),
+            FailTs = infer_br_value(V, false, FailTs0),
+            {SuccTs, FailTs}
+    end.
 
-infer_br_value(_V, _OldTs, _Bool, none) ->
+infer_br_value(_V, _Bool, none) ->
     none;
-infer_br_value(V, OldTs, Bool, NewTs) ->
-    case OldTs of
-        #{ V := T } ->
-            case beam_types:is_boolean_type(T) of
-                true ->
-                    NewTs#{ V := beam_types:make_atom(Bool) };
-                false ->
-                    %% V is a try/catch tag or similar, leave it alone.
-                    NewTs
-            end;
-        #{} ->
-            %% V is only used in this branch; don't propagate its type.
+infer_br_value(V, Bool, NewTs) ->
+    #{ V := T } = NewTs,
+    case beam_types:is_boolean_type(T) of
+        true ->
+            NewTs#{ V := beam_types:make_atom(Bool) };
+        false ->
+            %% V is a try/catch tag or similar, leave it alone.
             NewTs
     end.
 
-infer_types_switch(V, Lit, Ts, #d{ds=Ds}) ->
-    {PosTypes, _} = infer_type({bif,'=:='}, [V, Lit], Ts, Ds),
-    meet_types(PosTypes, Ts).
+infer_types_switch(V, Lit, Ts0, UsedOnce, #d{ds=Ds}) ->
+    {PosTypes, _} = infer_type({bif,'=:='}, [V, Lit], Ts0, Ds),
+    Ts = meet_types(PosTypes, Ts0),
+    case UsedOnce of
+        true -> ts_remove_var(V, Ts);
+        false -> Ts
+    end.
+
+ts_remove_var(_V, none) -> none;
+ts_remove_var(V, Ts) -> maps:remove(V, Ts).
 
 infer_type(succeeded, [#b_var{}=Src], Ts, Ds) ->
     #b_set{op=Op,args=Args} = maps:get(Src, Ds),
diff --git a/lib/compiler/src/beam_validator.erl b/lib/compiler/src/beam_validator.erl
index 923a0b4a4e..afede2b54d 100644
--- a/lib/compiler/src/beam_validator.erl
+++ b/lib/compiler/src/beam_validator.erl
@@ -1531,8 +1531,21 @@ validate_select_val(Fail, [Val,{f,L}|T], Src, Vst0) ->
                          update_ne_types(Src, Val, FailVst)
                  end),
     validate_select_val(Fail, T, Src, Vst);
-validate_select_val(Fail, [], _, Vst) ->
+validate_select_val(Fail, [], Src, Vst) ->
     branch(Fail, Vst,
+           fun(FailVst) ->
+                    FailType = get_term_type(Src, FailVst),
+                    case beam_types:get_singleton_value(FailType) of
+                       {ok, Value} ->
+                           %% This is the only possible value at the fail
+                           %% label, so we can infer types as if we matched it
+                           %% directly.
+                           Lit = value_to_literal(Value),
+                           update_eq_types(Src, Lit, FailVst);
+                       error ->
+                           FailVst
+                   end
+           end,
            fun(SuccVst) ->
                    %% The next instruction is never executed.
                    kill_state(SuccVst)
@@ -1921,6 +1934,13 @@ is_literal({integer,I}) when is_integer(I) -> true;
 is_literal({literal,_L}) -> true;
 is_literal(_) -> false.
 
+%% `dialyzer` complains about the float and general literal cases never being
+%% matched and I don't like suppressing warnings. Should they become possible
+%% I'm sure `dialyzer` will warn about it.
+value_to_literal([]) -> nil;
+value_to_literal(A) when is_atom(A) -> {atom,A};
+value_to_literal(I) when is_integer(I) -> {integer,I}.
+
 %% These are just wrappers around their equivalents in beam_types, which
 %% handle the validator-specific #t_abstract{} type.
 %%