Moved ITE and loops treatment to cps_conversion (#424)

Author: Keryan-dev

middle_end/flambda/flambda_middle_end.ml

@@ -113,15 +113,14 @@ let middle_end0 ppf ~prefixname ~backend ~filename ~module_ident
       ~module_block_size_in_words ~module_initializer =
   Misc.Color.setup !Clflags.color;
   Profile.record_call "flambda.0" (fun () ->
-    let prepared_lambda, recursive_static_catches =
+    let prepared_lambda =
       Profile.record_call "prepare_lambda" (fun () ->
         Prepare_lambda.run module_initializer)
     in
     print_prepared_lambda ppf prepared_lambda;
     let ilambda =
       Profile.record_call "cps_conversion" (fun () ->
-        Cps_conversion.lambda_to_ilambda prepared_lambda
-          ~recursive_static_catches)
+        Cps_conversion.lambda_to_ilambda prepared_lambda)
     in
     print_ilambda ppf ilambda;
     let ilambda =

middle_end/flambda/from_lambda/cps_conversion.ml

@@ -58,6 +58,14 @@ let try_stack_at_handler = ref Continuation.Map.empty
 let recursive_static_catches = ref Numbers.Int.Set.empty
 let mutable_variables = ref Ident.Set.empty
 
+let mark_as_recursive_static_catch cont =
+  if Numbers.Int.Set.mem cont !recursive_static_catches then begin
+    Misc.fatal_errorf "Static catch with continuation %d already marked as \
+        recursive -- is it being redefined?"
+      cont
+  end;
+  recursive_static_catches := Numbers.Int.Set.add cont !recursive_static_catches
+
 let _print_stack ppf stack =
   Format.fprintf ppf "%a"
     (Format.pp_print_list ~pp_sep:(fun ppf () -> Format.fprintf ppf "; ")
@@ -116,7 +124,7 @@ let compile_staticfail ~(continuation : Continuation.t) ~args =
   in
   mk_poptraps (I.Apply_cont (continuation, None, args))
 
-let switch_for_if_then_else ~cond ~ifso ~ifnot k =
+let switch_for_if_then_else ~cond ~ifso ~ifnot =
   (* CR mshinwell: We need to make sure that [cond] is {0, 1}-valued.
      The frontend should have been fixed on this branch for this. *)
   let switch : Lambda.lambda_switch =
@@ -128,7 +136,7 @@ let switch_for_if_then_else ~cond ~ifso ~ifnot k =
       sw_tags_to_sizes = Tag.Scannable.Map.empty;
     }
   in
-  k (L.Lswitch (cond, switch, Loc_unknown))
+  L.Lswitch (cond, switch, Loc_unknown)
 
 let transform_primitive (prim : L.primitive) args loc =
   match prim, args with
@@ -141,8 +149,7 @@ let transform_primitive (prim : L.primitive) args loc =
           switch_for_if_then_else
             ~cond:(L.Lvar cond)
             ~ifso:(L.Lvar const_true)
-            ~ifnot:arg2
-            (fun lam -> lam)))))
+            ~ifnot:arg2))))
   | Psequand, [arg1; arg2] ->
     let const_false = Ident.create_local "const_false" in
     let cond = Ident.create_local "cond_sequand" in
@@ -152,8 +159,7 @@ let transform_primitive (prim : L.primitive) args loc =
           switch_for_if_then_else
             ~cond:(L.Lvar cond)
             ~ifso:arg2
-            ~ifnot:(L.Lvar const_false)
-            (fun lam -> lam)))))
+            ~ifnot:(L.Lvar const_false)))))
   | (Psequand | Psequor), _ ->
     Misc.fatal_error "Psequand / Psequor must have exactly two arguments"
   (* Removed. Should be safe, but will no longer catch misuses.
@@ -165,8 +171,7 @@ let transform_primitive (prim : L.primitive) args loc =
       (switch_for_if_then_else
         ~cond:(L.Lprim (Pflambda_isint, [arg], loc))
         ~ifso:(L.Lconst (Const_base (Const_int 1)))
-        ~ifnot:(L.Lconst (Const_base (Const_int 0)))
-        (fun lam -> lam))
+        ~ifnot:(L.Lconst (Const_base (Const_int 0))))
   | (Pidentity | Pbytes_to_string | Pbytes_of_string), [arg] -> Transformed arg
   | Pignore, [arg] ->
     let ident = Ident.create_local "ignore" in
