Compiler: exit loop early

hhugo · hhugo · commit ac1985e671a5 · 2022-11-25T00:49:15.000+01:00
diff --git a/compiler/lib/generate.ml b/compiler/lib/generate.ml
@@ -20,9 +20,6 @@
 
 (*XXX
   Patterns:
-  => loops should avoid absorbing the whole continuation...
-     (detect when the continuation does not loop anymore and close
-      the loop at this point)
   => should have special code for switches that include the preceding
      if statement when possible
   => if e1 then {if e2 then P else Q} else {if e3 then P else Q}
@@ -741,6 +738,7 @@ let dominance_frontier st pc =
     | Old -> dominance_frontier_ref_implem st pc
     | New -> dominance_frontier_opt st pc
   in
+  if debug () then Format.eprintf "Dominance(%d): %s@." pc (string_of_set frontier);
   (if debug_dominance_frontier ()
   then
     let o, n =
@@ -1357,6 +1355,49 @@ and translate_instrs ctx expr_queue loc instr =
       let instrs, expr_queue = translate_instrs ctx expr_queue loc rem in
       st @ instrs, expr_queue
 
+and compute_exit_loop st pc_loop =
+  let cache = Hashtbl.create 17 in
+  let rec loop pc =
+    match Hashtbl.find cache pc with
+    | x -> x
+    | exception Not_found ->
+        let succs = Hashtbl.find st.succs pc in
+        let backs = Hashtbl.find st.backs pc in
+        let all_backs =
+          List.fold_left succs ~init:backs ~f:(fun acc pc -> Addr.Set.union acc (loop pc))
+        in
+        Hashtbl.replace cache pc all_backs;
+        all_backs
+  in
+  ignore (loop pc_loop);
+  let visited = Hashtbl.create 18 in
+  let rec find pc_loops exn_handler pc acc =
+    if Hashtbl.mem visited pc
+    then acc
+    else
+      let () = Hashtbl.add visited pc () in
+      let succs = Hashtbl.find st.succs pc in
+      let backs = Hashtbl.find cache pc in
+      if Addr.Set.cardinal (Addr.Set.inter pc_loops backs) > 0
+         || not (List.is_empty exn_handler)
+      then
+        let pc_loops =
+          if Addr.Set.mem pc st.loops then Addr.Set.add pc pc_loops else pc_loops
+        in
+        let block = Addr.Map.find pc st.blocks in
+        List.fold_left succs ~init:acc ~f:(fun acc pc' ->
+            let exn_handler =
+              match block.branch with
+              | Pushtrap (_, _, (pc_exn, _), _) when pc' <> pc_exn ->
+                  block.branch :: exn_handler
+              | Poptrap _ -> List.tl exn_handler
+              | _ -> exn_handler
+            in
+            find pc_loops exn_handler pc' acc)
+      else Addr.Set.add pc acc
+  in
+  find (Addr.Set.singleton pc_loop) [] pc_loop Addr.Set.empty
+
 and compile_block st queue (pc : Addr.t) frontier interm =
   if (not (List.is_empty queue))
      && (Addr.Set.mem pc st.loops || not (Config.Flag.inline ()))
@@ -1374,7 +1415,65 @@ and compile_block st queue (pc : Addr.t) frontier interm =
           | [] -> J.Label.zero
         in
         st.loop_stack <- (pc, (lab, ref false)) :: st.loop_stack;
-        let never_body, body = compile_block_no_loop st queue pc frontier interm in
+        let exit_grey = compute_exit_loop st pc in
+        let size_of the_pc =
+          let fold _blocs pc f acc =
+            let succs = Hashtbl.find st.succs pc in
+            List.fold_left succs ~init:acc ~f:(fun acc pc -> f pc acc)
+          in
+          Code.traverse
+            { fold }
+            (fun pc i ->
+              let b = Addr.Map.find pc st.blocks in
+              (* prefer to not nest for loops *)
+              let i = if Addr.Set.mem pc st.loops then i + 10000 else i in
+              let i =
+                if Addr.Set.is_empty (Hashtbl.find st.backs pc) then i else i + 10000
+              in
+              let i =
+                match b.branch with
+                | Cond _ | Switch _ -> i + 10
+                | Branch _ -> i
+                | Stop | Raise _ | Return _ | Poptrap _ -> i
+                | Pushtrap _ -> i + 10
+              in
+              i + List.length b.body)
+            the_pc
+            st.blocks
+            0
+        in
+        let list_max f = function
+          | [] -> None
+          | x :: xs ->
+              Some
+                (List.fold_left
+                   ~init:(f x, x)
+                   xs
+                   ~f:(fun (v, data) x ->
+                     let y = f x in
+                     if y > x then y, x else v, data))
+        in
+        let exit_grey =
+          Addr.Set.elements exit_grey
+          |> List.partition ~f:(fun pc -> Addr.Set.is_empty (dominance_frontier st pc))
+          |> function
+          | _, [ other ] -> Addr.Set.singleton other
+          | l, _ -> (
+              match list_max (fun x -> size_of x) l with
+              | None -> Addr.Set.empty
+              | Some (size, pc) ->
+                  if size < 50 then Addr.Set.empty else Addr.Set.singleton pc)
+        in
+        let exit_frontier = resolve_nodes interm exit_grey in
+        let exit_prefix, exit_cont, exit_interm =
+          colapse_frontier "for-loop" st exit_frontier interm
+        in
+        assert (Addr.Set.cardinal exit_cont <= 1);
+        let frontier_inner = Addr.Set.union frontier exit_cont in
+        Addr.Set.iter (protect_preds st) exit_frontier;
+        let never_body, body =
+          compile_block_no_loop st queue pc frontier_inner exit_interm
+        in
         let body =
           let rec remove_tailing_continue acc = function
             | [] -> body
@@ -1383,6 +1482,7 @@ and compile_block st queue (pc : Addr.t) frontier interm =
           in
           remove_tailing_continue [] body
         in
+        Addr.Set.iter (unprotect_preds st) exit_frontier;
         let for_loop =
           ( J.For_statement
               ( J.Left None
@@ -1405,9 +1505,20 @@ and compile_block st queue (pc : Addr.t) frontier interm =
               if !used then Some l else None
           | [] -> assert false
         in
-        match label with
-        | None -> never_body, [ for_loop ]
-        | Some label -> never_body, [ J.Labelled_statement (label, for_loop), J.N ])
+        let for_loop =
+          match label with
+          | None -> for_loop
+          | Some label -> J.Labelled_statement (label, for_loop), J.N
+        in
+        let for_loop = exit_prefix @ [ for_loop ] in
+        match Addr.Set.choose exit_cont with
+        | exception Not_found -> never_body, for_loop
+        | pc ->
+            if Addr.Set.mem pc frontier
+            then never_body, for_loop
+            else
+              let never_after, after = compile_block st [] pc frontier interm in
+              never_body || never_after, for_loop @ after)
 
 and compile_block_no_loop st queue (pc : Addr.t) frontier interm =
   if pc >= 0
@@ -1448,23 +1559,23 @@ and compile_block_no_loop st queue (pc : Addr.t) frontier interm =
       let pc3s = Addr.Set.filter (fun pc -> Hashtbl.mem st.succs pc) pc3s in
       (* no need to limit body for simple flow with no
          instruction.  eg return and branch *)
-      let rec limit pc =
-        if Addr.Set.mem pc exn_frontier
+      let rec keep_frontier st frontier pc =
+        if Addr.Set.mem pc frontier
         then false
         else
           match Addr.Map.find pc st.blocks with
           | { body = []; branch = Return _; _ } -> false
-          | { body = []; branch = Branch (pc', _); _ } -> limit pc'
+          | { body = []; branch = Branch (pc', _); _ } -> keep_frontier st frontier pc'
           | _ -> true
       in
-      let handler_frontier = Addr.Set.filter limit pc3s in
+      let handler_frontier = Addr.Set.filter (keep_frontier st exn_frontier) pc3s in
       (* TODO: Check that we are inside the [frontier/new_frontier] *)
       let handler_frontier =
         resolve_nodes interm (Addr.Set.union exn_frontier handler_frontier)
       in
       Addr.Set.iter (incr_preds st) handler_frontier;
       let prefix, handler_frontier_cont, handler_interm =
-        colapse_frontier st handler_frontier interm
+        colapse_frontier "try-catch" st handler_frontier interm
       in
       assert (Addr.Set.cardinal handler_frontier_cont <= 1);
       let try_catch_frontier = Addr.Set.union new_frontier handler_frontier_cont in
@@ -1550,7 +1661,9 @@ and compile_block_no_loop st queue (pc : Addr.t) frontier interm =
              , source_location st.ctx pc )
             :: after) )
   | _ ->
-      let prefix, frontier_cont, new_interm = colapse_frontier st new_frontier interm in
+      let prefix, frontier_cont, new_interm =
+        colapse_frontier "default" st new_frontier interm
+      in
       assert (Addr.Set.cardinal frontier_cont <= 1);
       List.iter succs ~f:(fun (pc, _) ->
           if Addr.Map.mem pc new_interm then decr_preds st pc);
@@ -1568,13 +1681,14 @@ and compile_block_no_loop st queue (pc : Addr.t) frontier interm =
       in
       never_cond || never_after, seq @ prefix @ cond @ after
 
-and colapse_frontier st new_frontier interm =
+and colapse_frontier name st new_frontier interm =
   if Addr.Set.cardinal new_frontier > 1
   then (
     if debug ()
     then
       Format.eprintf
-        "colapse frontier into %d: %s@."
+        "colapse frontier (%s) into %d: %s@."
+        name
         st.interm_idx
         (string_of_set new_frontier);
     let x = Code.Var.fresh_n "switch" in
@@ -1605,11 +1719,14 @@ and colapse_frontier st new_frontier interm =
     Addr.Set.iter (fun pc -> protect_preds st pc) new_frontier;
     Hashtbl.add st.succs idx (Addr.Set.elements new_frontier);
     Hashtbl.add st.backs idx Addr.Set.empty;
+    let interm =
+      List.fold_right pc_i ~init:interm ~f:(fun (pc, i) interm ->
+          Addr.Map.add pc (idx, (x, i, default = i)) interm)
+    in
     (* The [dominance_frontier_cache] is invalidated by [incr_preds] and [protect_preds] *)
     ( [ J.Variable_statement [ J.V x, Some (int default, J.N) ], J.N ]
     , Addr.Set.singleton idx
-    , List.fold_right pc_i ~init:interm ~f:(fun (pc, i) interm ->
-          Addr.Map.add pc (idx, (x, i, default = i)) interm) ))
+    , interm ))
   else [], new_frontier, interm
 
 and compile_decision_tree st backs frontier interm loc cx dtree =
diff --git a/compiler/tests-compiler/gh1007.ml b/compiler/tests-compiler/gh1007.ml
@@ -156,7 +156,6 @@ let ()  = M.myfun M.x
     {|
     function myfun(x)
      {var x$0=x;
-      a:
       for(;;)
        {if(! x$0)return 0;
         var
@@ -287,9 +286,9 @@ let ()  = M.myfun M.x
          {if(param)
            {var param$0=param[2],len$0=len + 1 | 0,len=len$0,param=param$0;
             continue}
-          if(2 <= len)sort(len,l);
-          var x$0=next;
-          continue a}}}
+          break}
+        if(2 <= len)sort(len,l);
+        var x$0=next}}
     //end |}]
 
 let%expect_test _ =
@@ -505,7 +504,6 @@ let ()  = M.run ()
     {|
     function run(param$0)
      {var i=0;
-      a:
       for(;;)
        {var
          closures=
@@ -549,10 +547,11 @@ let ()  = M.run ()
         for(;;)
          {if(759635106 > param$1[1])
            {var f=param$1[2],param$2=caml_call1(f,0),param$1=param$2;continue}
-          var _g_=i + 1 | 0;
-          if(4 !== i){var i=_g_;continue a}
-          var
-           _e_=caml_call1(Stdlib_List[9],delayed[1]),
-           _f_=function(f){return caml_call1(f,0)};
-          return caml_call2(Stdlib_List[17],_f_,_e_)}}}
+          break}
+        var _g_=i + 1 | 0;
+        if(4 !== i){var i=_g_;continue}
+        var
+         _e_=caml_call1(Stdlib_List[9],delayed[1]),
+         _f_=function(f){return caml_call1(f,0)};
+        return caml_call2(Stdlib_List[17],_f_,_e_)}}
     //end |}]
diff --git a/compiler/tests-compiler/loops.ml b/compiler/tests-compiler/loops.ml
