fixup! Reference unboxing

Author: vouillon

compiler/lib/driver.ml

@@ -216,8 +216,8 @@ let rec loop max name round i (p : 'a) : 'a =
 let round profile : 'a -> 'a =
   print
   +> tailcall
-  +> (flow +> specialize +> eval +> fst)
   +> Ref_unboxing.f
+  +> (flow +> specialize +> eval +> fst)
   +> inline profile
   +> phi
   +> deadcode

compiler/lib/phisimpl.ml

@@ -69,11 +69,6 @@ let program_deps { blocks; _ } =
     (fun _pc block ->
       List.iter block.body ~f:(fun i ->
           match i with
-          | Let (x, Prim (Extern "%identity", [ Pv y ])) ->
-              (* This is introduced by the reference unboxing pass *)
-              add_var vars x;
-              add_dep deps x y;
-              add_def vars defs x y
           | Let (x, e) ->
               add_var vars x;
               expr_deps blocks vars deps defs x e

compiler/lib/ref_unboxing.ml

@@ -37,45 +37,48 @@ let times = Debug.find "times"
 
 let stats = Debug.find "stats"
 
-let rewrite refs block m =
-  let m, l =
+let rewrite refs block m m' =
+  let m, m', l =
     List.fold_left
-      ~f:(fun (m, acc) i ->
+      ~f:(fun (m, m', acc) i ->
         match i with
         | Let (x, Block (0, [| y |], (NotArray | Unknown), Maybe_mutable))
-          when Var.Set.mem x refs -> Var.Map.add x y m, acc
+          when Var.Set.mem x refs -> Var.Map.add x y m, m', acc
         | Let (y, Field (x, 0, Non_float)) when Var.Map.mem x m ->
             (* Optimized away by Phisimpl *)
-            m, Let (y, Prim (Extern "%identity", [ Pv (Var.Map.find x m) ])) :: acc
+            m, Var.Map.add y (Var.Map.find x m) m', acc
         | Offset_ref (x, n) when Var.Map.mem x m ->
             let y = Var.fresh () in
             ( Var.Map.add x y m
+            , m'
             , Let
                 ( y
                 , Prim
                     ( Extern "%int_add"
                     , [ Pv (Var.Map.find x m); Pc (Int (Targetint.of_int_exn n)) ] ) )
               :: acc )
-        | Set_field (x, 0, Non_float, y) when Var.Map.mem x m -> Var.Map.add x y m, acc
+        | Set_field (x, 0, Non_float, y) when Var.Map.mem x m ->
+            Var.Map.add x y m, m', acc
         | Event _ -> (
             ( m
+            , m'
             , match acc with
               | Event _ :: prev ->
                   (* Avoid consecutive events (keep just the last one) *)
                   i :: prev
               | _ -> i :: acc ))
-        | _ -> m, i :: acc)
+        | _ -> m, m', i :: acc)
       block.body
-      ~init:(m, [])
+      ~init:(m, m', [])
   in
-  m, List.rev l
+  m, m', List.rev l
 
 let rewrite_cont relevant_vars vars (pc', args) =
   let refs, _ = Int.Hashtbl.find relevant_vars pc' in
   let vars = Var.Map.filter (fun x _ -> Var.Set.mem x refs) vars in
   pc', List.map ~f:snd (Var.Map.bindings vars) @ args
 
-let rewrite_function p variables pc =
+let rewrite_function p variables pc subst =
   let relevant_vars = Int.Hashtbl.create 16 in
   let g = Structure.(dominator_tree (build_graph p.blocks pc)) in
   let rec traverse_tree g pc refs =
@@ -94,14 +97,14 @@ let rewrite_function p variables pc =
     Addr.Set.iter (fun pc' -> traverse_tree g pc' refs') (Structure.get_edges g pc)
   in
   traverse_tree g pc Var.Set.empty;
-  let rec traverse_tree' g pc blocks =
+  let rec traverse_tree' g pc blocks subst =
     let block = Addr.Map.find pc p.blocks in
     let refs, refs' = Int.Hashtbl.find relevant_vars pc in
     let vars =
       Var.Set.fold (fun x m -> Var.Map.add x (Var.fork x) m) refs Var.Map.empty
     in
     let params = List.map ~f:snd (Var.Map.bindings vars) @ block.params in
-    let vars, body = rewrite refs' block vars in
+    let vars, subst, body = rewrite refs' block vars subst in
     let branch =
       match block.branch with
       | Return _ | Raise _ | Stop -> block.branch
@@ -122,12 +125,12 @@ let rewrite_function p variables pc =
     in
     let blocks = Addr.Map.add pc { params; body; branch } blocks in
     Addr.Set.fold
-      (fun pc' blocks -> traverse_tree' g pc' blocks)
+      (fun pc' (blocks, subst) -> traverse_tree' g pc' blocks subst)
       (Structure.get_edges g pc)
-      blocks
+      (blocks, subst)
   in
-  let blocks = traverse_tree' g pc p.blocks in
-  { p with blocks }
+  let blocks, subst = traverse_tree' g pc p.blocks subst in
+  { p with blocks }, subst
 
 let f p =
   let t = Timer.make () in
@@ -205,7 +208,17 @@ let f p =
   let functions =
     Var.Hashtbl.fold (fun _ pc s -> Addr.Set.add pc s) candidates Addr.Set.empty
   in
-  let p = Addr.Set.fold (fun pc p -> rewrite_function p candidates pc) functions p in
+  let p, subst =
+    Addr.Set.fold
+      (fun pc (p, subst) -> rewrite_function p candidates pc subst)
+      functions
+      (p, Var.Map.empty)
+  in
+  let p =
+    if Var.Map.is_empty subst
+    then p
+    else Subst.Excluding_Binders.program (Subst.from_map subst) p
+  in
   if times () then Format.eprintf "  reference unboxing: %a@." Timer.print t;
   if stats ()
   then Format.eprintf "Stats - reference unboxing: %d@." (Var.Hashtbl.length candidates);

compiler/tests-compiler/global_deadcode.ml

@@ -146,7 +146,10 @@ let%expect_test "Bug fix in PR #1681" =
           x.a <- 1; (* This has to be handled after [x] is returned *)
           {a = 3; b = 4}
         )
-      let g = ref (fun _ -> assert false)
+      let g =
+        (* Make sure the reference is not optimized away *)
+        let rec h n = if n = 0 then ref else h (n - 1) in
+        h 0 (fun _ -> assert false)
       let _ =
         (* We should not track that [f] is used below *)
         g := f; prerr_int ((!g true).b + (!g false).b)

compiler/tests-compiler/loops.ml

@@ -582,23 +582,22 @@ let () = print_endline (trim " ")
      var
       s$0 = copy(caml_bytes_of_string(x)),
       len = caml_ml_bytes_length(s$0),
-      ofs = 0;
+      i = 0;
      for(;;){
-      if(ofs >= len) break;
-      if(! is_space(caml_bytes_unsafe_get(s$0, ofs))) break;
-      ofs = ofs + 1 | 0;
+      if(i >= len) break;
+      if(! is_space(caml_bytes_unsafe_get(s$0, i))) break;
+      i = i + 1 | 0;
      }
      var j = len - 1 | 0;
      for(;;){
-      if(ofs <= j && is_space(caml_bytes_unsafe_get(s$0, j))){j = j - 1 | 0; continue;}
+      if(i <= j && is_space(caml_bytes_unsafe_get(s$0, j))){j = j - 1 | 0; continue;}
       a:
       {
-       if(ofs <= j){
-        var len$0 = (j - ofs | 0) + 1 | 0;
-        if
-         (0 <= ofs && 0 <= len$0 && (caml_ml_bytes_length(s$0) - len$0 | 0) >= ofs){
+       if(i <= j){
+        var len$0 = (j - i | 0) + 1 | 0;
+        if(0 <= i && 0 <= len$0 && (caml_ml_bytes_length(s$0) - len$0 | 0) >= i){
          var r = caml_create_bytes(len$0);
-         caml_blit_bytes(s$0, ofs, r, 0, len$0);
+         caml_blit_bytes(s$0, i, r, 0, len$0);
          var b = r;
          break a;
         }