WIP

compiler/bin-js_of_ocaml/cmd_arg.ml

@@ -54,7 +54,7 @@ type t =
   ; wrap_with_fun : [ `Iife | `Named of string | `Anonymous ]
   ; target_env : Target_env.t
   ; shape_files : string list
-  ; shapes : bool
+  ; write_shape : bool
   ; (* toplevel *)
     dynlink : bool
   ; linkall : bool
@@ -106,11 +106,11 @@ let options =
   in
   let shape_files =
     let doc = "load shape file [$(docv)]." in
-    Arg.(value & opt_all string [] & info [ "load" ] ~docv:"FILE" ~doc)
+    Arg.(value & opt_all string [] & info [ "load-shape" ] ~docv:"FILE" ~doc)
   in
-  let shapes =
+  let write_shape =
     let doc = "Emit shape files" in
-    Arg.(value & flag & info [ "shapes" ] ~doc)
+    Arg.(value & flag & info [ "write-shape" ] ~doc)
   in
   let input_file =
     let doc =
@@ -290,7 +290,7 @@ let options =
       input_file
       js_files
       shape_files
-      shapes
+      write_shape
       keep_unit_names =
     let inline_source_content = not sourcemap_don't_inline_content in
     let chop_extension s = try Filename.chop_extension s with Invalid_argument _ -> s in
@@ -354,7 +354,7 @@ let options =
       ; source_map
       ; keep_unit_names
       ; shape_files
-      ; shapes
+      ; write_shape
       }
   in
   let t =
@@ -386,7 +386,7 @@ let options =
       $ input_file
       $ js_files
       $ shape_files
-      $ shapes
+      $ write_shape
       $ keep_unit_names)
   in
   Term.ret t
@@ -584,7 +584,7 @@ let options_runtime_only =
       ; source_map
       ; keep_unit_names = false
       ; shape_files = []
-      ; shapes = false
+      ; write_shape = false
       }
   in
   let t =

compiler/bin-js_of_ocaml/cmd_arg.mli

@@ -38,7 +38,7 @@ type t =
       ]
   ; target_env : Target_env.t
   ; shape_files : string list
-  ; shapes : bool
+  ; write_shape : bool
   ; (* toplevel *)
     dynlink : bool
   ; linkall : bool

compiler/bin-js_of_ocaml/compile.ml

@@ -47,7 +47,7 @@ let source_map_enabled = function
   | Inline | File _ -> true
 
 let output_gen
-    ~write_shapes
+    ~write_shape
     ~standalone
     ~custom_header
     ~build_info
@@ -60,7 +60,7 @@ let output_gen
     if standalone then header ~custom_header fmt;
     if Config.Flag.header () then jsoo_header fmt build_info;
     let sm, shapes = f ~standalone ~source_map (k, fmt) in
-    (if write_shapes
+    (if write_shape
      then
        match output_file with
        | `Stdout -> ()
@@ -173,7 +173,7 @@ let run
     ; keep_unit_names
     ; include_runtime
     ; shape_files
-    ; shapes = write_shapes
+    ; write_shape
     } =
   let source_map_base = Option.map ~f:snd source_map in
   let source_map =
@@ -388,7 +388,7 @@ let run
         }
       in
       output_gen
-        ~write_shapes
+        ~write_shape
         ~standalone:true
         ~custom_header
         ~build_info:(Build_info.create `Runtime)
@@ -437,7 +437,7 @@ let run
           in
           if times () then Format.eprintf "  parsing: %a@." Timer.print t1;
           output_gen
-            ~write_shapes
+            ~write_shape
             ~standalone:true
             ~custom_header
             ~build_info:(Build_info.create `Exe)
@@ -476,7 +476,7 @@ let run
           in
           if times () then Format.eprintf "  parsing: %a@." Timer.print t1;
           output_gen
-            ~write_shapes
+            ~write_shape
             ~standalone:false
             ~custom_header
             ~build_info:(Build_info.create `Cmo)
@@ -507,7 +507,7 @@ let run
                    failwith "use [-o dirname/] or remove [--keep-unit-names]"
              in
              output_gen
-               ~write_shapes
+               ~write_shape
                ~standalone:false
                ~custom_header
                ~build_info:(Build_info.create `Runtime)
@@ -544,7 +544,7 @@ let run
                   t1
                   (Ocaml_compiler.Cmo_format.name cmo);
               output_gen
