Pattern matching on booleans

src/aeso_ast_to_fcode.erl

@@ -38,13 +38,15 @@
                | {nosplit, [var_name()], fexpr()}.
 
 -type fsplit_case() :: {'case', fsplit_pat(), fcase()}.
--type fsplit_pat()  :: {tuple, [var_name()]}.
+-type fsplit_pat()  :: {bool, false | true}
+                     | {tuple, [var_name()]}.
 
 -type fdefault() :: nodefault | {default, fcase()}.
 
 %% Intermediate format before case trees (fcase() and fsplit()).
 -type falt() :: {'case', [fpat()], fexpr()}.
 -type fpat() :: {var, var_name()}
+              | {bool, false | true}
               | {tuple, [fpat()]}.
 
 -type ftype() :: aeb_fate_data:fate_type_type().
@@ -238,6 +240,7 @@ split_tree(Env, Vars, Alts) ->
     end.
 
 -spec split_vars(fsplit_pat(), ftype()) -> [{var_name(), ftype()}].
+split_vars({bool, _}, boolean) -> [];
 split_vars({tuple, Xs}, {tuple, Ts}) ->
     lists:zip(Xs, Ts).
 
@@ -263,6 +266,7 @@ split_alt(I, {'case', Pats, Body}) ->
 
 -spec split_pat(fpat()) -> {fsplit_pat() | default, [fpat()]}.
 split_pat({var, X})      -> {default, [{var, X}]};
+split_pat({bool, B})     -> {{bool, B}, []};
 split_pat({tuple, Pats}) ->
     Var = fun({var, X}) -> X; (_) -> fresh_name() end,
     Xs = [Var(P) || P <- Pats],
@@ -271,7 +275,9 @@ split_pat({tuple, Pats}) ->
 -spec group_by_split_pat([{fsplit_pat() | default, falt()}]) -> [{fsplit_pat(), [falt()]}].
 group_by_split_pat(Alts) ->
     Tag = fun(default)    -> default;