@@ -247,6 +252,66 @@ let transform_primitive (prim : L.primitive) args loc =
     end
   | _, _ -> Primitive (prim, args, loc)
 
+let rec_catch_for_while_loop cond body =
+  let cont = L.next_raise_count () in
+  mark_as_recursive_static_catch cont;
+  let cond_result = Ident.create_local "while_cond_result" in
+  let lam : L.lambda =
+    Lstaticcatch (
+      Lstaticraise (cont, []),
+      (cont, []),
+      Llet (Strict, Pgenval, cond_result, cond,
+        Lifthenelse (Lvar cond_result,
+          Lsequence (
+            body,
+            Lstaticraise (cont, [])),
+          Lconst (Const_base (Const_int 0)))))
+  in lam
+
+let rec_catch_for_for_loop
+      ident start stop (dir : Asttypes.direction_flag) body =
+  let cont = L.next_raise_count () in
+  mark_as_recursive_static_catch cont;
+  let start_ident = Ident.create_local "for_start" in
+  let stop_ident = Ident.create_local "for_stop" in
+  let first_test : L.lambda =
+    match dir with
+    | Upto ->
+       Lprim (Pintcomp Cle,
+         [L.Lvar start_ident; L.Lvar stop_ident],
+         Loc_unknown)
+    | Downto ->
+       Lprim (Pintcomp Cge,
+         [L.Lvar start_ident; L.Lvar stop_ident],
+         Loc_unknown)
+  in
+  let subsequent_test : L.lambda =
+    Lprim (Pintcomp Cne, [L.Lvar ident; L.Lvar stop_ident], Loc_unknown)
+  in
+  let one : L.lambda = Lconst (Const_base (Const_int 1)) in
+  let next_value_of_counter =
+    match dir with
+    | Upto -> L.Lprim (Paddint, [L.Lvar ident; one], Loc_unknown)
+    | Downto -> L.Lprim (Psubint, [L.Lvar ident; one], Loc_unknown)
+  in
+  let lam : L.lambda =
+    (* Care needs to be taken here not to cause overflow if, for an
+       incrementing for-loop, the upper bound is [max_int]; likewise, for
+       a decrementing for-loop, if the lower bound is [min_int]. *)
+    Llet (Strict, Pgenval, start_ident, start,
+      Llet (Strict, Pgenval, stop_ident, stop,
+        Lifthenelse (first_test,
+          Lstaticcatch (
+            Lstaticraise (cont, [L.Lvar start_ident]),
+            (cont, [ident, Pgenval]),
+            Lsequence (
+              body,
+              Lifthenelse (subsequent_test,
+                Lstaticraise (cont, [next_value_of_counter]),
+                L.lambda_unit))),
+          L.lambda_unit)))
+  in lam
+
 let rec cps_non_tail (lam : L.lambda) (k : Ident.t -> Ilambda.t)
           (k_exn : Continuation.t) : Ilambda.t =
   match lam with
@@ -536,13 +601,25 @@ let rec cps_non_tail (lam : L.lambda) (k : Ident.t -> Ilambda.t)
         };
       handler = k result_var;
     }
+  | Lifthenelse (cond, ifso, ifnot) ->
+    let lam = switch_for_if_then_else ~cond ~ifso ~ifnot in
+    cps_non_tail lam k k_exn
+  | Lsequence (lam1, lam2) ->
+    let ident = Ident.create_local "sequence" in
+    cps_non_tail (L.Llet (Strict, Pgenval, ident, lam1, lam2)) k k_exn
+  | Lwhile (cond, body) ->
+    let loop = rec_catch_for_while_loop cond body in
+    cps_non_tail loop k k_exn
+  | Lfor (ident, start, stop, dir, body) ->
+    let loop = rec_catch_for_for_loop ident start stop dir body in
+    cps_non_tail loop k k_exn
   | Lassign (being_assigned, new_value) ->
     cps_non_tail_simple new_value (fun new_value ->
         name_then_cps_non_tail "assign"
           (I.Assign { being_assigned; new_value; })
           k k_exn)
       k_exn
-  | Lsequence _ | Lifthenelse _ | Lwhile _ | Lfor _ | Lifused _ | Levent _ ->
+  | Lifused _ | Levent _ ->
     Misc.fatal_errorf "Term should have been eliminated by [Prepare_lambda]: %a"
       Printlambda.lambda lam
 