-                ~write_shapes
+                ~write_shape
                 ~standalone:false
                 ~custom_header
                 ~build_info:(Build_info.create `Cma)
@@ -594,7 +594,7 @@ let run
             , shapes )
           in
           output_gen
-            ~write_shapes
+            ~write_shape
             ~standalone:false
             ~custom_header
             ~build_info:(Build_info.create `Cma)

compiler/lib/code.ml

@@ -822,6 +822,30 @@ let fold_closures_outermost_first { start; blocks; _ } f accu =
   let accu = f None [] (start, []) accu in
   visit blocks start f accu
 
+(* Compute the list of variables containing the return values of each
+   function *)
+let return_values p =
+  fold_closures
+    p
+    (fun name_opt _ (pc, _) rets ->
+      match name_opt with
+      | None -> rets
+      | Some name ->
+          let s =
+            traverse
+              { fold = fold_children }
+              (fun pc s ->
+                let block = Addr.Map.find pc p.blocks in
+                match block.branch with
+                | Return x -> Var.Set.add x s
+                | _ -> s)
+              pc
+              p.blocks
+              Var.Set.empty
+          in
+          Var.Map.add name s rets)
+    Var.Map.empty
+
 let eq p1 p2 =
   p1.start = p2.start
   && Addr.Map.cardinal p1.blocks = Addr.Map.cardinal p2.blocks

compiler/lib/code.mli

@@ -304,6 +304,8 @@ val fold_children_skip_try_body : 'c fold_blocs
 
 val poptraps : block Addr.Map.t -> Addr.t -> Addr.Set.t
 
+val return_values : program -> Var.Set.t Var.Map.t
+
 val traverse :
   fold_blocs_poly -> (Addr.t -> 'c -> 'c) -> Addr.t -> block Addr.Map.t -> 'c -> 'c
 

compiler/lib/config.ml

@@ -102,8 +102,6 @@ module Flag = struct
 
   let es6 = o ~name:"es6" ~default:false
 
-  let shapes = o ~name:"shapes" ~default:false
-
   let load_shapes_auto = o ~name:"load-shapes-auto" ~default:false
 end
 

compiler/lib/config.mli

@@ -75,8 +75,6 @@ module Flag : sig
 
   val es6 : unit -> bool
 
-  val shapes : unit -> bool
-
   val load_shapes_auto : unit -> bool
 
   val enable : string -> unit

compiler/lib/driver.ml

@@ -65,7 +65,10 @@ let inline p =
 
 let specialize_1 (p, info) =
   if debug () then Format.eprintf "Specialize...@.";
-  Specialize.f ~function_arity:(fun f -> Specialize.function_arity info f) p
+  let return_values = Code.return_values p in
+  Specialize.f
+    ~function_arity:(fun f -> Specialize.function_arity ~return_values info f)
+    p
 
 let specialize_js (p, info) =
   if debug () then Format.eprintf "Specialize js...@.";
@@ -673,6 +676,7 @@ if (typeof module === 'object' && module.exports) {
 let collects_shapes p =
   let _, info = Flow.f p in
   let pure = Pure_fun.f p in
+  let return_values = Code.return_values p in
   let l = ref StringMap.empty in
   Code.Addr.Map.iter
     (fun _ block ->
@@ -683,7 +687,7 @@ let collects_shapes p =
               , Prim
                   ( Extern "caml_register_global"
                   , [ _code; Pv block; Pc (NativeString name) ] ) ) ->
-              let shape = Flow.the_shape_of ~pure info block in
+              let shape = Flow.the_shape_of ~return_values ~pure info block in
               let name =
                 match name with
                 | Byte s -> s

compiler/lib/flow.ml

@@ -40,7 +40,6 @@ module Info = struct
     ; info_known_origins : Code.Var.Set.t Code.Var.Tbl.t
     ; info_maybe_unknown : bool Code.Var.Tbl.t
     ; info_possibly_mutable : Var.ISet.t
-    ; info_blocks : Code.block Addr.Map.t
     }
 
   let def t x =
@@ -431,7 +430,7 @@ let direct_approx (info : Info.t) x =
         y
   | _ -> None
 
