Compiler: exit loops early

Author: hhugo

compiler/lib/structure.ml

@@ -95,8 +95,8 @@ let build_graph blocks pc =
       if leave_try_body block_order preds blocks leave_pc
       then (
         (* Add an edge to limit the [try] body *)
-        Hashtbl.add succs enter_pc (Addr.Set.add leave_pc (Hashtbl.find succs enter_pc));
-        Hashtbl.add preds leave_pc (Addr.Set.add enter_pc (Hashtbl.find preds leave_pc))));
+        Hashtbl.replace succs enter_pc (Addr.Set.add leave_pc (Hashtbl.find succs enter_pc));
+        Hashtbl.replace preds leave_pc (Addr.Set.add enter_pc (Hashtbl.find preds leave_pc))));
   { succs; preds; reverse_post_order = !l; block_order }
 
 let dominator_tree g =
@@ -166,3 +166,82 @@ let dominance_frontier g idom =
     g.preds;
   frontiers
 *)
+
+
+
+(* Compute a map from each block to the set of loops it belongs to *)
+let mark_loops g =
+  let in_loop = Hashtbl.create 16 in
+  Hashtbl.iter
+    (fun pc preds ->
+      let rec mark_loop pc' =
+        if not (Addr.Set.mem pc (get_edges in_loop pc'))
+        then (
+          add_edge in_loop pc' pc;
+          if pc' <> pc then Addr.Set.iter mark_loop (Hashtbl.find g.preds pc'))
+      in
+      Addr.Set.iter (fun pc' -> if is_backward g pc' pc then mark_loop pc') preds)
+    g.preds;
+  in_loop
+
+let rec measure blocks g pc limit =
+  let b = Addr.Map.find pc blocks in
+  let limit = limit - List.length b.body in
+  if limit < 0
+  then limit
+  else
+    Addr.Set.fold
+      (fun pc limit -> if limit < 0 then limit else measure blocks g pc limit)
+      (get_edges g.succs pc)
+      limit
+
+let is_small blocks g pc = measure blocks g pc 20 >= 0
+
+let shrink_loops blocks ({ succs; preds; reverse_post_order; _ } as g) =
+  let add_edge pred succ =
+    Hashtbl.replace succs pred (Addr.Set.add succ (Hashtbl.find succs pred));
+    Hashtbl.replace preds succ (Addr.Set.add pred (Hashtbl.find preds succ))
+  in
+  let in_loop = mark_loops g in
+  let dom = dominator_tree g in
+  let root = List.hd reverse_post_order in
+  let rec traverse ignored pc =
+    let succs = get_edges dom pc in
+    let loops = get_edges in_loop pc in
+    let block = Addr.Map.find pc blocks in
+    Addr.Set.iter
+      (fun pc' ->
+        (* Whatever is in the scope of an exception handler should not be
+           moved outside *)
+        let ignored =
+          match fst block.branch with
+          | Pushtrap ((body_pc, _), _, _, _) when pc' = body_pc ->
+              Addr.Set.union ignored loops
+          | _ -> ignored
+        in
+        let loops' = get_edges in_loop pc' in
+        let left_loops = Addr.Set.diff (Addr.Set.diff loops loops') ignored in
+        (* If we leave a loop, we add an edge from a predecessor of
+           the loop header to the current block, so that it is
+           considered outside of the loop. *)
+        if not (Addr.Set.is_empty left_loops || is_small blocks g pc')
+        then
+          Addr.Set.iter
+            (fun pc0 ->
+              match
+                Addr.Set.find_first
+                  (fun pc -> is_forward g pc pc0)
+                  (get_edges g.preds pc0)
+              with
+              | pc -> add_edge pc pc'
+              | exception Not_found -> ())
+            left_loops;
+        traverse ignored pc')
+      succs
+  in
+  traverse Addr.Set.empty root
+
+let build_graph blocks pc =
+  let g = build_graph blocks pc in
+  shrink_loops blocks g;
+  g

compiler/tests-compiler/gh1007.ml

@@ -319,9 +319,11 @@ let ()  = M.myfun M.x
        len = 0,
        param = l;
       for(;;){
-       if(! param){if(2 <= len) sort(len, l); var x$0 = next; break;}
+       if(! param) break;
        var l$0 = param[2], len$0 = len + 1 | 0, len = len$0, param = l$0;
       }
+      if(2 <= len) sort(len, l);
+      var x$0 = next;
      }
     }
     //end |}]
@@ -625,14 +627,15 @@ let ()  = M.run ()
        even = closures$0[1],
        param$0 = even(i);
       for(;;){
-       if(759635106 <= param$0[1]){
-        var _g_ = i + 1 | 0;
-        if(4 !== i){var i = _g_; break;}
-        var _f_ = caml_call1(list_rev, delayed[1]);
-        return caml_call2(list_iter, function(f){return caml_call1(f, 0);}, _f_);
-       }
+       if(759635106 <= param$0[1]) break;
        var f = param$0[2], param$0 = f(0);
       }
+      var _g_ = i + 1 | 0;
+      if(4 === i){
+       var _f_ = caml_call1(list_rev, delayed[1]);
+       return caml_call2(list_iter, function(f){return caml_call1(f, 0);}, _f_);
+      }
+      var i = _g_;
      }
     }
     //end |}]

compiler/tests-compiler/loops.ml

@@ -100,9 +100,10 @@ let rec fun_with_loop acc = function
       for(;;){
        a[1] = [0, 1, a[1]];
        var _a_ = i + 1 | 0;
-       if(10 === i){var acc$1 = [0, x, a[1]], acc$0 = acc$1, param$0 = xs; break;}
+       if(10 === i) break;
        var i = _a_;
       }
+      var acc$1 = [0, x, a[1]], acc$0 = acc$1, param$0 = xs;
      }
     }
     //end
