cleaning and better tests

src/passes/2-simplify/pascaligo.ml

@@ -668,37 +668,9 @@ and simpl_single_instruction : Raw.instruction -> (_ -> expression result) resul
       let%bind body = body None in
       return_statement @@ e_loop cond body
   | Loop (For (ForInt fi)) -> 
-      (* cond part *)
-      let%bind var = ok @@ e_variable fi.value.assign.value.name.value in
-      let%bind value = simpl_expression fi.value.assign.value.expr in
-      let%bind bound = simpl_expression fi.value.bound in
-      let%bind comp = match fi.value.down with
-        | Some _ -> ok @@ e_annotation (e_constant "GE" [var ; bound]) t_bool
-        | None -> ok @@ e_annotation (e_constant "LE" [var ; bound]) t_bool
-      in
-
-      (* body part *)
-      let%bind body = simpl_block fi.value.block.value in
-      let%bind body = body None in
-      let%bind step =  match fi.value.step with
-        | Some (_,e) -> simpl_expression e
-        | None -> ok (e_int 1) in
-      let%bind ctrl = match fi.value.down with
-        | Some _ ->
-          let%bind _addi = ok @@ e_constant "SUB" [ var ; step ] in 
-          ok @@ e_assign fi.value.assign.value.name.value [] _addi 
-        | None ->
-          let%bind _subi = ok @@ e_constant "ADD" [ var ; step ] in 
-          ok @@ e_assign fi.value.assign.value.name.value [] _subi
-      in
-      let rec add_to_seq expr = match expr.expression with
-        | E_sequence (_,a) -> add_to_seq a
-        | _ -> e_sequence body ctrl in
-      let%bind body' = ok @@ add_to_seq body in
-      let%bind loop = ok @@ e_loop comp body' in
-
-      return_statement @@
-        e_let_in (fi.value.assign.value.name.value, Some t_int) value loop
+      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
   | Cond c -> (
@@ -993,5 +965,35 @@ and simpl_statements : Raw.statements -> (_ -> expression result) result =
 and simpl_block : Raw.block -> (_ -> expression result) result = fun t ->
   simpl_statements t.statements
 
+and simpl_for_int : Raw.for_int -> (_ -> expression result) result = fun fi ->
+  (* cond part *)
+  let%bind var = ok @@ e_variable fi.assign.value.name.value in
+  let%bind value = simpl_expression fi.assign.value.expr in
+  let%bind bound = simpl_expression fi.bound in
+  let%bind comp = match fi.down with
+    | Some _ -> ok @@ e_annotation (e_constant "GE" [var ; bound]) t_bool
+    | None -> ok @@ e_annotation (e_constant "LE" [var ; bound]) t_bool
+  in
+  (* body part *)
+  let%bind body = simpl_block fi.block.value in
+  let%bind body = body None in
+  let%bind step =  match fi.step with
+    | Some (_,e) -> simpl_expression e
+    | None -> ok (e_int 1) in
+  let%bind ctrl = match fi.down with
+    | Some _ ->
+      let%bind _addi = ok @@ e_constant "SUB" [ var ; step ] in 
+      ok @@ e_assign fi.assign.value.name.value [] _addi 
+    | None ->
+      let%bind _subi = ok @@ e_constant "ADD" [ var ; step ] in 
+      ok @@ e_assign fi.assign.value.name.value [] _subi
+  in
+  let rec add_to_seq expr = match expr.expression with
+    | E_sequence (_,a) -> add_to_seq a
+    | _ -> e_sequence body ctrl in
+  let%bind body' = ok @@ add_to_seq body in
+  let%bind loop = ok @@ e_loop comp body' in
+  return_statement @@ e_let_in (fi.assign.value.name.value, Some t_int) value loop
+
 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

@@ -16,12 +16,34 @@ function while_sum (var n : nat) : nat is block {
   }
 } with r
 
-function for_sum (var n : nat) : nat is block {
-  for i := 1 to 100 step 1
+function for_sum_up (var n : nat) : int is block {
+  var acc : int := 0 ;
+  for i := 1 to int(n) step 1
     begin 
-      n := n + 1n ;
-    end }
-  with n 
+      acc := acc + i ;
+    end
+} with acc 
+
+function for_sum_down (var n : nat) : int is block {
+  var acc : int := 0 ;
+  for i := int(n) down to 1 step 1
+    begin 
+      acc := acc + i ;
+    end
+} with acc 
+
+function for_sum_step (var n : nat) : int is block {
+  var acc : int := 0 ;
+  var mystep : int := 2 ;
+  for i := 1 to int(n) step mystep
+    begin 
+      acc := acc + i ;
+    end;
+  for i := 0 to int(n) step mystep
+    begin 
+      acc := acc + i ;
+    end;
+} with acc 
 
 function dummy (const n : nat) : nat is block {
   while (False) block { skip }

src/test/integration_tests.ml

@@ -667,8 +667,18 @@ let loop () : unit result =
   let%bind () = 
 >>>>>>> First version for ForInt loops
     let make_input = e_nat in
-    let make_expected = fun n -> e_nat (n + 100) in
-    expect_eq_n_pos_mid program "for_sum" make_input make_expected
+    let make_expected = fun n -> e_int (n * (n + 1) / 2) in
+    expect_eq_n_pos_mid program "for_sum_up" make_input make_expected
+  in
+  let%bind () = 
+    let make_input = e_nat in
+    let make_expected = fun n -> e_int (n * (n + 1) / 2) in
+    expect_eq_n_pos_mid program "for_sum_down" make_input make_expected
+  in
+  let%bind () = 
+    let make_input = e_nat in
+    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
   ok ()