Compiler: propagate arity across unit boundary (WIP)

Author: hhugo

compiler/lib/code.ml

@@ -339,7 +339,7 @@ type expr =
       ; args : Var.t list
       ; exact : bool
       }
-  | Block of int * Var.t array * array_or_not
+  | Block of int * Var.t array * array_or_not * bool
   | Field of Var.t * int
   | Closure of Var.t list * cont
   | Constant of constant
@@ -479,8 +479,18 @@ module Print = struct
         if exact
         then Format.fprintf f "%a!(%a)" Var.print g var_list args
         else Format.fprintf f "%a(%a)" Var.print g var_list args
-    | Block (t, a, _) ->
-        Format.fprintf f "{tag=%d" t;
+    | Block (t, a, k, imm) ->
+        Format.fprintf
+          f
+          "{%s%s:tag=%d"
+          (match imm with
+          | true -> "Imm"
+          | false -> "")
+          (match k with
+          | Array -> "A"
+          | NotArray -> "NA"
+          | Unknown -> "U")
+          t;
         for i = 0 to Array.length a - 1 do
           Format.fprintf f "; %d = %a" i Var.print a.(i)
         done;
@@ -732,7 +742,7 @@ let invariant { blocks; start; _ } =
     in
     let check_expr = function
       | Apply _ -> ()
-      | Block (_, _, _) -> ()
+      | Block (_, _, _, _) -> ()
       | Field (_, _) -> ()
       | Closure (l, cont) ->
           List.iter l ~f:define;

compiler/lib/code.mli

@@ -186,7 +186,7 @@ type expr =
       ; args : Var.t list
       ; exact : bool (* if true, then # of arguments = # of parameters *)
       }
-  | Block of int * Var.t array * array_or_not
+  | Block of int * Var.t array * array_or_not * bool
   | Field of Var.t * int
   | Closure of Var.t list * cont
   | Constant of constant

compiler/lib/deadcode.ml

@@ -70,7 +70,7 @@ and mark_expr st e =
   | Apply { f; args; _ } ->
       mark_var st f;
       List.iter args ~f:(fun x -> mark_var st x)
-  | Block (_, a, _) -> Array.iter a ~f:(fun x -> mark_var st x)
+  | Block (_, a, _, _) -> Array.iter a ~f:(fun x -> mark_var st x)
   | Field (x, _) -> mark_var st x
   | Closure (_, (pc, _)) -> mark_reachable st pc
   | Special _ -> ()

compiler/lib/driver.ml

@@ -625,6 +625,29 @@ if (typeof module === 'object' && module.exports) {
   if times () then Format.eprintf "  optimizing: %a@." Timer.print t;
   js
 
+let collects_shapes p =
+  let _, info = Flow.f p in
+  let l = ref StringMap.empty in
+  Code.Addr.Map.iter
+    (fun _ block ->
+      List.iter block.Code.body ~f:(fun (i, _) ->
+          match i with
+          | Code.Let
+              ( _
+              , Prim
+                  ( Extern "caml_register_global"
+                  , [ _code; Pv block; Pc (NativeString name) ] ) ) ->
+              let shape = Flow.the_shape_of info block in
+              let name =
+                match name with
+                | Byte s -> s
+                | Utf (Utf8 s) -> s
+              in
+              l := StringMap.add name shape !l
+          | _ -> ()))
+    p.blocks;
+  !l
+
 let configure formatter =
   let pretty = Config.Flag.pretty () in
   Pretty_print.set_compact formatter (not pretty);
@@ -663,7 +686,15 @@ let full ~standalone ~wrap_with_fun ~profile ~linkall ~source_map formatter d p
   in
   if times () then Format.eprintf "Start Optimizing...@.";
   let t = Timer.make () in
-  let r = opt p in
+  let (((prog, _), _) as r) = opt p in
+  let shapes = collects_shapes prog in
+  StringMap.iter
+    (fun name shape ->
+      Shape.set_shape ~name shape;
+      Pretty_print.string
+        formatter
+        (Printf.sprintf "//# shape: %s:%s\n" name (Shape.to_string shape)))
+    shapes;
   let () = if times () then Format.eprintf " optimizations : %a@." Timer.print t in
   emit r
 

compiler/lib/duplicate.ml

@@ -26,7 +26,7 @@ let expr s e =
   | Constant _ -> e
   | Apply { f; args; exact } ->
       Apply { f = s f; args = List.map args ~f:(fun x -> s x); exact }
-  | Block (n, a, k) -> Block (n, Array.map a ~f:(fun x -> s x), k)
+  | Block (n, a, k, imm) -> Block (n, Array.map a ~f:(fun x -> s x), k, imm)
   | Field (x, n) -> Field (s x, n)
   | Closure _ -> failwith "Inlining/Duplicating closure is currenly not supported"
   | Special x -> Special x

compiler/lib/eval.ml

@@ -177,7 +177,7 @@ let is_int info x =
         (fun x ->
           match Flow.Info.def info x with
           | Some (Constant (Int _)) -> Y
-          | Some (Block (_, _, _) | Constant _) -> N
+          | Some (Block (_, _, _, _) | Constant _) -> N
           | None | Some _ -> Unknown)
         Unknown
         (fun u v ->
@@ -196,7 +196,7 @@ let the_tag_of info x get =
         info
         (fun x ->
           match Flow.Info.def info x with
-          | Some (Block (j, _, _)) ->
+          | Some (Block (j, _, _, _)) ->
               if Flow.Info.possibly_mutable info x then None else get j
           | Some (Constant (Tuple (j, _, _))) -> get j
           | None | Some _ -> None)
@@ -278,7 +278,7 @@ let eval_instr info ((x, loc) as i) =
   | Let (x, Prim (Extern "caml_sys_const_backend_type", [ _ ])) ->
       let jsoo = Code.Var.fresh () in
       [ Let (jsoo, Constant (String "js_of_ocaml")), noloc
-      ; Let (x, Block (0, [| jsoo |], NotArray)), loc
+      ; Let (x, Block (0, [| jsoo |], NotArray, true)), loc
       ]
   | Let (_, Prim (Extern ("%resume" | "%perform" | "%reperform"), _)) ->
       [ i ] (* We need that the arguments to this primitives remain variables *)
@@ -338,7 +338,7 @@ let the_cond_of info x =
             | NativeString _
             | Float_array _
             | Int64 _ )) -> Non_zero
