Compiler: make tailcall optim more robust

We introduce an intermediate block at the beginning of the closure (if necessary) and
jump back to it when we see a tailcall.

With this, we can:
- remove the "skip_param" argument from "Flow.f"
- simplify the O3 profile

Author: hhugo

compiler/lib/driver.ml

@@ -87,10 +87,6 @@ let flow p =
   if debug () then Format.eprintf "Data flow...@.";
   Flow.f p
 
-let flow_simple p =
-  if debug () then Format.eprintf "Data flow...@.";
-  Flow.f ~skip_param:true p
-
 let phi p =
   if debug () then Format.eprintf "Variable passing simplification...@.";
   Phisimpl.f p
@@ -161,7 +157,7 @@ let identity x = x
 let o1 : 'a -> 'a =
   print
   +> tailcall
-  +> flow_simple (* flow simple to keep information for future tailcall opt *)
+  +> flow
   +> specialize'
   +> eval
   +> inline (* inlining may reveal new tailcall opt *)
@@ -190,19 +186,7 @@ let o2 : 'a -> 'a = loop 10 "o1" o1 1 +> print
 
 (* o3 *)
 
-let round1 : 'a -> 'a =
-  print
-  +> tailcall
-  +> inline (* inlining may reveal new tailcall opt *)
-  +> deadcode (* deadcode required before flow simple -> provided by constant *)
-  +> flow_simple (* flow simple to keep information for future tailcall opt *)
-  +> specialize'
-  +> eval
-  +> identity
-
-let round2 = flow +> specialize' +> eval +> deadcode +> o1
-
-let o3 = loop 10 "tailcall+inline" round1 1 +> loop 10 "flow" round2 1 +> print
+let o3 = loop 10 "o1" o1 1 +> print
 
 let generate
     d

compiler/lib/flow.ml

@@ -289,9 +289,9 @@ let program_escape defs known_origins { blocks; _ } =
 
 (****)
 
-let propagate2 ?(skip_param = false) defs known_origins possibly_mutable st x =
+let propagate2 defs known_origins possibly_mutable st x =
   match defs.(Var.idx x) with
-  | Param -> skip_param
+  | Param -> false
   | Phi s -> Var.Set.exists (fun y -> Var.Tbl.get st y) s
   | Expr e -> (
       match e with
@@ -318,11 +318,11 @@ end
 
 module Solver2 = G.Solver (Domain2)
 
-let solver2 ?skip_param vars deps defs known_origins possibly_mutable =
+let solver2 vars deps defs known_origins possibly_mutable =
   let g =
     { G.domain = vars; G.iter_children = (fun f x -> Var.Set.iter f deps.(Var.idx x)) }
   in
-  Solver2.f () g (propagate2 ?skip_param defs known_origins possibly_mutable)
+  Solver2.f () g (propagate2 defs known_origins possibly_mutable)
 
 let get_approx
     { Info.info_defs = _; info_known_origins; info_maybe_unknown; _ }
@@ -492,7 +492,7 @@ let build_subst (info : Info.t) vars =
 
 (****)
 
-let f ?skip_param p =
+let f p =
   Code.invariant p;
   let t = Timer.make () in
   let t1 = Timer.make () in
@@ -505,7 +505,7 @@ let f ?skip_param p =
   let possibly_mutable = program_escape defs known_origins p in
   if times () then Format.eprintf "    flow analysis 3: %a@." Timer.print t3;
   let t4 = Timer.make () in
-  let maybe_unknown = solver2 ?skip_param vars deps defs known_origins possibly_mutable in
+  let maybe_unknown = solver2 vars deps defs known_origins possibly_mutable in
   if times () then Format.eprintf "    flow analysis 4: %a@." Timer.print t4;
   if debug ()
   then

compiler/lib/flow.mli

@@ -65,4 +65,4 @@ val the_block_contents_of : Info.t -> Code.prim_arg -> Code.Var.t array option
 
 val the_int : Info.t -> Code.prim_arg -> Targetint.t option
 
-val f : ?skip_param:bool -> Code.program -> Code.program * Info.t
+val f : Code.program -> Code.program * Info.t

compiler/lib/tailcall.ml

@@ -39,53 +39,101 @@ let rec tail_call x f l =
     -> Some args
   | _ :: rem -> tail_call x f rem
 
-let rewrite_block (f, f_params, f_pc, args) pc blocks =
+let rewrite_block (f, f_params, f_pc, used, stats) pc blocks =
   let block = Addr.Map.find pc blocks in
   match block.branch with
   | Return x -> (
       match tail_call x f block.body with
-      | Some f_args when List.length f_params = List.length f_args ->
-          let m = Subst.build_mapping f_params f_args in
-          List.iter2 f_params f_args ~f:(fun p a -> Code.Var.propagate_name p a);
-          Addr.Map.add
-            pc
-            { params = block.params
-            ; body = remove_last block.body
-            ; branch = Branch (f_pc, List.map args ~f:(fun x -> Var.Map.find x m))
-            }
-            blocks
-      | _ -> blocks)
+      | Some f_args ->
+          if List.compare_lengths f_params f_args = 0
+          then (
+            incr stats;
+            List.iter2 f_params f_args ~f:(fun p a -> Code.Var.propagate_name p a);
+            used := true;
+            Addr.Map.add
+              pc
+              { params = block.params
+              ; body = remove_last block.body
+              ; branch = Branch (f_pc, f_args)
+              }
+              blocks)
+          else blocks
+      | None -> blocks)
   | _ -> blocks
 
 let rec traverse f pc visited blocks =
   if not (Addr.Set.mem pc visited)
   then
     let visited = Addr.Set.add pc visited in
