From f9c7ca20eace35fd5c332970a71654f9fc96eaa3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Vouillon?= <jerome.vouillon@gmail.com>
Date: Mon, 9 Dec 2024 10:58:43 +0100
Subject: [PATCH] Simplify function rewrite_instr

---
 compiler/lib/effects.ml | 29 +++++++++++------------------
 1 file changed, 11 insertions(+), 18 deletions(-)

diff --git a/compiler/lib/effects.ml b/compiler/lib/effects.ml
index 86a24105ba..f7bd176c4f 100644
--- a/compiler/lib/effects.ml
+++ b/compiler/lib/effects.ml
@@ -557,7 +557,7 @@ let cps_last ~st ~alloc_jump_closures pc (last : last) ~k : instr list * last =
             @ (Let (exn_handler, Prim (Extern "caml_pop_trap", [])) :: body)
           , branch ))
 
-let rewrite_instr ~st (instr : instr) : instr list =
+let rewrite_instr ~st (instr : instr) : instr =
   match instr with
   | Let (x, Closure (_, (pc, _))) when Var.Set.mem x st.cps_needed ->
       (* When CPS-transforming with double translation enabled, there are no closures in
@@ -567,16 +567,15 @@ let rewrite_instr ~st (instr : instr) : instr list =
          needed *)
       let cps_params, cps_cont = Hashtbl.find st.closure_info pc in
       st.in_cps := Var.Set.add x !(st.in_cps);
-      [ Let (x, Closure (cps_params, cps_cont)) ]
+      Let (x, Closure (cps_params, cps_cont))
   | Let (x, Prim (Extern "caml_alloc_dummy_function", [ size; arity ])) -> (
       match arity with
       | Pc (Int a) ->
-          [ Let
-              ( x
-              , Prim
-                  ( Extern "caml_alloc_dummy_function"
-                  , [ size; Pc (Int (Targetint.succ a)) ] ) )
-          ]
+          Let
+            ( x
+            , Prim
+                (Extern "caml_alloc_dummy_function", [ size; Pc (Int (Targetint.succ a)) ])
+            )
       | _ -> assert false)
   | Let (x, Apply { f; args; _ }) when not (Var.Set.mem x st.cps_needed) ->
       (* At the moment, we turn into CPS any function not called with
@@ -586,13 +585,13 @@ let rewrite_instr ~st (instr : instr) : instr list =
            introduced by the lambda lifting and we don't have exactness info any more. *)
         Var.idx f >= Var.Tbl.length st.flow_info.info_approximation
         || Global_flow.exact_call st.flow_info f (List.length args));
-      [ Let (x, Apply { f; args; exact = true }) ]
+      Let (x, Apply { f; args; exact = true })
   | Let (_, e) when effect_primitive_or_application e ->
       (* For the CPS target, applications of CPS functions and effect primitives require
          more work (allocating a continuation and/or modifying end-of-block branches) and
          are handled in a specialized function. *)
       assert false
-  | _ -> [ instr ]
+  | _ -> instr
 
 let cps_instr ~st (instr : instr) : instr list =
   match instr with
@@ -606,13 +605,7 @@ let cps_instr ~st (instr : instr) : instr list =
       [ Let (unit, Constant (Int Targetint.zero))
       ; Let (x, Apply { exact; f; args = [ unit ] })
       ]
-  | Let (_, Prim (Extern "caml_assume_no_perform", args)) ->
-      invalid_arg
-      @@ Format.sprintf
-           "Internal primitive `caml_assume_no_perform` takes exactly 1 argument (%d \
-            given)"
-           (List.length args)
-  | _ -> rewrite_instr ~st instr
+  | _ -> [ rewrite_instr ~st instr ]
 
 let cps_block ~st ~k ~orig_pc block =
   debug_print "cps_block %d\n" orig_pc;
@@ -770,7 +763,7 @@ let rewrite_direct_block ~st ~cps_needed ~closure_info ~pc block =
     in
     let body = List.concat_map block.body ~f:(fun i -> rewrite_instr i) in
     { block with body }
-  else { block with body = List.concat_map ~f:(rewrite_instr ~st) block.body }
+  else { block with body = List.map ~f:(rewrite_instr ~st) block.body }
 
 (* Apply a substitution in a set of blocks, including to bound variables *)
 let subst_bound_in_blocks blocks s =