Added if guards and let statements in list comp

src/aeso_ast_infer_types.erl

@@ -1093,6 +1093,41 @@ infer_expr(Env, {list_comp, As, Yield, [{comprehension_bind, Arg, BExpr}|Rest]})
     , As
     , {list_comp, As, TypedYield, [{comprehension_bind, {typed, Arg, BindVarType}, TypedBind}|TypedRest]}
     , ResType};
+infer_expr(Env, {list_comp, AttrsL, Yield, [{comprehension_if, AttrsIF, Cond}|Rest]}) ->
+    NewCond = check_expr(Env, Cond, {id, AttrsIF, "bool"}),
+    {typed, _, {list_comp, _, TypedYield, TypedRest}, ResType} =
+        infer_expr(Env, {list_comp, AttrsL, Yield, Rest}),
+    { typed
+    , AttrsL
+    , {list_comp, AttrsL, TypedYield, [{comprehension_if, AttrsIF, NewCond}|TypedRest]}
+    , ResType};
+infer_expr(Env, {list_comp, AsLC, Yield, [{letval, AsLV, Pattern, Type, E}|Rest]}) ->
+    NewE = {typed, _, _, PatType} = infer_expr(Env, {typed, AsLV, E, arg_type(Type)}),
+    BlockType = fresh_uvar(AsLV),
+    {'case', _, NewPattern, NewRest} =
+        infer_case( Env
+                  , AsLC
+                  , Pattern
+                  , PatType
+                  , {list_comp, AsLC, Yield, Rest}
+                  , BlockType),
+    {typed, _, {list_comp, _, TypedYield, TypedRest}, ResType} = NewRest,
+    { typed
+    , AsLC
+    , {list_comp, AsLC, TypedYield, [{letval, AsLV, NewPattern, Type, NewE}|TypedRest]}
+    , ResType
+    };
+infer_expr(Env, {list_comp, AsLC, Yield, [Def={letfun, AsLF, _, _, _, _}|Rest]}) ->
+    {{Name, TypeSig}, LetFun} = infer_letfun(Env, Def),
+    FunT = freshen_type(AsLF, typesig_to_fun_t(TypeSig)),
+    NewE = bind_var({id, AsLF, Name}, FunT, Env),
+    {typed, _, {list_comp, _, TypedYield, TypedRest}, ResType} =
+        infer_expr(NewE, {list_comp, AsLC, Yield, Rest}),
+    { typed
+    , AsLC
+    , {list_comp, AsLC, TypedYield, [LetFun|TypedRest]}
+    , ResType
+    };
 infer_expr(Env, {typed, As, Body, Type}) ->
     Type1 = check_type(Env, Type),
     {typed, _, NewBody, NewType} = check_expr(Env, Body, Type1),

src/aeso_ast_to_fcode.erl