-let rec the_shape_of ~pure info x =
+let rec the_shape_of ~return_values ~pure info x =
   let rec merge (u : Shape.t) (v : Shape.t) =
     match u, v with
     | ( Function { arity = a1; pure = p1; res = r1 }
@@ -455,30 +454,28 @@ let rec the_shape_of ~pure info x =
         | None -> (
             match info.info_defs.(Var.idx x) with
             | Expr (Block (_, a, _, Immutable)) ->
-                Shape.Block (List.map ~f:(the_shape_of ~pure info) (Array.to_list a))
-            | Expr (Closure (l, (pc, _))) ->
+                Shape.Block
+                  (List.map ~f:(the_shape_of ~return_values ~pure info) (Array.to_list a))
+            | Expr (Closure (l, _)) ->
                 let pure = Code.Var.Set.mem x pure in
-                let blocks = info.info_blocks in
                 let res =
-                  Code.traverse
-                    { fold = fold_children }
-                    (fun pc res ->
-                      let block = Addr.Map.find pc blocks in
-                      match block.branch with
-                      | Return x -> (
-                          let s2 = loop info x acc in
-                          match res with
-                          | None -> Some s2
-                          | Some s1 -> Some (merge s1 s2))
-                      | _ -> res)
-                    pc
-                    blocks
-                    None
-                in
-                let res : Shape.t =
-                  match res with
-                  | None -> Top "no return"
-                  | Some res -> res
+                  match Var.Map.find x return_values with
+                  | exception Not_found -> Shape.Top "not return_values found"
+                  | set -> (
+                      match
+                        Var.Set.fold
+                          (fun x res ->
+                            let s2 = loop info x acc in
+                            match res with
+                            | None -> Some s2
+                            | Some s1 -> Some (merge s1 s2))
+                          set
+                          None
+                      with
+                      | None ->
+                          assert (Var.Set.is_empty set);
+                          Shape.Top "no return"
+                      | Some res -> res)
                 in
                 Shape.Function { arity = List.length l; pure; res }
             | Expr (Special (Alias_prim name)) -> (
@@ -565,7 +562,6 @@ let f ?skip_param p =
     ; info_known_origins = known_origins
     ; info_maybe_unknown = maybe_unknown
     ; info_possibly_mutable = possibly_mutable
-    ; info_blocks = p.blocks
     }
   in
   let s = build_subst info vars in

compiler/lib/flow.mli

@@ -64,6 +64,11 @@ val the_native_string_of :
 val the_int :
   target:[ `JavaScript | `Wasm ] -> Info.t -> Code.prim_arg -> Targetint.t option
 
-val the_shape_of : pure:Code.Var.Set.t -> Info.t -> Code.Var.t -> Shape.t
+val the_shape_of :
+     return_values:Code.Var.Set.t Code.Var.Map.t
+  -> pure:Code.Var.Set.t
+  -> Info.t
+  -> Code.Var.t
+  -> Shape.t
 
 val f : ?skip_param:bool -> Code.program -> Code.program * Info.t

compiler/lib/global_flow.ml

@@ -33,32 +33,6 @@ open Code
 
 (****)
 
-(* Compute the list of variables containing the return values of each
-   function *)
-let return_values p =
-  Code.fold_closures
-    p
-    (fun name_opt _ (pc, _) rets ->
-      match name_opt with
-      | None -> rets
-      | Some name ->
-          let s =
-            Code.traverse
-              { fold = fold_children }
-              (fun pc s ->
-                let block = Addr.Map.find pc p.blocks in
-                match block.branch with
-                | Return x -> Var.Set.add x s
-                | _ -> s)
-              pc
-              p.blocks
-              Var.Set.empty
-          in
-          Var.Map.add name s rets)
-    Var.Map.empty
-
-(****)
-
 (* A variable is either let-bound, or a parameter, to which we
    associate a set of possible arguments.
 *)

compiler/lib/specialize.ml

@@ -20,8 +20,8 @@
 open! Stdlib
 open Code
 
-let function_arity info x =
-  match Flow.the_shape_of ~pure:Code.Var.Set.empty info x with
+let function_arity ~return_values info x =
+  match Flow.the_shape_of ~return_values ~pure:Code.Var.Set.empty info x with
   | Top _ | Block _ -> None
   | Function { arity; _ } -> Some arity
 

compiler/lib/specialize.mli

@@ -18,6 +18,7 @@
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  *)
 
-val function_arity : Flow.Info.t -> Code.Var.t -> int option
+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