@@ -818,7 +895,19 @@ and cps_tail (lam : L.lambda) (k : Continuation.t) (k_exn : Continuation.t)
         };
       handler;
     }
-  | Lsequence _ | Lifthenelse _ | Lwhile _ | Lfor _ | Lifused _ | Levent _ ->
+  | Lifthenelse (cond, ifso, ifnot) ->
+    let lam = switch_for_if_then_else ~cond ~ifso ~ifnot in
+    cps_tail lam k k_exn
+  | Lsequence (lam1, lam2) ->
+    let ident = Ident.create_local "sequence" in
+    cps_tail (L.Llet (Strict, Pgenval, ident, lam1, lam2)) k k_exn
+  | Lwhile (cond, body) ->
+    let loop = rec_catch_for_while_loop cond body in
+    cps_tail loop k k_exn
+  | Lfor (ident, start, stop, dir, body) ->
+    let loop = rec_catch_for_for_loop ident start stop dir body in
+    cps_tail loop k k_exn
+  | Lifused _ | Levent _ ->
     Misc.fatal_errorf "Term should have been eliminated by [Prepare_lambda]: %a"
       Printlambda.lambda lam
 
@@ -943,12 +1032,11 @@ and cps_switch (switch : proto_switch) ~scrutinee (k : Continuation.t)
         wrappers)
     k_exn
 
-let lambda_to_ilambda lam ~recursive_static_catches:recursive_static_catches'
-      : Ilambda.program =
+let lambda_to_ilambda lam : Ilambda.program =
   static_exn_env := Numbers.Int.Map.empty;
   try_stack := [];
   try_stack_at_handler := Continuation.Map.empty;
-  recursive_static_catches := recursive_static_catches';
+  recursive_static_catches := Numbers.Int.Set.empty;
   mutable_variables := Ident.Set.empty;
   let the_end = Continuation.create ~sort:Define_root_symbol () in
   let the_end_exn = Continuation.create ~sort:Exn () in

middle_end/flambda/from_lambda/cps_conversion.mli

@@ -18,7 +18,4 @@
 
 [@@@ocaml.warning "+a-4-30-40-41-42"]
 
-val lambda_to_ilambda
-   : Lambda.lambda
-  -> recursive_static_catches:Numbers.Int.Set.t
-  -> Ilambda.program
+val lambda_to_ilambda : Lambda.lambda -> Ilambda.program

middle_end/flambda/from_lambda/prepare_lambda.ml

@@ -62,31 +62,6 @@ end = struct
     }
 end
 
-(* CR-soon mshinwell: Remove mutable state *)
-let recursive_static_catches = ref Numbers.Int.Set.empty
-
-let mark_as_recursive_static_catch cont =
-  if Numbers.Int.Set.mem cont !recursive_static_catches then begin
-    Misc.fatal_errorf "Static catch with continuation %d already marked as \
-        recursive -- is it being redefined?"
-      cont
-  end;
-  recursive_static_catches := Numbers.Int.Set.add cont !recursive_static_catches
-
-let switch_for_if_then_else ~cond ~ifso ~ifnot k =
-  (* CR mshinwell: We need to make sure that [cond] is {0, 1}-valued.
-     The frontend should have been fixed on this branch for this. *)
-  let switch : Lambda.lambda_switch =
-    { sw_numconsts = 2;
-      sw_consts = [0, ifnot; 1, ifso];
-      sw_numblocks = 0;
-      sw_blocks = [];
-      sw_failaction = None;
-      sw_tags_to_sizes = Tag.Scannable.Map.empty;
-    }
-  in
-  k (L.Lswitch (cond, switch, Loc_unknown))
-
 (*
 let simplify_primitive (prim : L.primitive) args loc =
   match prim, args with
@@ -349,65 +324,20 @@ let rec prepare env (lam : L.lambda) (k : L.lambda -> L.lambda) =
     prepare env cond (fun cond ->
       prepare env ifso (fun ifso ->
         prepare env ifnot (fun ifnot ->
-          switch_for_if_then_else ~cond ~ifso ~ifnot k)))
+          k (L.Lifthenelse(cond, ifso, ifnot)))))
   | Lsequence (lam1, lam2) ->
