Compiler: no longer rely on iife

compiler/lib/driver.ml

@@ -603,7 +603,8 @@ let full ~standalone ~wrap_with_fun ~profile ~linkall ~source_map formatter d p
        | O3 -> o3)
     +> exact_calls ~deadcode_sentinal profile
     +> effects ~deadcode_sentinal
-    +> map_fst (Generate_closure.f +> deadcode')
+    +> map_fst (if Config.Flag.effects () then fun x -> x else Generate_closure.f)
+    +> map_fst deadcode'
   in
   let emit =
     generate

compiler/lib/generate.ml

@@ -281,13 +281,15 @@ module Ctx = struct
     ; effect_warning : bool ref
     ; cps_calls : Effects.cps_calls
     ; deadcode_sentinal : Var.t
+    ; mutated_vars : Code.Var.Set.t Code.Addr.Map.t
     }
 
   let initial
       ~warn_on_unhandled_effect
       ~exported_runtime
       ~should_export
       ~deadcode_sentinal
+      ~mutated_vars
       blocks
       live
       cps_calls
@@ -302,6 +304,7 @@ module Ctx = struct
     ; effect_warning = ref (not warn_on_unhandled_effect)
     ; cps_calls
     ; deadcode_sentinal
+    ; mutated_vars
     }
 end
 
@@ -1383,9 +1386,73 @@ and translate_instr ctx expr_queue instr =
         mutator_p
         [ J.Expression_statement (J.EBin (J.Eq, Mlvalue.Array.field cx cy, cz)), loc ]
 
-and translate_instrs ctx expr_queue instr last =
+and translate_instrs (ctx : Ctx.t) expr_queue instr last =
   match instr with
   | [] -> [], expr_queue
