Compiler: propagate information during optcall optimiastion.

Author: hhugo

compiler/lib/driver.ml

@@ -68,6 +68,7 @@ let specialize_1 (p, info) =
   let return_values = Code.Var.Map.empty in
   Specialize.f
     ~function_arity:(fun f -> Specialize.function_arity ~return_values info f)
+    ~update_def:(fun x expr -> Flow.Info.update_def info x expr)
     p
 
 let specialize_js (p, info) =
@@ -177,7 +178,10 @@ let effects_and_exact_calls
       p, trampolined_calls, in_cps, None, shapes
   | `Disabled | `Jspi ->
       let p =
-        Specialize.f ~function_arity:(fun f -> Global_flow.function_arity info f) p
+        Specialize.f
+          ~function_arity:(fun f -> Global_flow.function_arity info f)
+          ~update_def:(fun x expr -> Global_flow.update_def info x expr)
+          p
       in
       let shapes = collects_shapes ~shapes p in
       ( p

compiler/lib/global_flow.ml

@@ -804,6 +804,10 @@ let get_unique_closure info f =
         | None -> None
         | Some kind -> kind)
 
+let update_def info x expr =
+  let idx = Code.Var.idx x in
+  info.info_defs.(idx) <- Expr expr
+
 let function_arity info f =
   match Var.Tbl.get info.info_approximation f with
   | Top | Values { others = true; _ } -> None

compiler/lib/global_flow.mli

@@ -80,6 +80,8 @@ type state =
 
 val f : fast:bool -> Code.program -> state * info
 
+val update_def : info -> Code.Var.t -> Code.expr -> unit
+
 val exact_call : info -> Var.t -> int -> bool
 
 val get_unique_closure : info -> Var.t -> Var.t option

compiler/lib/specialize.ml

@@ -40,7 +40,7 @@ let unknown_apply = function
   | Let (_, Apply { f = _; args = _; exact = false }) -> true
   | _ -> false
 
-let specialize_apply opt_count function_arity ((acc, free_pc, extra), loc) i =
+let specialize_apply opt_count function_arity update_def ((acc, free_pc, extra), loc) i =
   match i with
   | Let (x, Apply { f; args; exact = false }) -> (
       let n' = List.length args in
@@ -74,13 +74,13 @@ let specialize_apply opt_count function_arity ((acc, free_pc, extra), loc) i =
             ; branch = Return return'
             }
           in
-          ( Let (x, Closure (missing, (free_pc, missing), None)) :: acc
-          , free_pc + 1
-          , (free_pc, block) :: extra )
+          let expr = Closure (missing, (free_pc, missing), None) in
+          update_def x expr;
+          Let (x, expr) :: acc, free_pc + 1, (free_pc, block) :: extra
       | Some _ -> assert false)
   | _ -> i :: acc, free_pc, extra
 
-let specialize_instrs ~function_arity opt_count p =
+let specialize_instrs ~function_arity ~update_def opt_count p =
   let blocks, free_pc =
     Addr.Map.fold
       (fun pc block (blocks, free_pc) ->
@@ -95,7 +95,7 @@ let specialize_instrs ~function_arity opt_count p =
                 | Event loc ->
                     let (body, free_pc, extra), _ = acc in
                     (i :: body, free_pc, extra), Some loc
-                | _ -> specialize_apply opt_count function_arity acc i, None)
+                | _ -> specialize_apply opt_count function_arity update_def acc i, None)
           in
           let blocks =
             List.fold_left extra ~init:blocks ~f:(fun blocks (pc, b) ->
@@ -108,13 +108,15 @@ let specialize_instrs ~function_arity opt_count p =
   in
   { p with blocks; free_pc }
 
-let f ~function_arity p =
+let f ~function_arity ~update_def p =
   Code.invariant p;
   let previous_p = p in
   let t = Timer.make () in
   let opt_count = ref 0 in
   let p =
-    if Config.Flag.optcall () then specialize_instrs ~function_arity opt_count p else p
+    if Config.Flag.optcall ()
+    then specialize_instrs ~function_arity ~update_def opt_count p
+    else p
   in
   if times () then Format.eprintf "  optcall: %a@." Timer.print t;
   if stats () then Format.eprintf "Stats - optcall: %d@." !opt_count;

compiler/lib/specialize.mli

@@ -21,6 +21,10 @@
 val function_arity :
   return_values:Code.Var.Set.t Code.Var.Map.t -> Flow.Info.t -> Code.Var.t -> int option
 
-val f : function_arity:(Code.Var.t -> int option) -> Code.program -> Code.program
+val f :
+     function_arity:(Code.Var.t -> int option)
+  -> update_def:(Code.Var.t -> Code.expr -> unit)
+  -> Code.program
+  -> Code.program
 
 val switches : Code.program -> Code.program