WIP

compiler/lib/driver.ml

@@ -629,6 +629,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 l = ref StringMap.empty in
   Code.Addr.Map.iter
     (fun _ block ->
@@ -639,7 +640,7 @@ let collects_shapes p =
               , Prim
                   ( Extern "caml_register_global"
                   , [ _code; Pv block; Pc (NativeString name) ] ) ) ->
-              let shape = Flow.the_shape_of info block in
+              let shape = Flow.the_shape_of ~pure info block in
               let name =
                 match name with
                 | Byte s -> s

compiler/lib/flow.ml

@@ -397,7 +397,7 @@ let direct_approx (info : Info.t) x =
         y
   | _ -> None
 
-let rec the_shape_of info x =
+let rec the_shape_of ~pure info x =
   let rec loop info x acc : Shape.t =
     get_approx
       info
@@ -407,9 +407,10 @@ let rec the_shape_of info x =
         | None -> (
             match info.info_defs.(Var.idx x) with
             | Expr (Block (_, a, _, Immutable)) ->
-                Shape.Block (List.map ~f:(the_shape_of info) (Array.to_list a))
+                Shape.Block (List.map ~f:(the_shape_of ~pure info) (Array.to_list a))
             | Expr (Closure (l, _)) ->
-                Shape.Function { arity = List.length l; pure = false; res = Top "unk" }
+                let pure = Code.Var.Set.mem x pure in
+                Shape.Function { arity = List.length l; pure; res = Top "unk" }
             | Expr (Special (Alias_prim name)) -> (
                 try
                   let arity = Primitive.arity name in
@@ -429,7 +430,22 @@ let rec the_shape_of info x =
                   | Shape.Block _ | Shape.Top _ -> Shape.Top "apply2")
             | _ -> Shape.Top "other"))
       (Top "init")
-      (fun _u _v -> Shape.Top "merge")
+      (fun u v ->
+        let rec merge (u : Shape.t) (v : Shape.t) =
+          match u, v with
+          | ( Function { arity = a1; pure = p1; res = r1 }
+            , Function { arity = a2; pure = p2; res = r2 } ) ->
+              if a1 = a2
+              then Shape.Function { arity = a1; pure = p1 && p2; res = merge r1 r2 }
+              else Shape.Top "merge"
+          | Block b1, Block b2 ->
+              if List.length b1 = List.length b2
+              then Block (List.map2 b1 b2 ~f:merge)
+              else Top "merge block"
+          | (Top _ as a), _ | _, (Top _ as a) -> a
+          | Function _, Block _ | Block _, Function _ -> Shape.Top "merge block/fun"
+        in
+        merge u v)
       x
   in
   loop info x []

compiler/lib/flow.mli

@@ -60,6 +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 the_shape_of : 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/pure_fun.ml

@@ -26,7 +26,8 @@ let pure_expr pure_funs e =
   match e with
   | Block _ | Field _ | Closure _ | Constant _ -> true
   | Special (Alias_prim _ | Undefined) -> true
-  | Apply { f; exact; _ } -> exact && Var.Set.mem f pure_funs
+  | Apply { f; exact; _ } ->
+      exact && (Var.Set.mem f pure_funs || Shape.State.is_pure_fun f)
   | Prim (p, _l) -> (
       match p with
       | Extern f -> Primitive.is_pure f

compiler/lib/shape.ml

@@ -46,7 +46,8 @@ let rec to_string (shape : t) =
   match shape with
   | Top s -> if true then "N" else Printf.sprintf "N(%s)" s
   | Block l -> "[" ^ String.concat ~sep:"," (List.map ~f:to_string l) ^ "]"
-  | Function { arity; _ } -> Printf.sprintf "F(%d)" arity
+  | Function { arity; pure; _ } ->
+      Printf.sprintf "F(%d)%s" arity (if pure then "" else "")
 
 module Store = struct
   module T = Hashtbl.Make (struct
@@ -153,5 +154,11 @@ module State = struct
 
   let get x = T.find_opt t x
 
+  let is_pure_fun x =
+    match T.find_opt t x with
+    | None -> false
+    | Some (Top _ | Block _) -> false
+    | Some (Function { pure; _ }) -> pure
+
   let reset () = T.clear t
 end

compiler/lib/shape.mli

@@ -53,5 +53,7 @@ module State : sig
 
   val get : Code.Var.t -> t option
 
+  val is_pure_fun : Code.Var.t -> bool
+
   val reset : unit -> unit
 end

compiler/lib/specialize.ml

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

toplevel/examples/lwt_toplevel/dune

@@ -71,6 +71,7 @@
  (action
   (run
    %{bin:js_of_ocaml}
+   --no-runtime
    --pretty
    --toplevel
    %{read-strings:effects_flags.txt}
@@ -81,6 +82,7 @@
  (action
   (run
    %{bin:js_of_ocaml}
+   --no-runtime
    --pretty
    --toplevel
    %{read-strings:effects_flags.txt}