-    let ident = Ident.create_local "sequence" in
-    prepare env (L.Llet (Strict, Pgenval, ident, lam1, lam2)) k
+    prepare env lam1 (fun lam1 ->
+      prepare env lam2 (fun lam2 ->
+        k (L.Lsequence(lam1, lam2))))
   | Lwhile (cond, body) ->
-    let cont = L.next_raise_count () in
-    mark_as_recursive_static_catch cont;
-    let cond_result = Ident.create_local "cond_result" in
-    let lam : L.lambda =
-      Lstaticcatch (
-        Lstaticraise (cont, []),
-        (cont, []),
-        Llet (Strict, Pgenval, cond_result, cond,
-          Lifthenelse (Lvar cond_result,
-            Lsequence (
-              body,
-              Lstaticraise (cont, [])),
-            Lconst (Const_base (Const_int 0)))))
-    in
-    prepare env lam k
+    prepare env cond (fun cond ->
+      prepare env body (fun body ->
+        k (Lwhile (cond, body))))
   | Lfor (ident, start, stop, dir, body) ->
-    let cont = L.next_raise_count () in
-    mark_as_recursive_static_catch cont;
-    let start_ident = Ident.create_local "start" in
-    let stop_ident = Ident.create_local "stop" in
-    let first_test : L.lambda =
-      match dir with
-      | Upto ->
-        Lprim (Pintcomp Cle, [L.Lvar start_ident; L.Lvar stop_ident], Loc_unknown)
-      | Downto ->
-        Lprim (Pintcomp Cge, [L.Lvar start_ident; L.Lvar stop_ident], Loc_unknown)
-    in
-    let subsequent_test : L.lambda =
-      Lprim (Pintcomp Cne, [L.Lvar ident; L.Lvar stop_ident], Loc_unknown)
-    in
-    let one : L.lambda = Lconst (Const_base (Const_int 1)) in
-    let next_value_of_counter =
-      match dir with
-      | Upto -> L.Lprim (Paddint, [L.Lvar ident; one], Loc_unknown)
-      | Downto -> L.Lprim (Psubint, [L.Lvar ident; one], Loc_unknown)
-    in
-    let lam : L.lambda =
-      (* Care needs to be taken here not to cause overflow if, for an
-         incrementing for-loop, the upper bound is [max_int]; likewise, for
-         a decrementing for-loop, if the lower bound is [min_int]. *)
-      Llet (Strict, Pgenval, start_ident, start,
-        Llet (Strict, Pgenval, stop_ident, stop,
-          Lifthenelse (first_test,
-            Lstaticcatch (
-              Lstaticraise (cont, [start]),
-              (cont, [ident, Pgenval]),
-              Lsequence (
-                body,
-                Lifthenelse (subsequent_test,
-                  Lstaticraise (cont, [next_value_of_counter]),
-                  L.lambda_unit))),
-            L.lambda_unit)))
-    in
-    prepare env lam k
+    prepare env start (fun start ->
+      prepare env stop (fun stop ->
+        prepare env body (fun body ->
+          k (L.Lfor (ident, start, stop, dir, body)))))
   | Lassign (ident, lam) ->
     if not (Env.is_mutable env ident) then begin
       Misc.fatal_errorf "Lassign on non-mutable variable %a"
@@ -456,11 +386,9 @@ and prepare_option env lam_opt k =
   | Some lam -> prepare env lam (fun lam -> k (Some lam))
 
 let run lam =
-  recursive_static_catches := Numbers.Int.Set.empty;
   let current_unit_id =
     Compilation_unit.get_persistent_ident
       (Compilation_unit.get_current_exn ())
   in
   let env = Env.create ~current_unit_id in
-  let lam = prepare env lam (fun lam -> lam) in
-  lam, !recursive_static_catches
+  prepare env lam (fun lam -> lam)

middle_end/flambda/from_lambda/prepare_lambda.mli

@@ -16,10 +16,4 @@
 
 (** Preparation of [Lambda] code before CPS and closure conversion. *)
 
-(** The set of integers returned by [run] identifies all those [Lstaticcatch]
-    handlers which are to be treated as recursive.  (This is rather more
-    straightforward than changing the type in [Lambda] to accommodate
-    this). *)
-val run
-   : Lambda.lambda
-  -> Lambda.lambda * Numbers.Int.Set.t
+val run : Lambda.lambda -> Lambda.lambda