small refactoring

Author: hhugo

compiler/lib/ref_unboxing.ml

@@ -74,30 +74,51 @@ let rewrite_body unboxed_refs body ref_contents subst =
   in
   ref_contents, subst, List.rev l
 
-let rewrite_cont relevant_vars vars (pc', args) =
+let rewrite_cont relevant_vars ref_contents (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
+  let vars = Var.Map.filter (fun x _ -> Var.Set.mem x refs) ref_contents in
   pc', List.map ~f:snd (Var.Map.bindings vars) @ args
 
-let rewrite_function p variables pc subst =
-  let relevant_vars = Int.Hashtbl.create 16 in
+let rewrite_branch relevant_vars ref_contents branch =
+  match branch with
+  | Return _ | Raise _ | Stop -> branch
+  | Branch cont -> Branch (rewrite_cont relevant_vars ref_contents cont)
+  | Cond (x, cont, cont') ->
+      Cond
+        ( x
+        , rewrite_cont relevant_vars ref_contents cont
+        , rewrite_cont relevant_vars ref_contents cont' )
+  | Switch (x, a) ->
+      Switch (x, Array.map ~f:(fun cont -> rewrite_cont relevant_vars ref_contents cont) a)
+  | Pushtrap (cont, x, cont') ->
+      Pushtrap
+        ( rewrite_cont relevant_vars ref_contents cont
+        , x
+        , rewrite_cont relevant_vars ref_contents cont' )
+  | Poptrap cont -> Poptrap (rewrite_cont relevant_vars ref_contents cont)
+
+let rewrite_function p ~unboxed_refs pc subst =
   let g = Structure.(dominator_tree (build_graph p.blocks pc)) in
-  let rec traverse_tree g pc refs =
-    let block = Addr.Map.find pc p.blocks in
-    let refs' =
-      List.fold_left
-        ~f:(fun s i ->
-          match i with
-          | Let (x, Block (0, [| _ |], (NotArray | Unknown), Maybe_mutable))
-            when Var.Hashtbl.mem variables x -> Var.Set.add x s
-          | _ -> s)
-        ~init:refs
-        block.body
+  let relevant_vars =
+    let relevant_vars = Int.Hashtbl.create 16 in
+    let rec traverse_tree g pc refs =
+      let block = Addr.Map.find pc p.blocks in
+      let refs' =
+        List.fold_left
+          ~f:(fun s i ->
+            match i with
+            | Let (x, Block (0, [| _ |], (NotArray | Unknown), Maybe_mutable))
+              when Var.Hashtbl.mem unboxed_refs x -> Var.Set.add x s
+            | _ -> s)
+          ~init:refs
+          block.body
+      in
+      Int.Hashtbl.add relevant_vars pc (refs, refs');
+      Addr.Set.iter (fun pc' -> traverse_tree g pc' refs') (Structure.get_edges g pc)
     in
-    Int.Hashtbl.add relevant_vars pc (refs, refs');
-    Addr.Set.iter (fun pc' -> traverse_tree g pc' refs') (Structure.get_edges g pc)
+    traverse_tree g pc Var.Set.empty;
+    relevant_vars
   in
-  traverse_tree g pc Var.Set.empty;
   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
@@ -106,25 +127,7 @@ let rewrite_function p variables pc subst =
     in
     let params = List.map ~f:snd (Var.Map.bindings ref_contents) @ block.params in
     let ref_contents, subst, body = rewrite_body refs' block.body ref_contents subst in
-    let branch =
-      match block.branch with
-      | Return _ | Raise _ | Stop -> block.branch
-      | Branch cont -> Branch (rewrite_cont relevant_vars ref_contents cont)
-      | Cond (x, cont, cont') ->
-          Cond
-            ( x
-            , rewrite_cont relevant_vars ref_contents cont
-            , rewrite_cont relevant_vars ref_contents cont' )
-      | Switch (x, a) ->
-          Switch
-            (x, Array.map ~f:(fun cont -> rewrite_cont relevant_vars ref_contents cont) a)
-      | Pushtrap (cont, x, cont') ->
-          Pushtrap
-            ( rewrite_cont relevant_vars ref_contents cont
-            , x
-            , rewrite_cont relevant_vars ref_contents cont' )
-      | Poptrap cont -> Poptrap (rewrite_cont relevant_vars ref_contents cont)
-    in
+    let branch = rewrite_branch relevant_vars ref_contents block.branch in
     let blocks = Addr.Map.add pc { params; body; branch } blocks in
     Addr.Set.fold
       (fun pc' (blocks, subst) -> traverse_tree' g pc' blocks subst)
@@ -212,7 +215,7 @@ let f p =
   in
   let p, subst =
     Addr.Set.fold
-      (fun pc (p, subst) -> rewrite_function p candidates pc subst)
+      (fun pc (p, subst) -> rewrite_function p ~unboxed_refs:candidates pc subst)
       functions
       (p, Var.Map.empty)
   in