diff --git a/src/passes/2-simplify/pascaligo.ml b/src/passes/2-simplify/pascaligo.ml index f7b5b4211..d1ffb5672 100644 --- a/src/passes/2-simplify/pascaligo.ml +++ b/src/passes/2-simplify/pascaligo.ml @@ -68,7 +68,7 @@ module Errors = struct ] in error ~data title message - let unsupported_for_loops region = + (* let unsupported_for_loops region = let title () = "bounded iterators" in let message () = Format.asprintf "only simple for loops are supported for now" in @@ -76,7 +76,7 @@ module Errors = struct ("loop_loc", fun () -> Format.asprintf "%a" Location.pp_lift @@ region) ] in - error ~data title message + error ~data title message *) let unsupported_non_var_pattern p = let title () = "pattern is not a variable" in @@ -671,8 +671,10 @@ and simpl_single_instruction : Raw.instruction -> (_ -> expression result) resul let%bind loop = simpl_for_int fi.value in let%bind loop = loop None in return_statement @@ loop - | Loop (For (ForCollect {region ; _})) -> - fail @@ unsupported_for_loops region + | Loop (For (ForCollect fc)) -> + let%bind loop = simpl_for_collect fc.value in + let%bind loop = loop None in + return_statement @@ loop | Cond c -> ( let (c , loc) = r_split c in let%bind expr = simpl_expression c.test in @@ -995,5 +997,18 @@ and simpl_for_int : Raw.for_int -> (_ -> expression result) result = fun fi -> let%bind loop = ok @@ e_loop comp body' in return_statement @@ e_let_in (fi.assign.value.name.value, Some t_int) value loop +and simpl_for_collect : Raw.for_collect -> (_ -> expression result) result = fun fc -> + let%bind col = simpl_expression fc.expr in + + let%bind body = simpl_block fc.block.value in + let%bind body = body None in + + let%bind invar = ok @@ e_variable fc.var.value in + let%bind letin = ok @@ e_let_in (fc.var.value, None) invar body in + let%bind lambda = ok @@ e_lambda fc.var.value None (Some t_unit) letin in + (* let%bind lambda = ok @@ e_lambda fc.var.value None (Some t_unit) body in *) + + return_statement @@ e_constant "SET_ITER" [col ; lambda] + let simpl_program : Raw.ast -> program result = fun t -> bind_list @@ List.map simpl_declaration @@ nseq_to_list t.decl diff --git a/src/test/contracts/loop.ligo b/src/test/contracts/loop.ligo index f3f0db33c..27af27a72 100644 --- a/src/test/contracts/loop.ligo +++ b/src/test/contracts/loop.ligo @@ -45,6 +45,14 @@ function for_sum_step (var n : nat) : int is block { end; } with acc +function for_collection (var n : set(int)) : int is block { + var acc : int := 0; + for i in n + begin + acc := acc + i ; + end; +} with acc + function dummy (const n : nat) : nat is block { while (False) block { skip } } with n diff --git a/src/test/integration_tests.ml b/src/test/integration_tests.ml index d28f628b4..f180ff7d5 100644 --- a/src/test/integration_tests.ml +++ b/src/test/integration_tests.ml @@ -680,6 +680,11 @@ let loop () : unit result = let make_expected = fun n -> e_int (n * (n + 1) / 2) in expect_eq_n_pos_mid program "for_sum_step" make_input make_expected in + let%bind () = + let make_input = fun _n -> e_set [e_int 1; e_int 2] in + let make_expected = fun _n -> e_int 3 in + expect_eq_n_pos_mid program "for_sum_step" make_input make_expected + in ok () (* Don't know how to assert parse error happens in this test framework