Compiler: propagate information during optcall optimiastion.

Author: hhugo

compiler/lib/driver.ml

@@ -65,6 +65,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) =
@@ -98,7 +99,6 @@ let ( +> ) f g x = g (f x)
 
 let map_fst4 f (x, y, z, t) = f x, y, z, t
 
-
 let collects_shapes ~shapes (p : Code.program) =
   if debug_shapes () || shapes
   then (
@@ -139,7 +139,6 @@ let collects_shapes ~shapes (p : Code.program) =
     map)
   else StringMap.empty
 
-
 let effects_and_exact_calls ~deadcode_sentinal ~shapes (profile : Profile.t) p =
   let fast =
     match Config.effects (), profile with
@@ -157,21 +156,27 @@ let effects_and_exact_calls ~deadcode_sentinal ~shapes (profile : Profile.t) p =
   in
   match Config.effects () with
   | `Cps | `Double_translation ->
-     if debug () then Format.eprintf "Effects...@.";
-     let shapes = collects_shapes ~shapes p in
-     let p, tramp, cps = Effects.f ~flow_info:info ~live_vars p in
-     let p = match Config.target () with
-           | `Wasm -> p
-           | `JavaScript -> Lambda_lifting.f p in
-     p, tramp, cps, shapes
+      if debug () then Format.eprintf "Effects...@.";
+      let shapes = collects_shapes ~shapes p in
+      let p, tramp, cps = Effects.f ~flow_info:info ~live_vars p in
+      let p =
+        match Config.target () with
+        | `Wasm -> p
+        | `JavaScript -> Lambda_lifting.f p
+      in
+      p, tramp, cps, 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
       , (Code.Var.Set.empty : Effects.trampolined_calls)
-      , (Code.Var.Set.empty : Effects.in_cps), shapes )
+      , (Code.Var.Set.empty : Effects.in_cps)
+      , shapes )
 
 let print p =
   if debug () then Code.Print.program Format.err_formatter (fun _ _ -> "") p;
@@ -644,7 +649,6 @@ let simplify_js js =
   if times () then Format.eprintf "  optimizing: %a@." Timer.print t;
   js
 
-
 let configure formatter =
   let pretty = Config.Flag.pretty () in
   Pretty_print.set_compact formatter (not pretty)
@@ -673,7 +677,7 @@ let optimize ~shapes ~profile p =
        | O2 -> o2
        | O3 -> o3)
     +> specialize_js_once_after
-    +> effects_and_exact_calls ~deadcode_sentinal ~shapes profile 
+    +> effects_and_exact_calls ~deadcode_sentinal ~shapes profile
     +> map_fst4
          (match Config.target (), Config.effects () with
          | `JavaScript, `Disabled -> Generate_closure.f

compiler/lib/global_flow.ml

@@ -716,6 +716,10 @@ let exact_call info f n =
           | Expr _ | Phi _ -> assert false)
         known
 
+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

@@ -44,6 +44,8 @@ type info =
   ; info_return_vals : Var.Set.t Var.Map.t
   }
 
+val update_def : info -> Code.Var.t -> Code.expr -> unit
+
 val f : fast:bool -> Code.program -> info
 
 val exact_call : info -> Var.t -> int -> bool

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