Small changes in the internal APIs for Wasm

compiler/lib/code.ml

@@ -680,6 +680,39 @@ let is_empty p =
       | _ -> false)
   | _ -> false
 
+let poptraps blocks pc =
+  let rec loop blocks pc visited depth acc =
+    if Addr.Set.mem pc visited
+    then acc, visited
+    else
+      let visited = Addr.Set.add pc visited in
+      let block = Addr.Map.find pc blocks in
+      match fst block.branch with
+      | Return _ | Raise _ | Stop -> acc, visited
+      | Branch (pc', _) -> loop blocks pc' visited depth acc
+      | Poptrap (pc', _) ->
+          if depth = 0
+          then Addr.Set.add pc' acc, visited
+          else loop blocks pc' visited (depth - 1) acc
+      | Pushtrap ((pc', _), _, (pc_h, _)) ->
+          let acc, visited = loop blocks pc' visited (depth + 1) acc in
+          let acc, visited = loop blocks pc_h visited depth acc in
+          acc, visited
+      | Cond (_, (pc1, _), (pc2, _)) ->
+          let acc, visited = loop blocks pc1 visited depth acc in
+          let acc, visited = loop blocks pc2 visited depth acc in
+          acc, visited
+      | Switch (_, a) ->
+          let acc, visited =
+            Array.fold_right
+              ~init:(acc, visited)
+              ~f:(fun (pc, _) (acc, visited) -> loop blocks pc visited depth acc)
+              a
+          in
+          acc, visited
+  in
+  loop blocks pc Addr.Set.empty 0 Addr.Set.empty |> fst
+
 let fold_children blocks pc f accu =
   let block = Addr.Map.find pc blocks in
   match fst block.branch with
@@ -697,6 +730,23 @@ let fold_children blocks pc f accu =
       let accu = Array.fold_right ~init:accu ~f:(fun (pc, _) accu -> f pc accu) a1 in
       accu
 
+let fold_children_skip_try_body blocks pc f accu =
+  let block = Addr.Map.find pc blocks in
+  match fst block.branch with
+  | Return _ | Raise _ | Stop -> accu
+  | Branch (pc', _) | Poptrap (pc', _) -> f pc' accu
+  | Pushtrap ((pc', _), _, (pc_h, _)) ->
+      let accu = Addr.Set.fold f (poptraps blocks pc') accu in
+      let accu = f pc_h accu in
+      accu
+  | Cond (_, (pc1, _), (pc2, _)) ->
+      let accu = f pc1 accu in
+      let accu = f pc2 accu in
+      accu
+  | Switch (_, a1) ->
+      let accu = Array.fold_right ~init:accu ~f:(fun (pc, _) accu -> f pc accu) a1 in
+      accu
+
 type 'c fold_blocs = block Addr.Map.t -> Addr.t -> (Addr.t -> 'c -> 'c) -> 'c -> 'c
 
 type fold_blocs_poly = { fold : 'a. 'a fold_blocs } [@@unboxed]
@@ -720,39 +770,6 @@ let rec traverse' { fold } f pc visited blocks acc =
 
 let traverse fold f pc blocks acc = snd (traverse' fold f pc Addr.Set.empty blocks acc)
 
-let poptraps blocks pc =
-  let rec loop blocks pc visited depth acc =
-    if Addr.Set.mem pc visited
-    then acc, visited
-    else
-      let visited = Addr.Set.add pc visited in
-      let block = Addr.Map.find pc blocks in
-      match fst block.branch with
-      | Return _ | Raise _ | Stop -> acc, visited
-      | Branch (pc', _) -> loop blocks pc' visited depth acc
-      | Poptrap (pc', _) ->
-          if depth = 0
-          then Addr.Set.add pc' acc, visited
-          else loop blocks pc' visited (depth - 1) acc
-      | Pushtrap ((pc', _), _, (pc_h, _)) ->
-          let acc, visited = loop blocks pc' visited (depth + 1) acc in
-          let acc, visited = loop blocks pc_h visited depth acc in
-          acc, visited
-      | Cond (_, (pc1, _), (pc2, _)) ->
-          let acc, visited = loop blocks pc1 visited depth acc in
-          let acc, visited = loop blocks pc2 visited depth acc in
-          acc, visited
-      | Switch (_, a) ->
-          let acc, visited =
-            Array.fold_right
-              ~init:(acc, visited)
-              ~f:(fun (pc, _) (acc, visited) -> loop blocks pc visited depth acc)
-              a
-          in
-          acc, visited
-  in
-  loop blocks pc Addr.Set.empty 0 Addr.Set.empty |> fst
-
 let rec preorder_traverse' { fold } f pc visited blocks acc =
   if not (Addr.Set.mem pc visited)
   then
@@ -789,6 +806,25 @@ let fold_closures_innermost_first { start; blocks; _ } f accu =
   let accu = visit blocks start f accu in
   f None [] (start, []) accu
 
+let fold_closures_outermost_first { start; blocks; _ } f accu =
+  let rec visit blocks pc f accu =
+    traverse
+      { fold = fold_children }
+      (fun pc accu ->
+        let block = Addr.Map.find pc blocks in
+        List.fold_left block.body ~init:accu ~f:(fun accu i ->
+            match i with
+            | Let (x, Closure (params, cont)), _ ->
+                let accu = f (Some x) params cont accu in
+                visit blocks (fst cont) f accu
+            | _ -> accu))
+      pc
+      blocks
+      accu
+  in
+  let accu = f None [] (start, []) accu in
+  visit blocks start f accu
+
 let eq p1 p2 =
   p1.start = p2.start
   && Addr.Map.cardinal p1.blocks = Addr.Map.cardinal p2.blocks

compiler/lib/code.mli

@@ -293,8 +293,16 @@ val fold_closures_innermost_first :
     innermost closures first. Unlike with {!fold_closures}, only the closures
     reachable from [p.start] are considered. *)
 
+val fold_closures_outermost_first :
+  program -> (Var.t option -> Var.t list -> cont -> 'd -> 'd) -> 'd -> 'd
+(** Similar to {!fold_closures}, but applies the fold function to the
+    outermost closures first. Unlike with {!fold_closures}, only the closures
+    reachable from [p.start] are considered. *)
+
 val fold_children : 'c fold_blocs
 
+val fold_children_skip_try_body : 'c fold_blocs
+
 val poptraps : block Addr.Map.t -> Addr.t -> Addr.Set.t
 
 val traverse :

compiler/lib/freevars.ml

@@ -100,8 +100,8 @@ type st =
   ; mutable revisited : bool
   }
 
-let find_loops p =
-  let in_loop = ref Addr.Map.empty in
+let find_loops p in_loop pc =
+  let in_loop = ref in_loop in
   let index = ref 0 in
   let state = ref Addr.Map.empty in
   let stack = Stack.create () in
@@ -141,9 +141,17 @@ let find_loops p =
       if st.revisited
       then List.iter !l ~f:(fun pc' -> in_loop := Addr.Map.add pc' pc !in_loop))
   in
-  Code.fold_closures p (fun _ _ (pc, _) () -> traverse pc) ();
+  traverse pc;
   !in_loop
 
+let find_loops_in_closure p pc = find_loops p Addr.Map.empty pc
+
+let find_all_loops p =
+  Code.fold_closures
+    p
+    (fun _ _ (pc, _) (in_loop : _ Addr.Map.t) -> find_loops p in_loop pc)
+    Addr.Map.empty
+
 let mark_variables in_loop p =
   let vars = Var.Tbl.make () (-1) in
   let visited = BitSet.create' p.free_pc in
@@ -245,7 +253,7 @@ let f p =
 let f_mutable p =
   Code.invariant p;
   let t = Timer.make () in
-  let in_loop = find_loops p in
+  let in_loop = find_all_loops p in
   let vars = mark_variables in_loop p in
   let free_vars = free_variables vars in_loop p in
   if times () then Format.eprintf "  free vars 1: %a@." Timer.print t;

compiler/lib/freevars.mli

@@ -23,6 +23,12 @@ val iter_block_free_vars : (Code.Var.t -> unit) -> Code.block -> unit
 
 val iter_block_bound_vars : (Code.Var.t -> unit) -> Code.block -> unit
 
+val iter_instr_free_vars : (Code.Var.t -> unit) -> Code.instr -> unit
+
+val iter_last_free_var : (Code.Var.t -> unit) -> Code.last -> unit
+
+val find_loops_in_closure : Code.program -> Code.Addr.t -> Code.Addr.t Code.Addr.Map.t
+
 val f_mutable : Code.program -> Code.Var.Set.t Code.Addr.Map.t
 
 val f : Code.program -> Code.Var.Set.t Code.Addr.Map.t

compiler/lib/fs.ml

@@ -48,3 +48,29 @@ let read_file f =
     Bytes.unsafe_to_string s
   with e ->
     failwith (Printf.sprintf "Cannot read content of %s.\n%s" f (Printexc.to_string e))
+
+let write_file ~name ~contents =
+  let ch = open_out_bin name in
+  output_string ch contents;
+  close_out ch
+
+let remove_file file = try Sys.remove file with Sys_error _ -> ()
+
+let gen_file file f =
+  let f_tmp =
+    Filename.temp_file_name
+      ~temp_dir:(Filename.dirname file)
+      (Filename.basename file)
+      ".tmp"
+  in
+  try
+    let res = f f_tmp in
+    remove_file file;
+    Sys.rename f_tmp file;
+    res
+  with exc ->
+    remove_file f_tmp;
+    raise exc
+
+let with_intermediate_file name f =
+  Fun.protect ~finally:(fun () -> remove_file name) (fun () -> f name)

compiler/lib/fs.mli

@@ -21,3 +21,9 @@ val find_in_path : string list -> string -> string option
 val absolute_path : string -> string
 
 val read_file : string -> string
+
+val write_file : name:string -> contents:string -> unit
+
+val gen_file : string -> (string -> 'a) -> 'a
+
+val with_intermediate_file : string -> (string -> 'a) -> 'a

compiler/lib/generate.ml

@@ -345,11 +345,14 @@ let bool e = J.ECond (e, one, zero)
 
 (****)
 
-let source_location ctx ?force (pc : Code.loc) =
-  match Parse_bytecode.Debug.find_loc ctx.Ctx.debug ?force pc with
+let source_location debug ?force (pc : Code.loc) =
+  match Parse_bytecode.Debug.find_loc debug ?force pc with
   | Some pi -> J.Pi pi
   | None -> J.N
 
+let source_location_ctx ctx ?force (pc : Code.loc) =
+  source_location ctx.Ctx.debug ?force pc
+
 (****)
 
 let float_const f = J.ENum (J.Num.of_float f)
@@ -1069,14 +1072,14 @@ let rec translate_expr ctx queue loc x e level : _ * J.statement_list =
       let (px, cx), queue = access_queue queue x in
       (Mlvalue.Block.field cx n, or_p px mutable_p, queue), []
   | Closure (args, ((pc, _) as cont)) ->
-      let loc = source_location ctx ~force:After (After pc) in
+      let loc = source_location_ctx ctx ~force:After (After pc) in
       let fv = Addr.Map.find pc ctx.freevars in
       let clo = compile_closure ctx cont in
       let clo =
         match clo with
         | (st, x) :: rem ->
             let loc =
-              match x, source_location ctx (Before pc) with
+              match x, source_location_ctx ctx (Before pc) with
               | (J.U | J.N), (J.U | J.N) -> J.U
               | x, (J.U | J.N) -> x
               | (J.U | J.N), x -> x
@@ -1341,14 +1344,14 @@ and translate_instr ctx expr_queue instr =
   let instr, pc = instr in
   match instr with
   | Assign (x, y) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       let (_py, cy), expr_queue = access_queue expr_queue y in
       flush_queue
         expr_queue
         mutator_p
         [ J.Expression_statement (J.EBin (J.Eq, J.EVar (J.V x), cy)), loc ]
   | Let (x, e) -> (
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       let (ce, prop, expr_queue), instrs = translate_expr ctx expr_queue loc x e 0 in
       let keep_name x =
         match Code.Var.get_name x with
@@ -1374,15 +1377,15 @@ and translate_instr ctx expr_queue instr =
             prop
             (instrs @ [ J.variable_declaration [ J.V x, (ce, loc) ], loc ]))
   | Set_field (x, n, y) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       let (_px, cx), expr_queue = access_queue expr_queue x in
       let (_py, cy), expr_queue = access_queue expr_queue y in
       flush_queue
         expr_queue
         mutator_p
         [ J.Expression_statement (J.EBin (J.Eq, Mlvalue.Block.field cx n, cy)), loc ]
   | Offset_ref (x, n) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       (* FIX: may overflow.. *)
       let (_px, cx), expr_queue = access_queue expr_queue x in
       let expr = Mlvalue.Block.field cx 0 in
@@ -1395,7 +1398,7 @@ and translate_instr ctx expr_queue instr =
       in
       flush_queue expr_queue mutator_p [ J.Expression_statement expr', loc ]
   | Array_set (x, y, z) ->
-      let loc = source_location ctx pc in
+      let loc = source_location_ctx ctx pc in
       let (_px, cx), expr_queue = access_queue expr_queue x in
       let (_py, cy), expr_queue = access_queue expr_queue y in
       let (_pz, cz), expr_queue = access_queue expr_queue z in
@@ -1557,7 +1560,7 @@ and compile_block st queue (pc : Addr.t) scope_stack ~fall_through =
           if debug () then Format.eprintf "}@]@,";
           let for_loop =
             ( J.For_statement (J.Left None, None, None, Js_simpl.block body)
-            , source_location st.ctx (Code.location_of_pc pc) )
+            , source_location_ctx st.ctx (Code.location_of_pc pc) )
           in
           let label = if !lab_used then Some lab else None in
           let for_loop =
@@ -1720,7 +1723,7 @@ and compile_conditional st queue ~fall_through last scope_stack : _ * _ =
      | Stop -> Format.eprintf "stop;@;"
      | Cond (x, _, _) -> Format.eprintf "@[<hv 2>cond(%a){@;" Code.Var.print x
      | Switch (x, _) -> Format.eprintf "@[<hv 2>switch(%a){@;" Code.Var.print x);
-  let loc = source_location st.ctx pc in
+  let loc = source_location_ctx st.ctx pc in
   let res =
     match last with
     | Return x ->

compiler/lib/generate.mli

@@ -30,3 +30,9 @@ val f :
   -> Javascript.program
 
 val init : unit -> unit
+
+val source_location :
+     Parse_bytecode.Debug.t
+  -> ?force:Parse_bytecode.Debug.force
+  -> Code.loc
+  -> Javascript.location

compiler/lib/inline.ml

@@ -151,24 +151,13 @@ let rewrite_block pc' pc blocks =
   in
   Addr.Map.add pc block blocks
 
-(* Skip try body *)
-let fold_children blocks pc f accu =
-  let block = Addr.Map.find pc blocks in
-  match fst block.branch with
-  | Return _ | Raise _ | Stop -> accu
-  | Branch (pc', _) | Poptrap (pc', _) -> f pc' accu
-  | Pushtrap ((try_body, _), _, (pc1, _)) ->
-      f pc1 (Addr.Set.fold f (Code.poptraps blocks try_body) accu)
-  | Cond (_, (pc1, _), (pc2, _)) ->
-      let accu = f pc1 accu in
-      let accu = f pc2 accu in
-      accu
-  | Switch (_, a1) ->
-      let accu = Array.fold_right a1 ~init:accu ~f:(fun (pc, _) accu -> f pc accu) in
-      accu
-
 let rewrite_closure blocks cont_pc clos_pc =
-  Code.traverse { fold = fold_children } (rewrite_block cont_pc) clos_pc blocks blocks
+  Code.traverse
+    { fold = Code.fold_children_skip_try_body }
+    (rewrite_block cont_pc)
+    clos_pc
+    blocks
+    blocks
 
 (****)