diff --git a/src/passes/2-simplify/pascaligo.ml b/src/passes/2-simplify/pascaligo.ml index 3a438f68b..f7b5b4211 100644 --- a/src/passes/2-simplify/pascaligo.ml +++ b/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 diff --git a/src/test/contracts/loop.ligo b/src/test/contracts/loop.ligo index 8fc4cd254..f3f0db33c 100644 --- a/src/test/contracts/loop.ligo +++ b/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 } diff --git a/src/test/integration_tests.ml b/src/test/integration_tests.ml index 0644f1189..d28f628b4 100644 --- a/src/test/integration_tests.ml +++ b/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 ()