-      | Some (Block (_, _, _)) -> Non_zero
+      | Some (Block (_, _, _, _)) -> Non_zero
       | Some (Field _ | Closure _ | Prim _ | Apply _ | Special _) -> Unknown
       | None -> Unknown)
     Unknown
@@ -381,7 +381,7 @@ let rec do_not_raise pc visited blocks =
         | Array_set (_, _, _) | Offset_ref (_, _) | Set_field (_, _, _) | Assign _ -> ()
         | Let (_, e) -> (
             match e with
-            | Block (_, _, _) | Field (_, _) | Constant _ | Closure _ -> ()
+            | Block (_, _, _, _) | Field (_, _) | Constant _ | Closure _ -> ()
             | Apply _ -> raise May_raise
             | Special _ -> ()
             | Prim (Extern name, _) when Primitive.is_pure name -> ()

compiler/lib/flow.ml

@@ -103,7 +103,7 @@ let expr_deps blocks vars deps defs x e =
   | Closure (l, cont) ->
       List.iter l ~f:(fun x -> add_param_def vars defs x);
       cont_deps blocks vars deps defs cont
-  | Block (_, a, _) -> Array.iter a ~f:(fun y -> add_dep deps x y)
+  | Block (_, a, _, _) -> Array.iter a ~f:(fun y -> add_dep deps x y)
   | Field (y, _) -> add_dep deps x y
 
 let program_deps { blocks; _ } =
@@ -152,7 +152,7 @@ let propagate1 deps defs st x =
           var_set_lift
             (fun z ->
               match defs.(Var.idx z) with
-              | Expr (Block (_, a, _)) when n < Array.length a ->
+              | Expr (Block (_, a, _, _)) when n < Array.length a ->
                   let t = a.(n) in
                   add_dep deps x t;
                   Var.Tbl.get st t
@@ -186,15 +186,17 @@ type mutability_state =
   ; possibly_mutable : Code.Var.ISet.t
   }
 
-let rec block_escape st x =
+let rec block_escape st ?(immutable = false) x =
   Var.Set.iter
     (fun y ->
       if not (Code.Var.ISet.mem st.may_escape y)
       then (
         Code.Var.ISet.add st.may_escape y;
-        Code.Var.ISet.add st.possibly_mutable y;
+        if not immutable then Code.Var.ISet.add st.possibly_mutable y;
+
         match st.defs.(Var.idx y) with
-        | Expr (Block (_, l, _)) -> Array.iter l ~f:(fun z -> block_escape st z)
+        | Expr (Block (_, l, _, immutable)) ->
+            Array.iter l ~f:(fun z -> block_escape st ~immutable z)
         | _ -> ()))
     (Var.Tbl.get st.known_origins x)
 
@@ -226,15 +228,18 @@ let expr_escape st _x e =
             | Pv v, `Shallow_const -> (
                 match st.defs.(Var.idx v) with
                 | Expr (Constant (Tuple _)) -> ()
-                | Expr (Block (_, a, _)) -> Array.iter a ~f:(fun x -> block_escape st x)
+                | Expr (Block (_, a, _, true)) ->
+                    Array.iter a ~f:(fun x -> block_escape st ~immutable:true x)
+                | Expr (Block (_, a, _, false)) ->
+                    Array.iter a ~f:(fun x -> block_escape st x)
                 | _ -> block_escape st v)
             | Pv v, `Object_literal -> (
                 match st.defs.(Var.idx v) with
                 | Expr (Constant (Tuple _)) -> ()
-                | Expr (Block (_, a, _)) ->
+                | Expr (Block (_, a, _, _)) ->
                     Array.iter a ~f:(fun x ->
                         match st.defs.(Var.idx x) with
-                        | Expr (Block (_, [| _k; v |], _)) -> block_escape st v
+                        | Expr (Block (_, [| _k; v |], _, _)) -> block_escape st v
                         | Expr (Constant _) -> ()
                         | _ -> block_escape st x)
                 | _ -> block_escape st v)
@@ -282,7 +287,7 @@ let propagate2 ?(skip_param = false) defs known_origins possibly_mutable st x =
           || Var.Set.exists
                (fun z ->
                  match defs.(Var.idx z) with
-                 | Expr (Block (_, a, _)) ->
+                 | Expr (Block (_, a, _, _)) ->
                      n >= Array.length a
                      || Var.ISet.mem possibly_mutable z
                      || Var.Tbl.get st a.(n)
@@ -382,7 +387,7 @@ let direct_approx (info : Info.t) x =
           then None
           else
             match info.info_defs.(Var.idx z) with
-            | Expr (Block (_, a, _)) when n < Array.length a -> Some a.(n)
+            | Expr (Block (_, a, _, _)) when n < Array.length a -> Some a.(n)
             | _ -> None)
         None
         (fun u v ->
@@ -392,6 +397,29 @@ let direct_approx (info : Info.t) x =
         y
   | _ -> None
 
+let rec the_shape_of info x =
+  get_approx
+    info
+    (fun x ->
+      if Var.ISet.mem info.info_possibly_mutable x
+      then Shape.Bot "possibly_mutable"
+      else
+        match info.info_defs.(Var.idx x) with
+        | Expr (Block (_, a, _, true)) ->
+            Shape.Block (List.map ~f:(the_shape_of info) (Array.to_list a))
+        | Expr (Closure (l, _)) ->
+            Shape.Function { arity = List.length l; pure = false; res = Bot "unk" }
+        | Expr (Special (Alias_prim name)) -> (
+            try
+              let arity = Primitive.arity name in
+              let pure = Primitive.is_pure name in
+              Shape.Function { arity; pure; res = Bot "unk" }
+            with _ -> Bot "other")
+        | _ -> Shape.Bot "other")
+    (Bot "init")
+    (fun _u _v -> Shape.Bot "merge")
+    x
+
 let build_subst (info : Info.t) vars =
   let nv = Var.count () in
   let subst = Array.init nv ~f:(fun i -> Var.of_idx i) in

compiler/lib/flow.mli

@@ -60,4 +60,6 @@ val the_native_string_of : Info.t -> Code.prim_arg -> Code.Native_string.t optio
 
 val the_int : Info.t -> Code.prim_arg -> int32 option
 
+val the_shape_of : Info.t -> Code.Var.t -> Shape.t
+
 val f : ?skip_param:bool -> Code.program -> Code.program * Info.t

compiler/lib/freevars.ml

@@ -33,7 +33,7 @@ let iter_expr_free_vars f e =
   | Apply { f = x; args; _ } ->
       f x;
       List.iter ~f args
-  | Block (_, a, _) -> Array.iter ~f a
+  | Block (_, a, _, _) -> Array.iter ~f a
   | Field (x, _) -> f x
   | Closure _ -> ()
   | Special _ -> ()

compiler/lib/generate.ml

@@ -1042,7 +1042,7 @@ let rec translate_expr ctx queue loc x e level : _ * J.statement_list =
       let prop = or_p prop prop' in
       let e = apply_fun ctx f args exact cps loc in
       (e, prop, queue), []
-  | Block (tag, a, array_or_not) ->
+  | Block (tag, a, array_or_not, _imm) ->
       let contents, prop, queue =
         List.fold_right
           ~f:(fun x (args, prop, queue) ->