@@ -130,19 +131,15 @@ let for_for_while () =
      var k = 1;
      for(;;){
       var j = 1;
-      a:
-      for(;;)
-       for(;;){
-        if(10 <= runtime.caml_mul(k, j)){
-         var _b_ = j + 1 | 0;
-         if(10 !== j){var j = _b_; break;}
-         var _a_ = k + 1 | 0;
-         if(10 === k) return 0;
-         var k = _a_;
-         break a;
-        }
-        id[1]++;
-       }
+      for(;;){
+       for(;;){if(10 <= runtime.caml_mul(k, j)) break; id[1]++;}
+       var _b_ = j + 1 | 0;
+       if(10 === j) break;
+       var j = _b_;
+      }
+      var _a_ = k + 1 | 0;
+      if(10 === k) return 0;
+      var k = _a_;
      }
     }
     //end |}]
@@ -314,26 +311,25 @@ in loop x
      var x$1 = x;
      for(;;){
       if(0 === x$1) return 1;
-      if(1 === x$1){
-       var x$0 = 2;
-       for(;;){
-        a:
-        {
-         if(3 >= x$0 >>> 0)
-          switch(x$0){
-            case 0:
-             var _a_ = 1; break a;
-            case 2:
-             var n = caml_call1(Stdlib_Random[5], 2), _a_ = n + n | 0; break a;
-            case 3:
-             var n$0 = caml_call1(Stdlib_Random[5], 2), x$0 = n$0; continue;
-          }
-         var _a_ = 2;
+      if(1 === x$1) break;
+      var x$2 = x$1 + 1 | 0, x$1 = x$2;
+     }
+     var x$0 = 2;
+     for(;;){
+      a:
+      {
+       if(3 >= x$0 >>> 0)
+        switch(x$0){
+          case 0:
+           var _a_ = 1; break a;
+          case 2:
+           var n = caml_call1(Stdlib_Random[5], 2), _a_ = n + n | 0; break a;
+          case 3:
+           var n$0 = caml_call1(Stdlib_Random[5], 2), x$0 = n$0; continue;
         }
-        return _a_ + 2 | 0;
-       }
+       var _a_ = 2;
       }
-      var x$2 = x$1 + 1 | 0, x$1 = x$2;
+      return _a_ + 2 | 0;
      }
     }
     //end |}]
@@ -455,42 +451,39 @@ let add_substitute =
         a:
         {
          if(40 !== opening && 123 !== opening){
-          var
-           start = start$0 + 1 | 0,
-           lim$0 = caml_ml_string_length(s),
-           i$2 = start;
-          for(;;){
-           b:
+          var start = start$0 + 1 | 0, lim$0 = caml_ml_string_length(s);
+          b:
+          {
+           c:
            {
-            if(lim$0 > i$2){
-             var match = caml_string_get(s, i$2);
-             c:
-             {
+            d:
+            {
+             var i$2 = start;
+             for(;;){
+              if(lim$0 <= i$2) break c;
+              var match = caml_string_get(s, i$2);
               if(91 <= match){
                if(97 <= match){
-                if(123 > match) break c;
+                if(123 <= match) break d;
                }
-               else if(95 === match) break c;
+               else if(95 !== match) break d;
               }
               else if(58 <= match){
-               if(65 <= match) break c;
+               if(65 > match) break;
               }
-              else if(48 <= match) break c;
-              var stop$0 = i$2;
-              break b;
+              else if(48 > match) break d;
+              var i$3 = i$2 + 1 | 0, i$2 = i$3;
              }
-             var i$3 = i$2 + 1 | 0, i$2 = i$3;
-             continue;
             }
-            var stop$0 = lim$0;
+            var stop$0 = i$2;
+            break b;
            }
-           var
-            match$0 =
-              [0,
-               caml_call3(string_sub, s, start$0, stop$0 - start$0 | 0),
-               stop$0];
-           break a;
+           var stop$0 = lim$0;
           }
+          var
+           match$0 =
+             [0, caml_call3(string_sub, s, start$0, stop$0 - start$0 | 0), stop$0];
+          break a;
          }
          var new_start = start$0 + 1 | 0, k$2 = 0;
          if(40 === opening)
@@ -506,19 +499,17 @@ let add_substitute =
           if(caml_string_get(s, stop) === opening)
            var i = stop + 1 | 0, k$0 = k + 1 | 0, k = k$0, stop = i;
           else if(caml_string_get(s, stop) === closing){
-           if(0 === k){
-            var
-             match$0 =
-               [0,
-                caml_call3(string_sub, s, new_start, (stop - start$0 | 0) - 1 | 0),
-                stop + 1 | 0];
-            break;
-           }
+           if(0 === k) break;
            var i$0 = stop + 1 | 0, k$1 = k - 1 | 0, k = k$1, stop = i$0;
           }
           else
            var i$1 = stop + 1 | 0, stop = i$1;
          }
+         var
+          match$0 =
+            [0,
+             caml_call3(string_sub, s, new_start, (stop - start$0 | 0) - 1 | 0),
+             stop + 1 | 0];
         }
         var next_i = match$0[2], ident = match$0[1];
         caml_call2(add_string, b, caml_call1(f, ident));