@@ -451,6 +451,15 @@ expr_to_fcode(Env, _Type, {list_comp, As, Yield, [{comprehension_bind, {typed, {
     Bind = {lam, [Arg], expr_to_fcode(Env1, {list_comp, As, Yield, Rest})},
     {def_u, FlatMap, _} = resolve_fun(Env, ["List", "flat_map"]),
     {def, FlatMap, [Bind, expr_to_fcode(Env, BindExpr)]};
+expr_to_fcode(Env, Type, {list_comp, As, Yield, [{comprehension_if, _, Cond}|Rest]}) ->
+    make_if(expr_to_fcode(Env, Cond),
+            expr_to_fcode(Env, Type, {list_comp, As, Yield, Rest}),
+            nil
+           );
+expr_to_fcode(Env, Type, {list_comp, As, Yield, [LV = {letval, _, _, _, _}|Rest]}) ->
+    expr_to_fcode(Env, Type, {block, As, [LV, {list_comp, As, Yield, Rest}]});
+expr_to_fcode(Env, Type, {list_comp, As, Yield, [LF = {letfun, _, _, _, _, _}|Rest]}) ->
+    expr_to_fcode(Env, Type, {block, As, [LF, {list_comp, As, Yield, Rest}]});
 
 %% Conditionals
 expr_to_fcode(Env, _Type, {'if', _, Cond, Then, Else}) ->

src/aeso_ast_to_icode.erl

@@ -531,6 +531,12 @@ ast_body({list_comp, As, Yield, [{comprehension_bind, {typed, Arg, ArgType}, Bin
               , ast_body(BindExpr, Icode)
               ]
         };
+ast_body({list_comp, As, Yield, [{comprehension_if, AsIF, Cond}|Rest]}, Icode) ->
+    ast_body({'if', AsIF, Cond, {list_comp, As, Yield, Rest}, {list, As, []}}, Icode);
+ast_body({list_comp, As, Yield, [LV = {letval, _, _, _, _}|Rest]}, Icode) ->
+    ast_body({block, As, [LV, {list_comp, As, Yield, Rest}]}, Icode);
+ast_body({list_comp, As, Yield, [LF = {letfun, _, _, _, _, _}|Rest]}, Icode) ->
+    ast_body({block, As, [LF, {list_comp, As, Yield, Rest}]}, Icode);
 ast_body({'if',_,Dec,Then,Else}, Icode) ->
     #ifte{decision = ast_body(Dec, Icode)
          ,then     = ast_body(Then, Icode)

src/aeso_parser.erl

@@ -246,12 +246,22 @@ exprAtom() ->
         , {bool, keyword(false), false}
         , ?LET_P(Fs, brace_list(?LAZY_P(field_assignment())), record(Fs))
         , {list, [], bracket_list(Expr)}
-        , ?RULE(keyword('['), Expr, token('|'), comma_sep(?LAZY_P(comprehension_bind())), tok(']'), list_comp_e(_1, _2, _4))
+        , ?RULE(keyword('['), Expr, token('|'), comma_sep(comprehension_exp()), tok(']'), list_comp_e(_1, _2, _4))
         , ?RULE(tok('['), Expr, binop('..'), Expr, tok(']'), _3(_2, _4))
         , ?RULE(keyword('('), comma_sep(Expr), tok(')'), tuple_e(_1, _2))
         ])
     end).
 
+comprehension_exp() ->
+    ?LAZY_P(choice(
+      [ comprehension_bind()
+      , letdecl()
+      , comprehension_if()
+      ])).
+
+comprehension_if() ->
+    ?RULE(keyword('if'), parens(expr()), {comprehension_if, _1, _2}).
+
 comprehension_bind() ->
     ?RULE(id(), tok('<-'), expr(), {comprehension_bind, _1, _3}).
 

src/aeso_stdlib.erl

@@ -67,7 +67,7 @@ namespace Func =
 std_list() ->"
 namespace List =
 
-  function empty(l : list('a)) : bool = switch(l)
+  function is_empty(l : list('a)) : bool = switch(l)
     [] => true
     _  => false
 

src/aeso_syntax.erl

@@ -92,7 +92,7 @@
      | {proj, ann(), expr(), id()}
      | {tuple, ann(), [expr()]}
      | {list, ann(), [expr()]}
-     | {list_comp, ann(), expr(), [comprehension_bind()]}
+     | {list_comp, ann(), expr(), [comprehension_exp()]}
      | {typed, ann(), expr(), type()}
      | {record, ann(), [field(expr())]}
      | {record, ann(), expr(), [field(expr())]} %% record update
@@ -105,7 +105,9 @@
      | id() | qid() | con() | qcon()
      | constant().
 
--type comprehension_bind() :: [{comprehension_bind, ann(), id(), type()}].
+-type comprehension_exp() :: [{ comprehension_bind, ann(), id(), expr()}
+                             | {comprehension_if, expr()}
+                             | letbind()].
 
 -type arg_expr() :: expr() | {named_arg, ann(), id(), expr()}.
 

test/contracts/list_comp.aes

@@ -12,3 +12,12 @@ contract ListComp =
     entrypoint l3() = [x ++ y | x <- [[":)"] | x <- [1,2]]
                               , y <- [[":("]]]  
     entrypoint l3_true() = [[":)", ":("], [":)", ":("]]
+
+    entrypoint l4() = [(a, b, c) | let is_pit(a, b, c) = a*a + b*b == c*c
+                                 , let base = [1,2,3,4,5,6,7,8,9,10]
+                                 , a <- base
+                                 , b <- base, if (b >= a)
+                                 , c <- base, if (c >= b)
+                                 , if (is_pit(a, b, c))
+                                 ]
+    entrypoint l4_true() = [(3, 4, 5), (6, 8, 10)]