-             ({tuple, _}) -> tuple end,
+             ({tuple, _}) -> tuple;
+             ({bool, B})  -> B
+          end,
     Grouped = maps:values(lists:foldr(
         fun({Pat, _} = Alt, Map) ->
             maps:update_with(Tag(Pat), fun(As) -> [Alt | As] end, [Alt], Map)
@@ -292,6 +298,8 @@ pat_to_fcode(Env, Pat) ->
 pat_to_fcode(_Env, _Type, {id, _, X}) -> {var, X};
 pat_to_fcode(Env, _Type, {tuple, _, Pats}) ->
     {tuple, [ pat_to_fcode(Env, Pat) || Pat <- Pats ]};
+pat_to_fcode(_Env, _Type, {bool, _, B}) ->
+    {bool, B};
 pat_to_fcode(_Env, Type, Pat) -> {todo, Pat, ':', Type}.
 
 -spec stmts_to_fcode(env(), [aeso_syntax:stmt()]) -> fexpr().
@@ -307,6 +315,7 @@ stmts_to_fcode(Env, [Expr]) ->
 %% - Deadcode elimination
 %% - Unused variable analysis (replace by _)
 %% - Simplified case trees (FATE has special instructions for shallow matching)
+%% - Case specialization
 %% - Constant propagation
 
 %% -- Helper functions -------------------------------------------------------

src/aeso_fcode_to_fate.erl

@@ -111,6 +111,9 @@ lookup_var(Env = #env{ args = Args, locals = Locals }, X) ->
 to_scode(_Env, {integer, N}) ->
     [aeb_fate_code:push(?i(N))];    %% Doesn't exist (yet), translated by desugaring
 
+to_scode(_Env, {bool, B}) ->
+    [aeb_fate_code:push(?i(B))];
+
 to_scode(Env, {var, X}) ->
     [aeb_fate_code:push(lookup_var(Env, X))];
 
@@ -146,6 +149,16 @@ case_to_scode(Env, {split, _Type, X, [{'case', {tuple, Xs}, Case}], nodefault})
     {Code, Env1} = match_tuple(Env, Xs),
     [aeb_fate_code:push(lookup_var(Env, X)),
      Code, case_to_scode(Env1, Case)];
+case_to_scode(Env, Split = {split, boolean, X, Cases, nodefault}) ->
+    Then = lists:keyfind({bool, true},  2, Cases),
+    Else = lists:keyfind({bool, false}, 2, Cases),
+    case {Then, Else} of
+        {{'case', _, ThenSplit}, {'case', _, ElseSplit}} ->
+            [aeb_fate_code:push(lookup_var(Env, X)),
+             {ifte, case_to_scode(Env, ThenSplit),
+                    case_to_scode(Env, ElseSplit)}];
+        _ -> ?TODO({'case', Split})
+    end;
 case_to_scode(_, Split = {split, _, _, _, _}) ->
     ?TODO({'case', Split}).
 
@@ -473,7 +486,8 @@ apply_rules_once([{RName, Rule} | Rules], I, Code) ->
 merge_rules() ->
     [?RULE(r_push_consume),
      ?RULE(r_one_shot_var),
-     ?RULE(r_write_to_dead_var)
+     ?RULE(r_write_to_dead_var),
+     ?RULE(r_write_single_branch)
     ].
 
 rules() ->
@@ -481,7 +495,7 @@ rules() ->
     [?RULE(r_dup_to_push),
      ?RULE(r_swap_push),
      ?RULE(r_swap_write),
-     ?RULE(r_inline)
+     ?RULE(r_inline_store)
     ].
 
 %% Removing pushes that are immediately consumed.
@@ -530,27 +544,28 @@ r_swap_write(IA = {_, I}, [JA = {_, J} | Code]) ->
     end;
 r_swap_write(_, _) -> false.
 
-r_swap_write(Pre, IA = {_, I}, Code0 = [JA = {_, J} | Code]) ->
-    case apply_rules_once(merge_rules(), IA, Code0) of
-        {_Rule, New, Rest} ->
+r_swap_write(Pre, IA = {_, I}, Code0 = [JA | Code]) ->
+    case {apply_rules_once(merge_rules(), IA, Code0), JA} of
+        {{_Rule, New, Rest}, _} ->
             {lists:reverse(Pre) ++ New, Rest};
-        false ->
+        {false, {_, J}} ->
             case independent(I, J) of
                 false -> false;
                 true  ->
                     {J1, I1} = swap_instrs(IA, JA),
                     r_swap_write([J1 | Pre], I1, Code)
-            end
+            end;
+        _ -> false
     end;
-r_swap_write(_, _, []) -> false.
+r_swap_write(_, _, _) -> false.
 
 %% Inline stores
-r_inline(I = {_, {'STORE', R = {var, _}, A = {arg, _}}}, Code) ->
+r_inline_store(I = {_, {'STORE', R = {var, _}, A = {arg, _}}}, Code) ->
     %% Not when A is var unless updating the annotations properly.
-    r_inline([I], R, A, Code);
-r_inline(_, _) -> false.
+    r_inline_store([I], R, A, Code);
+r_inline_store(_, _) -> false.
 
-r_inline(Acc, R, A, [{Ann, I} | Code]) ->
+r_inline_store(Acc, R, A, [{Ann, I} | Code]) ->
     #{ write := W, pure := Pure } = attributes(I),
     Inl = fun(X) when X == R -> A; (X) -> X end,
     case not live_in(R, Ann) orelse not Pure orelse lists:member(W, [R, A]) of
@@ -559,14 +574,14 @@ r_inline(Acc, R, A, [{Ann, I} | Code]) ->
             case I of
                 {Op, S, B, C} when ?IsBinOp(Op), B == R orelse C == R ->
                     Acc1 = [{Ann, {Op, S, Inl(B), Inl(C)}} | Acc],
-                    case r_inline(Acc1, R, A, Code) of
+                    case r_inline_store(Acc1, R, A, Code) of
                         false -> {lists:reverse(Acc1), Code};
                         {New, Rest} -> {New, Rest}
                     end;
-                _ -> r_inline([{Ann, I} | Acc], R, A, Code)
+                _ -> r_inline_store([{Ann, I} | Acc], R, A, Code)
             end
     end;
-r_inline(_Acc, _, _, []) -> false.
+r_inline_store(_Acc, _, _, _) -> false.
 
 %% Shortcut write followed by final read
 r_one_shot_var({Ann1, {Op, R = {var, _}, A, B}}, [{Ann2, J} | Code]) when ?IsBinOp(Op) ->
@@ -604,6 +619,22 @@ r_write_to_dead_var({Ann, {'STORE', R = {var, _}, A}}, Code) when A /= ?a ->
     end;
 r_write_to_dead_var(_, _) -> false.
 
+%% Push variable writes that are only needed in a single branch inside the branch.
+r_write_single_branch(IA = {_Ann, I}, [{ifte, Then = [{AnnThen, _} | _], Else = [{AnnElse, _} | _]} | Code]) ->
+    #{ write := R } = attributes(I),
+    case R of
+        {var, _} ->
+            case {live_in(R, AnnThen), live_in(R, AnnElse)} of
+                {true, false} ->
+                    {[], [{ifte, [IA | Then], Else} | Code]};
+                {false, true} ->
+                    {[], [{ifte, Then, [IA | Else]} | Code]};
+                _ -> false
+            end;
+        _ -> false
+    end;
+r_write_single_branch(_, _) -> false.
+
 
 %% Desugar and specialize and remove annotations
 unannotate({ifte, Then, Else}) -> [{ifte, unannotate(Then), unannotate(Else)}];