From f437ee564e865695dcb0bbd3e3760adf45dcaea7 Mon Sep 17 00:00:00 2001
From: Gaith Hallak <gaithhallak@gmail.com>
Date: Mon, 30 Aug 2021 11:41:12 +0300
Subject: [PATCH] Add pattern guards to infer types and fcode generation

---
 src/aeso_ast_infer_types.erl | 23 ++++++++++++++++++++---
 src/aeso_ast_to_fcode.erl    | 21 ++++++++++++++++-----
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/src/aeso_ast_infer_types.erl b/src/aeso_ast_infer_types.erl
index 1e6baf8..1c4a068 100644
--- a/src/aeso_ast_infer_types.erl
+++ b/src/aeso_ast_infer_types.erl
@@ -1616,8 +1616,16 @@ infer_expr(Env, {'if', Attrs, Cond, Then, Else}) ->
 infer_expr(Env, {switch, Attrs, Expr, Cases}) ->
     NewExpr = {typed, _, _, ExprType} = infer_expr(Env, Expr),
     SwitchType = fresh_uvar(Attrs),
-    NewCases = [infer_case(Env, As, Pattern, ExprType, Branch, SwitchType)
-                || {'case', As, Pattern, Branch} <- Cases],
+    ApplyInferCase =
+        fun(Case) ->
+            case Case of
+                {'case', As, Pattern, Branch} ->
+                    infer_case(Env, As, Pattern, ExprType, Branch, SwitchType);
+                {'case', As, Pattern, Guard, Branch} ->
+                    infer_case(Env, As, Pattern, Guard, ExprType, Branch, SwitchType)
+            end
+        end,
+    NewCases = lists:map(ApplyInferCase, Cases),
     {typed, Attrs, {switch, Attrs, NewExpr, NewCases}, SwitchType};
 infer_expr(Env, {record, Attrs, Fields}) ->
     RecordType = fresh_uvar(Attrs),
@@ -1837,10 +1845,19 @@ infer_pattern(Env, Pattern) ->
     {NewEnv#env{ in_pattern = Env#env.in_pattern }, NewPattern}.
 
 infer_case(Env, Attrs, Pattern, ExprType, Branch, SwitchType) ->
+    infer_case(Env, Attrs, Pattern, none, ExprType, Branch, SwitchType).
+
+infer_case(Env, Attrs, Pattern, Guard, ExprType, Branch, SwitchType) ->
     {NewEnv, NewPattern = {typed, _, _, PatType}} = infer_pattern(Env, Pattern),
     NewBranch  = check_expr(NewEnv#env{ in_pattern = false }, Branch, SwitchType),
     unify(Env, PatType, ExprType, {case_pat, Pattern, PatType, ExprType}),
-    {'case', Attrs, NewPattern, NewBranch}.
+    case Guard of
+        none ->
+            {'case', Attrs, NewPattern, NewBranch};
+        _ ->
+            NewGuard = check_expr(NewEnv, Guard, {id, Attrs, "bool"}),
+            {'case', Attrs, NewPattern, NewGuard, NewBranch}
+    end.
 
 %% NewStmts = infer_block(Env, Attrs, Stmts, BlockType)
 infer_block(_Env, Attrs, [], BlockType) ->
diff --git a/src/aeso_ast_to_fcode.erl b/src/aeso_ast_to_fcode.erl
index 69378d5..b5a43b8 100644
--- a/src/aeso_ast_to_fcode.erl
+++ b/src/aeso_ast_to_fcode.erl
@@ -683,9 +683,9 @@ expr_to_fcode(Env, _Type, {'if', _, Cond, Then, Else}) ->
             expr_to_fcode(Env, Else));
 
 %% Switch
-expr_to_fcode(Env, _, {switch, _, Expr = {typed, _, E, Type}, Alts}) ->
+expr_to_fcode(Env, _, S = {switch, _, Expr = {typed, _, E, Type}, Alts}) ->
     Switch = fun(X) ->
-                {switch, alts_to_fcode(Env, type_to_fcode(Env, Type), X, Alts)}
+                {switch, alts_to_fcode(Env, type_to_fcode(Env, Type), X, Alts, S)}
              end,
     case E of
         {id, _, X} -> Switch(X);
@@ -863,9 +863,9 @@ is_first_order(_)                      -> true.
 
 %% -- Pattern matching --
 
--spec alts_to_fcode(env(), ftype(), var_name(), [aeso_syntax:alt()]) -> fsplit().
-alts_to_fcode(Env, Type, X, Alts) ->
-    FAlts = [alt_to_fcode(Env, Alt) || Alt <- Alts],
+-spec alts_to_fcode(env(), ftype(), var_name(), [aeso_syntax:alt()], aeso_syntax:expr()) -> fsplit().
+alts_to_fcode(Env, Type, X, Alts, Switch) ->
+    FAlts = remove_guards(Env, Alts, Switch),
     split_tree(Env, [{X, Type}], FAlts).
 
 %% Intermediate format before case trees (fcase() and fsplit()).
@@ -879,6 +879,17 @@ alts_to_fcode(Env, Type, X, Alts) ->
               | {con, arities(), tag(), [fpat()]}
               | {assign, fpat(), fpat()}.
 
+remove_guards(_Env, [], _Switch) ->
+    [];
+remove_guards(Env, [Alt = {'case', _, _, _} | Rest], Switch) ->
+    [alt_to_fcode(Env, Alt) | remove_guards(Env, Rest, Switch)];
+remove_guards(Env, [{'case', _, Pat, Guard, Body} | Rest], {switch, Ann, Expr, _}) ->
+    FPat  = pat_to_fcode(Env, Pat),
+    FSwitch = expr_to_fcode(Env, {switch, Ann, Expr, Rest}),
+    FGuard = expr_to_fcode(bind_vars(Env, pat_vars(FPat)), Guard),
+    FBody = expr_to_fcode(bind_vars(Env, pat_vars(FPat)), Body),
+    [{'case', [FPat], make_if(FGuard, FBody, FSwitch)}].
+
 %% %% Invariant: the number of variables matches the number of patterns in each falt.
 -spec split_tree(env(), [{var_name(), ftype()}], [falt()]) -> fsplit().
 split_tree(_Env, _Vars, []) ->