not complete for collect tryout

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

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

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