-    let blocks = rewrite_block f pc blocks in
-    let visited, blocks =
-      Code.fold_children_skip_try_body
-        blocks
-        pc
-        (fun pc (visited, blocks) ->
-          let visited, blocks = traverse f pc visited blocks in
-          visited, blocks)
-        (visited, blocks)
-    in
-    visited, blocks
+    let blocks' = rewrite_block f pc blocks in
+    if not (phys_equal blocks blocks')
+    then (* the block was rewritten *)
+      visited, blocks'
+    else
+      let blocks = blocks' in
+      let visited, blocks =
+        Code.fold_children_skip_try_body
+          blocks
+          pc
+          (fun pc (visited, blocks) ->
+            let visited, blocks = traverse f pc visited blocks in
+            visited, blocks)
+          (visited, blocks)
+      in
+      visited, blocks
   else visited, blocks
 
 let f p =
+  let free_pc = ref p.free_pc in
+  let blocks = ref p.blocks in
+  let stats = ref 0 in
   let t = Timer.make () in
-  let blocks =
-    fold_closures
-      p
-      (fun f params (pc, args) blocks ->
-        match f with
-        | Some f when List.length params = List.length args ->
-            let _, blocks = traverse (f, params, pc, args) pc Addr.Set.empty blocks in
-            blocks
-        | _ -> blocks)
-      p.blocks
-  in
-  if times () then Format.eprintf "  tail calls: %a@." Timer.print t;
-  { p with blocks }
+  Addr.Map.iter
+    (fun pc _ ->
+      let block = Addr.Map.find pc !blocks in
+      let rewrite_body = ref false in
+      let body =
+        List.map block.body ~f:(function
+          | Let (f, Closure (params, (pc_head, args))) as i ->
+              if List.equal ~eq:Code.Var.equal params args
+              then (
+                blocks :=
+                  snd
+                    (traverse
+                       (f, params, pc_head, ref false, stats)
+                       pc_head
+                       Addr.Set.empty
+                       !blocks);
+                i)
+              else
+                let intermediate_pc = !free_pc in
+                let need_to_create_intermediate_block = ref false in
+                blocks :=
+                  snd
+                    (traverse
+                       ( f
+                       , params
+                       , intermediate_pc
+                       , need_to_create_intermediate_block
+                       , stats )
+                       pc_head
+                       Addr.Set.empty
+                       !blocks);
+                if !need_to_create_intermediate_block
+                then (
+                  incr free_pc;
+                  let new_params = List.map params ~f:Code.Var.fork in
+                  blocks :=
+                    Addr.Map.add
+                      intermediate_pc
+                      { params; body = []; branch = Branch (pc_head, args) }
+                      !blocks;
+                  rewrite_body := true;
+                  Let (f, Closure (new_params, (intermediate_pc, new_params))))
+                else i
+          | i -> i)
+      in
+      if !rewrite_body then blocks := Addr.Map.add pc { block with body } !blocks)
+    p.blocks;
+  if times () then Format.eprintf "  tail calls: %a #%d@." Timer.print t !stats;
+  { p with blocks = !blocks; free_pc = !free_pc }

compiler/tests-compiler/double-translation/effects_continuations.ml

@@ -101,42 +101,44 @@ let%expect_test "test-compiler/lib-effects/test1.ml" =
   [%expect
     {|
     function exceptions$0(s){
-     try{var _K_ = caml_int_of_string(s), n = _K_;}
-     catch(_N_){
-      var _F_ = caml_wrap_exception(_N_);
-      if(_F_[1] !== Stdlib[7]) throw caml_maybe_attach_backtrace(_F_, 0);
+     try{var _C_ = caml_int_of_string(s), n = _C_;}
+     catch(exn$0){
+      var exn = caml_wrap_exception(exn$0);
+      if(exn[1] !== Stdlib[7]) throw caml_maybe_attach_backtrace(exn, 0);
       var n = 0;
      }
      try{
       if(caml_string_equal(s, cst$0))
        throw caml_maybe_attach_backtrace(Stdlib[8], 1);
-      var _J_ = 7, m = _J_;
+      var _B_ = 7, m = _B_;
      }
-     catch(_M_){
-      var _G_ = caml_wrap_exception(_M_);
-      if(_G_ !== Stdlib[8]) throw caml_maybe_attach_backtrace(_G_, 0);
+     catch(exn){
+      var exn$0 = caml_wrap_exception(exn);
+      if(exn$0 !== Stdlib[8]) throw caml_maybe_attach_backtrace(exn$0, 0);
       var m = 0;
      }
      try{
       if(caml_string_equal(s, cst))
        throw caml_maybe_attach_backtrace(Stdlib[8], 1);
-      var _I_ = [0, [0, caml_call1(Stdlib[79], cst_toto), n, m]];
-      return _I_;
+      var _A_ = [0, [0, caml_call1(Stdlib[79], cst_toto), n, m]];
+      return _A_;
      }
-     catch(_L_){
-      var _H_ = caml_wrap_exception(_L_);
-      if(_H_ === Stdlib[8]) return 0;
-      throw caml_maybe_attach_backtrace(_H_, 0);
+     catch(exn){
+      var exn$1 = caml_wrap_exception(exn);
+      if(exn$1 === Stdlib[8]) return 0;
+      throw caml_maybe_attach_backtrace(exn$1, 0);
      }
     }
     //end
     function exceptions$1(s, cont){
-     try{var _z_ = caml_int_of_string(s), n = _z_;}
-     catch(_E_){
-      var _A_ = caml_wrap_exception(_E_);
-      if(_A_[1] !== Stdlib[7]){
-       var raise$1 = caml_pop_trap();
-       return raise$1(caml_maybe_attach_backtrace(_A_, 0));
+     try{var _y_ = caml_int_of_string(s), n = _y_;}
+     catch(exn){
+      var exn$2 = caml_wrap_exception(exn);
+      if(exn$2[1] !== Stdlib[7]){
+       var
+        raise$1 = caml_pop_trap(),
+        exn$0 = caml_maybe_attach_backtrace(exn$2, 0);
+       return raise$1(exn$0);
       }
       var n = 0;
      }
@@ -145,25 +147,25 @@ let%expect_test "test-compiler/lib-effects/test1.ml" =
        throw caml_maybe_attach_backtrace(Stdlib[8], 1);
       var _x_ = 7, m = _x_;
      }
-     catch(_D_){
-      var _y_ = caml_wrap_exception(_D_);
-      if(_y_ !== Stdlib[8]){
-       var raise$0 = caml_pop_trap();
-       return raise$0(caml_maybe_attach_backtrace(_y_, 0));
+     catch(exn$0){
+      var exn$1 = caml_wrap_exception(exn$0);
+      if(exn$1 !== Stdlib[8]){
+       var raise$0 = caml_pop_trap(), exn = caml_maybe_attach_backtrace(exn$1, 0);
+       return raise$0(exn);
       }
       var m = 0;
      }
      runtime.caml_push_trap
-      (function(_C_){
-        if(_C_ === Stdlib[8]) return cont(0);
-        var raise = caml_pop_trap();
-        return raise(caml_maybe_attach_backtrace(_C_, 0));
+      (function(exn$0){
+        if(exn$0 === Stdlib[8]) return cont(0);
+        var raise = caml_pop_trap(), exn = caml_maybe_attach_backtrace(exn$0, 0);
+        return raise(exn);
        });
      if(! caml_string_equal(s, cst))
       return caml_trampoline_cps_call2
               (Stdlib[79],
                cst_toto,
-               function(_B_){caml_pop_trap(); return cont([0, [0, _B_, n, m]]);});
+               function(_z_){caml_pop_trap(); return cont([0, [0, _z_, n, m]]);});
      var _w_ = Stdlib[8], raise = caml_pop_trap();
      return raise(caml_maybe_attach_backtrace(_w_, 1));
     }