Avoid additional block splitting during CPS transformation

We ignore any event at the end of a block. This happens when the block
ends with a return or a branch. In case of a return, it's a tail call,
so we don't have an event anyway. For a branch, the target block will
start with an event which is would take precedence anyway.

Author: vouillon

compiler/lib/effects.ml

@@ -139,6 +139,28 @@ let dominance_frontier g idom =
     g.preds;
   frontiers
 
+(* Last instruction of a block, ignoring events *)
+let rec last_instr l =
+  match l with
+  | [] -> None
+  | [ i ] | [ i; (Event _, _) ] -> Some i
+  | _ :: rem -> last_instr rem
+
+(* Split a block, separating the last instruction from the preceeding
+   ones, ignoring events *)
+let block_split_last xs =
+  let rec aux acc = function
+    | [] -> None
+    | [ x ] | [ x; (Event _, _) ] -> Some (List.rev acc, x)
+    | x :: xs -> aux (x :: acc) xs
+  in
+  aux [] xs
+
+let empty_body b =
+  match b with
+  | [] | [ (Event _, _) ] -> true
+  | _ -> false
+
 (****)
 
 (*
@@ -176,7 +198,7 @@ let compute_needed_transformations ~cfg ~idom ~cps_needed ~blocks ~start =
       let block = Addr.Map.find pc blocks in
       (match fst block.branch with
       | Branch (dst, _) -> (
-          match List.last block.body with
+          match last_instr block.body with
           | Some
               ( Let
                   (x, (Apply _ | Prim (Extern ("%resume" | "%perform" | "%reperform"), _)))
@@ -572,7 +594,7 @@ let cps_block ~st ~k pc block =
   in
 
   let rewritten_block =
-    match List.split_last block.body, block.branch with
+    match block_split_last block.body, block.branch with
     | Some (body_prefix, (Let (x, e), loc)), (Return ret, _loc_ret) ->
         Option.map (rewrite_instr x e loc) ~f:(fun f ->
             assert (List.is_empty alloc_jump_closures);
@@ -847,7 +869,7 @@ let split_blocks ~cps_needed (p : Code.program) =
     let is_split_point i r branch =
       match i with
       | Let (x, (Apply _ | Prim (Extern ("%resume" | "%perform" | "%reperform"), _))) ->
-          ((not (List.is_empty r))
+          ((not (empty_body r))
           ||
           match fst branch with
           | Branch _ -> false
@@ -901,14 +923,6 @@ let remove_empty_blocks ~live_vars (p : Code.program) : Code.program =
       | None -> cont
   in
   let resolve cont = resolve_rec Addr.Set.empty cont in
-  let empty_body b =
-    List.for_all
-      ~f:(fun (i, _) ->
-        match i with
-        | Event _ -> true
-        | _ -> false)
-      b
-  in
   Addr.Map.iter
     (fun pc block ->
       match block with

compiler/lib/partial_cps_analysis.ml

@@ -39,9 +39,18 @@ let add_tail_dep deps x y =
       (fun s -> Some (Var.Set.add x (Option.value ~default:Var.Set.empty s)))
       !deps
 
+let rec block_iter_last ~f l =
+  match l with
+  | [] -> ()
+  | [ i ] -> f true i
+  | [ i; (Event _, _) ] -> f true i
+  | i :: l ->
+      f false i;
+      block_iter_last ~f l
+
 let block_deps ~info ~vars ~tail_deps ~deps ~blocks ~fun_name pc =
   let block = Addr.Map.find pc blocks in
-  List.iter_last block.body ~f:(fun is_last (i, _) ->
+  block_iter_last block.body ~f:(fun is_last (i, _) ->
       match i with
       | Let (x, Apply { f; _ }) -> (
           add_var vars x;

compiler/tests-compiler/effects_continuations.ml

@@ -103,86 +103,79 @@ let%expect_test "test-compiler/lib-effects/test1.ml" =
     {|
 
     function exceptions(s, cont){
-     try{var _w_ = runtime.caml_int_of_string(s), n = _w_;}
-     catch(_A_){
-      var _s_ = caml_wrap_exception(_A_);
-      if(_s_[1] !== Stdlib[7]){
+     try{var _t_ = runtime.caml_int_of_string(s), n = _t_;}
+     catch(_x_){
+      var _p_ = caml_wrap_exception(_x_);
+      if(_p_[1] !== Stdlib[7]){
        var raise$1 = caml_pop_trap();
-       return raise$1(caml_maybe_attach_backtrace(_s_, 0));
+       return raise$1(caml_maybe_attach_backtrace(_p_, 0));
       }
       var n = 0;
      }
      try{
       if(caml_string_equal(s, cst$0))
        throw caml_maybe_attach_backtrace(Stdlib[8], 1);
-      var _v_ = 7, m = _v_;
+      var _s_ = 7, m = _s_;
      }
-     catch(_z_){
-      var _t_ = caml_wrap_exception(_z_);
-      if(_t_ !== Stdlib[8]){
+     catch(_w_){
+      var _q_ = caml_wrap_exception(_w_);
+      if(_q_ !== Stdlib[8]){
        var raise$0 = caml_pop_trap();
-       return raise$0(caml_maybe_attach_backtrace(_t_, 0));
+       return raise$0(caml_maybe_attach_backtrace(_q_, 0));
       }
       var m = 0;
      }
      runtime.caml_push_trap
-      (function(_y_){
-        if(_y_ === Stdlib[8]) return cont(0);
+      (function(_v_){
+        if(_v_ === Stdlib[8]) return cont(0);
         var raise = caml_pop_trap();
-        return raise(caml_maybe_attach_backtrace(_y_, 0));
+        return raise(caml_maybe_attach_backtrace(_v_, 0));
        });
      if(! caml_string_equal(s, cst))
       return caml_cps_call2
               (Stdlib[79],
                cst_toto,
-               function(_x_){caml_pop_trap(); return cont([0, [0, _x_, n, m]]);});
-     var _u_ = Stdlib[8], raise = caml_pop_trap();
-     return raise(caml_maybe_attach_backtrace(_u_, 1));
+               function(_u_){caml_pop_trap(); return cont([0, [0, _u_, n, m]]);});
+     var _r_ = Stdlib[8], raise = caml_pop_trap();
+     return raise(caml_maybe_attach_backtrace(_r_, 1));
     }
     //end
     function cond1(b, cont){
-     function _p_(ic){return cont([0, ic, 7]);}
+     function _o_(ic){return cont([0, ic, 7]);}
      return b
-             ? caml_cps_call2
-               (Stdlib[79], cst_toto$0, function(_q_){return _p_(_q_);})
-             : caml_cps_call2
-               (Stdlib[79], cst_titi, function(_r_){return _p_(_r_);});
+             ? caml_cps_call2(Stdlib[79], cst_toto$0, _o_)
+             : caml_cps_call2(Stdlib[79], cst_titi, _o_);
     }
     //end
     function cond2(b, cont){
-     function _m_(){return cont(7);}
+     function _m_(_n_){return cont(7);}
      return b
-             ? caml_cps_call2(Stdlib_Printf[3], _a_, function(_n_){return _m_();})
-             : caml_cps_call2(Stdlib_Printf[3], _b_, function(_o_){return _m_();});
+             ? caml_cps_call2(Stdlib_Printf[3], _a_, _m_)
+             : caml_cps_call2(Stdlib_Printf[3], _b_, _m_);
     }
     //end
     function cond3(b, cont){
      var x = [0, 0];
-     function _k_(){return cont(x[1]);}
-     return b
-             ? (x[1] = 1, _k_())
-             : caml_cps_call2(Stdlib_Printf[3], _c_, function(_l_){return _k_();});
+     function _k_(_l_){return cont(x[1]);}
+     return b ? (x[1] = 1, _k_(0)) : caml_cps_call2(Stdlib_Printf[3], _c_, _k_);
     }
     //end
     function loop1(b, cont){
      return caml_cps_call2
              (Stdlib[79],
               cst_static_examples_ml,
               function(ic){
-               function _i_(){
+               function _i_(_j_){
                 return caml_cps_call2
                         (Stdlib[83],
                          ic,
                          function(line){
                           return b
-                                  ? caml_cps_call2
-                                    (Stdlib[53],
-                                     line,
-                                     function(_j_){return caml_cps_exact_call0(_i_);})
-                                  : caml_cps_exact_call0(_i_);
+                                  ? caml_cps_call2(Stdlib[53], line, _i_)
+                                  : caml_cps_exact_call1(_i_, 0);
                          });
                }
-               return _i_();
+               return _i_(0);
               });
     }
     //end
@@ -191,23 +184,15 @@ let%expect_test "test-compiler/lib-effects/test1.ml" =
              (Stdlib[79],
               cst_static_examples_ml$0,
               function(ic){
-               return caml_cps_call2
-                       (Stdlib_Printf[3],
-                        _d_,
-                        function(_f_){
-                         function _g_(){
-                          return caml_cps_call2
-                                  (Stdlib[83],
-                                   ic,
-                                   function(line){
-                                    return caml_cps_call2
-                                            (Stdlib[53],
-                                             line,
-                                             function(_h_){return caml_cps_exact_call0(_g_);});
-                                   });
-                         }
-                         return _g_();
-                        });
+               function _g_(_h_){
+                return caml_cps_call2
+                        (Stdlib[83],
+                         ic,
+                         function(line){
+                          return caml_cps_call2(Stdlib[53], line, _g_);
+                         });
+               }
+               return caml_cps_call2(Stdlib_Printf[3], _d_, _g_);
               });
     }
     //end
@@ -216,8 +201,12 @@ let%expect_test "test-compiler/lib-effects/test1.ml" =
              (list_rev,
               _e_,
               function(l){
-               var x = l;
-               for(;;){if(! x) return cont(l); var r = x[2]; x = r;}
+               function _f_(x){
+                if(! x) return cont(l);
+                var r = x[2];
+                return caml_cps_exact_call1(_f_, r);
+               }
+               return _f_(l);
               });
     }
     //end |}]