+  | (Let (_, Closure _), _) :: _ -> (
+      let names, mut, pcs, all, rem = collect_closures ctx instr in
+      match Code.Var.Set.cardinal mut with
+      | 0 ->
+          let st, expr_queue =
+            List.fold_left all ~init:([], expr_queue) ~f:(fun (st, expr_queue) i ->
+                let l, expr_queue = translate_instr ctx expr_queue i in
+                List.rev_append l st, expr_queue)
+          in
+          let instrs, expr_queue = translate_instrs ctx expr_queue rem last in
+          List.rev_append st instrs, expr_queue
+      | _ ->
+          let muts =
+            Code.Var.Set.diff mut names
+            |> Code.Var.Set.elements
+            |> List.map ~f:(fun x -> x, Code.Var.fork x)
+          in
+          (* Rewrite blocks using well-scoped closure variables *)
+          let ctx =
+            let map =
+              List.fold_left muts ~init:Var.Map.empty ~f:(fun acc (x, x') ->
+                  Var.Map.add x x' acc)
+            in
+            let p, _visited =
+              List.fold_left
+                pcs
+                ~init:(ctx.blocks, Addr.Set.empty)
+                ~f:(fun (blocks, visited) pc ->
+                  Subst.cont' (Subst.from_map map) pc blocks visited)
+            in
+            { ctx with blocks = p }
+          in
+          (* Let bind mutable variables that are part of closures *)
+          let let_bindings_rev, expr_queue =
+            let expr_queue, st_rev, l_rev =
+              List.fold_left
+                muts
+                ~init:(expr_queue, [], [])
+                ~f:(fun (expr_queue, st_rev, l_rev) (v, v') ->
+                  let instrs, ((_px, cx), expr_queue) =
+                    access_queue_may_flush expr_queue v' v
+                  in
+                  let l_rev = (J.V v', (cx, J.N)) :: l_rev in
+                  expr_queue, List.rev_append instrs st_rev, l_rev)
+            in
+            (J.variable_declaration ~kind:Let (List.rev l_rev), J.N) :: st_rev, expr_queue
+          in
+          (* Mutually recursive need to be properly scoped. *)
+          let st_rev, expr_queue =
+            List.fold_left all ~init:([], expr_queue) ~f:(fun (st_rev, expr_queue) i ->
+                let l, expr_queue = translate_instr ctx expr_queue i in
+                let l_rev =
+                  List.rev_map l ~f:(fun (e, loc') ->
+                      match e with
+                      (* FIXME: This pattern is too fragile *)
+                      | J.Variable_statement
+                          (Var, [ DeclIdent (x, Some (J.EFun (None, dcl), _loc)) ]) ->
+                          J.Function_declaration (x, dcl), loc'
+                      | _ -> e, loc')
+                in
+                List.rev_append l_rev st_rev, expr_queue)
+          in
+          let instrs, expr_queue = translate_instrs ctx expr_queue rem last in
+          List.rev_append let_bindings_rev (List.rev_append st_rev instrs), expr_queue)
   | instr :: rem ->
       let st, expr_queue = translate_instr ctx expr_queue instr in
       let instrs, expr_queue = translate_instrs ctx expr_queue rem last in
@@ -1750,6 +1817,14 @@ and compile_closure ctx (pc, args) =
   if debug () then Format.eprintf "}@]@;";
   res
 
+and collect_closures ctx l =
+  match l with
+  | ((Let (x, Closure (_, (pc, _))), _loc) as i) :: rem ->
+      let names', mut', pcs', i', rem' = collect_closures ctx rem in
+      let mut = Code.Addr.Map.find pc ctx.Ctx.mutated_vars in
+      Code.Var.Set.add x names', Code.Var.Set.union mut mut', pc :: pcs', i :: i', rem'
+  | _ -> Code.Var.Set.empty, Code.Var.Set.empty, [], [], l
+
 let generate_shared_value ctx =
   let strings =
     ( J.variable_declaration
@@ -1808,12 +1883,14 @@ let f
   let exported_runtime =
     if exported_runtime then Some (Code.Var.fresh_n "runtime", ref false) else None
   in
+  let mutated_vars = Freevars.f p in
   let ctx =
     Ctx.initial
       ~warn_on_unhandled_effect
       ~exported_runtime
       ~should_export
       ~deadcode_sentinal
+      ~mutated_vars
       p.blocks
       live_vars
       cps_calls

compiler/lib/generate_closure.ml

@@ -27,7 +27,6 @@ type closure_info =
   ; args : Code.Var.t list
   ; cont : Code.cont
   ; tc : Code.Addr.Set.t Code.Var.Map.t
-  ; mutated_vars : Code.Var.Set.t
   ; loc : Code.loc
   }
 
@@ -67,24 +66,22 @@ let rec collect_apply pc blocks visited tc =
           (fun pc (visited, tc) -> collect_apply pc blocks visited tc)
           (visited, tc)
 
-let rec collect_closures blocks mutated_vars l =
+let rec collect_closures blocks l =
   match l with
   | (Let (f_name, Closure (args, ((pc, _) as cont))), loc) :: rem ->
       let _, tc = collect_apply pc blocks Addr.Set.empty Var.Map.empty in
-      let l, rem = collect_closures blocks mutated_vars rem in
-      let mutated_vars = Addr.Map.find pc mutated_vars in
-      { f_name; args; cont; tc; mutated_vars; loc } :: l, rem
+      let l, rem = collect_closures blocks rem in
+      { f_name; args; cont; tc; loc } :: l, rem
   | rem -> [], rem
 
-let group_closures ~tc_only closures_map =
+let group_closures closures_map =
   let names =
     Var.Map.fold (fun _ x names -> Var.Set.add x.f_name names) closures_map Var.Set.empty
   in
   let graph =
     Var.Map.fold
       (fun _ x graph ->
         let calls = Var.Map.fold (fun x _ tc -> Var.Set.add x tc) x.tc Var.Set.empty in
-        let calls = if tc_only then calls else Var.Set.union calls x.mutated_vars in
         Var.Map.add x.f_name (Var.Set.inter names calls) graph)
       closures_map
       Var.Map.empty
@@ -278,8 +275,6 @@ end
 let rewrite_tc free_pc blocks closures_map component =
   let open Config.Param in
   let trampoline =
-    (not (Config.Flag.effects ()))
-    &&
     match tailcall_optim () with
     | TcTrampoline -> true
     | TcNone -> false
@@ -288,189 +283,52 @@ let rewrite_tc free_pc blocks closures_map component =
   then Trampoline.f free_pc blocks closures_map component
   else Ident.f free_pc blocks closures_map component
 
-let rewrite_mutable
-    free_pc
-    blocks
-    mutated_vars
-    rewrite_list
-    { int = closures_intern; ext = closures_extern } =
-  let internal_and_external = closures_intern @ closures_extern in
-  assert (not (List.is_empty closures_extern));
-  let all_mut, names =
-    List.fold_left
-      internal_and_external
-      ~init:(Var.Set.empty, Var.Set.empty)
-      ~f:(fun (all_mut, names) i ->
-        match i with
-        | Let (x, Closure (_, (pc, _))), _ ->
-            let all_mut =
-              try Var.Set.union all_mut (Addr.Map.find pc mutated_vars)
-              with Not_found -> all_mut
-            in
-            let names = Var.Set.add x names in
-            all_mut, names
-        | _ -> assert false)
-  in
-  let vars = Var.Set.elements (Var.Set.diff all_mut names) in
-  if List.is_empty vars
-  then free_pc, blocks, internal_and_external
-  else
-    match internal_and_external with
-    | [ (Let (x, Closure (params, (pc, pc_args))), loc) ] ->
-        let new_pc = free_pc in
-        let free_pc = free_pc + 1 in
-        let closure = Code.Var.fork x in
-        let args = List.map vars ~f:Code.Var.fork in
-        let new_x = Code.Var.fork x in
-        let mapping = Subst.from_map (Subst.build_mapping (x :: vars) (new_x :: args)) in
-        rewrite_list := (mapping, pc) :: !rewrite_list;
-        let new_block =
-          { params = []
-          ; body =
-              [ Let (new_x, Closure (params, (pc, List.map pc_args ~f:mapping))), loc ]
-          ; branch = Return new_x, loc
-          }
-        in
-        let blocks = Addr.Map.add new_pc new_block blocks in
-        let body =
-          [ Let (closure, Closure (args, (new_pc, []))), noloc
-          ; Let (x, Apply { f = closure; args = vars; exact = true }), loc
-          ]
-        in
-        free_pc, blocks, body
-    | _ ->
-        let new_pc = free_pc in
-        let free_pc = free_pc + 1 in
-        let closure = Code.Var.fresh_n "closures" in
-        let closure' = Code.Var.fresh_n "closures" in
-        let b = Code.Var.fresh_n "block" in
-        let args = List.map vars ~f:Code.Var.fork in
-        let pcs =
-          List.map internal_and_external ~f:(function
-              | Let (_, Closure (_, (pc, _))), _ -> pc
-              | _ -> assert false)
-        in
-        let old_xs =
-          List.map closures_extern ~f:(function
-              | Let (x, Closure _), _ -> x
-              | _ -> assert false)
-        in
-        let new_xs = List.map old_xs ~f:Code.Var.fork in
-        let mapping =
-          Subst.from_map (Subst.build_mapping (old_xs @ vars) (new_xs @ args))
-        in
-        rewrite_list := List.map pcs ~f:(fun pc -> mapping, pc) @ !rewrite_list;
-        let new_block =
-          let proj =
-            List.map2 closures_extern new_xs ~f:(fun cl new_x ->
-                match cl with
-                | Let (_, Closure (params, (pc, pc_args))), loc ->
-                    Let (new_x, Closure (params, (pc, List.map pc_args ~f:mapping))), loc
-                | _ -> assert false)
-          in
-          { params = []
-          ; body =
-              closures_intern
-              @ proj
-              @ [ Let (b, Block (0, Array.of_list new_xs, NotArray)), noloc ]
-          ; branch = Return b, noloc
-          }
-        in
-        let blocks = Addr.Map.add new_pc new_block blocks in
-        let body =
-          [ Let (closure, Closure (args, (new_pc, []))), noloc
-          ; Let (closure', Apply { f = closure; args = vars; exact = true }), noloc
-          ]
-          @ List.mapi closures_extern ~f:(fun i x ->
-                match x with
-                | Let (x, Closure _), loc -> Let (x, Field (closure', i)), loc
-                | _ -> assert false)
-        in
-        free_pc, blocks, body
-
-let rec rewrite_closures mutated_vars rewrite_list free_pc blocks body : int * _ * _ list
-    =
+let rec rewrite_closures free_pc blocks body : int * _ * _ list =
   match body with
   | (Let (_, Closure _), _) :: _ ->
-      let closures, rem = collect_closures blocks mutated_vars body in
+      let closures, rem = collect_closures blocks body in
       let closures_map =
         List.fold_left closures ~init:Var.Map.empty ~f:(fun closures_map x ->
             Var.Map.add x.f_name x closures_map)
       in
-      let components = group_closures ~tc_only:false closures_map in
+      let components = group_closures closures_map in
       let free_pc, blocks, closures =
         List.fold_left
           (Array.to_list components)
           ~init:(free_pc, blocks, [])
           ~f:(fun (free_pc, blocks, acc) component ->
             let free_pc, blocks, closures =
-              let components =
-                match component with
-                | SCC.No_loop _ as one -> [ one ]
-                | SCC.Has_loop all ->
-                    group_closures
-                      ~tc_only:true
-                      (Var.Map.filter
-                         (fun v _ -> List.exists all ~f:(Var.equal v))
-                         closures_map)
-                    |> Array.to_list
-              in
-              List.fold_left
-                ~init:(free_pc, blocks, { int = []; ext = [] })
-                components
-                ~f:(fun (free_pc, blocks, acc) component ->
-                  let free_pc, blocks, ie =
-                    rewrite_tc free_pc blocks closures_map component
-                  in
-                  free_pc, blocks, { int = ie.int :: acc.int; ext = ie.ext :: acc.ext })
-            in
-            let closures =
-              { int = List.concat (List.rev closures.int)
-              ; ext = List.concat (List.rev closures.ext)
-              }
-            in
-            let free_pc, blocks, intrs =
-              rewrite_mutable free_pc blocks mutated_vars rewrite_list closures
+              rewrite_tc free_pc blocks closures_map component
             in
-            free_pc, blocks, intrs :: acc)
-      in
-      let free_pc, blocks, rem =
-        rewrite_closures mutated_vars rewrite_list free_pc blocks rem
+            let intrs = closures.int :: closures.ext :: acc in
+            free_pc, blocks, intrs)
       in
+      let free_pc, blocks, rem = rewrite_closures free_pc blocks rem in
       free_pc, blocks, List.flatten closures @ rem
   | i :: rem ->
-      let free_pc, blocks, rem =
-        rewrite_closures mutated_vars rewrite_list free_pc blocks rem
-      in
+      let free_pc, blocks, rem = rewrite_closures free_pc blocks rem in
       free_pc, blocks, i :: rem
   | [] -> free_pc, blocks, []
 
 let f p : Code.program =
   Code.invariant p;
-  let mutated_vars = Freevars.f p in
-  let rewrite_list = ref [] in
   let blocks, free_pc =
     Addr.Map.fold
       (fun pc _ (blocks, free_pc) ->
         (* make sure we have the latest version *)
         let block = Addr.Map.find pc blocks in
-        let free_pc, blocks, body =
-          rewrite_closures mutated_vars rewrite_list free_pc blocks block.body
-        in
+        let free_pc, blocks, body = rewrite_closures free_pc blocks block.body in
         Addr.Map.add pc { block with body } blocks, free_pc)
       p.blocks
       (p.blocks, p.free_pc)
   in
   (* Code.invariant (pc, blocks, free_pc); *)
   let p = { p with blocks; free_pc } in
-  let p =
-    List.fold_left !rewrite_list ~init:p ~f:(fun program (mapping, pc) ->
-        Subst.cont mapping pc program)
-  in
   Code.invariant p;
   p
 
 let f p =
+  assert (not (Config.Flag.effects ()));
   let t = Timer.make () in
   let p' = f p in
   if Debug.find "times" () then Format.eprintf "  generate closures: %a@." Timer.print t;

compiler/lib/javascript.ml

@@ -540,8 +540,8 @@ end)
 
 let dot e l = EDot (e, ANormal, l)
 
-let variable_declaration l =
-  Variable_statement (Var, List.map l ~f:(fun (i, e) -> DeclIdent (i, Some e)))
+let variable_declaration ?(kind = Var) l =
+  Variable_statement (kind, List.map l ~f:(fun (i, e) -> DeclIdent (i, Some e)))
 
 let array l = EArr (List.map l ~f:(fun x -> Element x))