Compiler: disconnect closure cont from closure location. (#1947)

* Compiler: use explicit locations for closures
- Store Parse_info.t inside closure
- Debug infos are no longer needed after parsing bytecode
---------

Co-authored-by: Jérôme Vouillon <[email protected]>

Author: hhugo

CHANGES.md

@@ -25,6 +25,7 @@
 * Lib: make the Wasm version of Json.output work with native ints and JavaScript objects (#1872)
 * Compiler: static evaluation of more primitives (#1912)
 * Compiler: faster compilation by stopping sooner when optimizations become unproductive (#1939)
+* Compiler: improve debug/sourcemap location of closures (#1947)
 * Runtime: use Dataview to convert between floats and bit representation
 
 ## Bug fixes

compiler/bin-js_of_ocaml/build_fs.ml

@@ -82,7 +82,6 @@ function jsoo_create_file_extern(name,content){
           ~link:`Needed
           ~formatter:pfs_fmt
           ~source_map:false
-          (Parse_bytecode.Debug.create ~include_cmis:false false)
           code
       in
       ())

compiler/bin-js_of_ocaml/compile.ml

@@ -277,7 +277,6 @@ let run
             ~wrap_with_fun
             ~source_map:(source_map_enabled source_map)
             ~formatter
-            one.debug
             code
       | `File, formatter ->
           let fs_instr1, fs_instr2 =
@@ -301,22 +300,14 @@ let run
               ~wrap_with_fun
               ~source_map:(source_map_enabled source_map)
               ~formatter
-              one.debug
               code
           in
           Option.iter fs_output ~f:(fun file ->
               Filename.gen_file file (fun chan ->
                   let instr = fs_instr2 in
                   let code = Code.prepend Code.empty instr in
                   let pfs_fmt = Pretty_print.to_out_channel chan in
-                  Driver.f'
-                    ~standalone
-                    ~link:`Needed
-                    ?profile
-                    ~wrap_with_fun
-                    pfs_fmt
-                    one.debug
-                    code));
+                  Driver.f' ~standalone ~link:`Needed ?profile ~wrap_with_fun pfs_fmt code));
           res
     in
     if times () then Format.eprintf "compilation: %a@." Timer.print t;

compiler/bin-wasm_of_ocaml/compile.ml

@@ -240,15 +240,13 @@ let generate_prelude ~out_file =
     Driver.optimize ~profile code
   in
   let context = Generate.start () in
-  let debug = Parse_bytecode.Debug.create ~include_cmis:false false in
   let _ =
     Generate.f
       ~context
       ~unit_name:(Some "prelude")
       ~live_vars:variable_uses
       ~in_cps
       ~deadcode_sentinal
-      ~debug
       program
   in
   Generate.output ch ~context;
@@ -404,15 +402,13 @@ let run
       Driver.optimize ~profile code
     in
     let context = Generate.start () in
-    let debug = one.debug in
     let toplevel_name, generated_js =
       Generate.f
         ~context
         ~unit_name
         ~live_vars:variable_uses
         ~in_cps
         ~deadcode_sentinal
-        ~debug
         program
     in
     if standalone then Generate.add_start_function ~context toplevel_name;

compiler/lib-wasm/closure_conversion.ml

@@ -35,15 +35,15 @@ let iter_closures ~f instrs =
         let l = f clos_acc in
         List.rev_map
           ~f:(fun g ->
-            let params, cont = Var.Map.find g clos_acc in
-            Let (g, Closure (params, cont)))
+            let params, cont, cloc = Var.Map.find g clos_acc in
+            Let (g, Closure (params, cont, cloc)))
           l
         @ instr_acc
     in
     match instrs with
     | [] -> List.rev (push_closures clos_acc instr_acc)
-    | Let (g, Closure (params, cont)) :: rem ->
-        iter_closures_rec f instr_acc (Var.Map.add g (params, cont) clos_acc) rem
+    | Let (g, Closure (params, cont, cloc)) :: rem ->
+        iter_closures_rec f instr_acc (Var.Map.add g (params, cont, cloc) clos_acc) rem
     | i :: rem ->
         iter_closures_rec f (i :: push_closures clos_acc instr_acc) Var.Map.empty rem
   in
@@ -79,7 +79,7 @@ let mark_bound_variables var_depth block depth =
   Freevars.iter_block_bound_vars (fun x -> var_depth.(Var.idx x) <- depth) block;
   List.iter block.body ~f:(fun i ->
       match i with
-      | Let (_, Closure (params, _)) ->
+      | Let (_, Closure (params, _, _)) ->
           List.iter params ~f:(fun x -> var_depth.(Var.idx x) <- depth + 1)
       | _ -> ())
 
@@ -93,7 +93,7 @@ let rec traverse var_depth closures program pc depth =
         List.fold_left
           ~f:(fun program i ->
             match i with
-            | Let (_, Closure (_, (pc', _))) ->
+            | Let (_, Closure (_, (pc', _), _)) ->
                 traverse var_depth closures program pc' (depth + 1)
             | _ -> program)
           ~init:program
@@ -103,7 +103,7 @@ let rec traverse var_depth closures program pc depth =
         iter_closures block.body ~f:(fun l ->
             let free_vars =
               Var.Map.fold
-                (fun f (_, (pc', _)) free_vars ->
+                (fun f (_, (pc', _), _) free_vars ->
                   Var.Map.add
                     f
                     (collect_free_vars program var_depth (depth + 1) pc' !closures)
@@ -136,7 +136,7 @@ let rec traverse var_depth closures program pc depth =
                   let functions =
                     let arities =
                       Var.Map.fold
-                        (fun f (params, _) m -> Var.Map.add f (List.length params) m)
+                        (fun f (params, _, _) m -> Var.Map.add f (List.length params) m)
                         l
                         Var.Map.empty
                     in

compiler/lib-wasm/generate.ml

@@ -39,7 +39,6 @@ module Generate (Target : Target_sig.S) = struct
     ; blocks : block Addr.Map.t
     ; closures : Closure_conversion.closure Var.Map.t
     ; global_context : Code_generation.context
-    ; debug : Parse_bytecode.Debug.t
     }
 
   let label_index context pc =
@@ -909,6 +908,7 @@ module Generate (Target : Target_sig.S) = struct
       ~unit_name
       params
       ((pc, _) as cont)
+      cloc
       acc =
     let g = Structure.build_graph ctx.blocks pc in
     let dom = Structure.dominator_tree g in
@@ -1088,8 +1088,7 @@ module Generate (Target : Target_sig.S) = struct
                (fun ~result_typ ~fall_through ~context ->
                  translate_branch result_typ fall_through (-1) cont context)
            in
-           let end_loc = Parse_bytecode.Debug.find_loc ctx.debug ~position:After pc in
-           match end_loc with
+           match cloc with
            | Some loc -> event loc
            | None -> return ())
     in
@@ -1176,8 +1175,7 @@ module Generate (Target : Target_sig.S) = struct
     ~should_export
     ~warn_on_unhandled_effect
 *)
-      ~deadcode_sentinal
-      ~debug =
+      ~deadcode_sentinal =
     global_context.unit_name <- unit_name;
     let p, closures = Closure_conversion.f p in
     (*
@@ -1190,15 +1188,14 @@ module Generate (Target : Target_sig.S) = struct
       ; blocks = p.blocks
       ; closures
       ; global_context
-      ; debug
       }
     in
     let toplevel_name = Var.fresh_n "toplevel" in
     let functions =
       Code.fold_closures_outermost_first
         p
-        (fun name_opt params cont ->
-          translate_function p ctx name_opt ~toplevel_name ~unit_name params cont)
+        (fun name_opt params cont cloc ->
+          translate_function p ctx name_opt ~toplevel_name ~unit_name params cont cloc)
         []
     in
     let functions =
@@ -1293,10 +1290,10 @@ let init = G.init
 
 let start () = make_context ~value_type:Gc_target.Value.value
 
-let f ~context ~unit_name p ~live_vars ~in_cps ~deadcode_sentinal ~debug =
+let f ~context ~unit_name p ~live_vars ~in_cps ~deadcode_sentinal =
   let t = Timer.make () in
   let p = fix_switch_branches p in
-  let res = G.f ~context ~unit_name ~live_vars ~in_cps ~deadcode_sentinal ~debug p in
+  let res = G.f ~context ~unit_name ~live_vars ~in_cps ~deadcode_sentinal p in
   if times () then Format.eprintf "  code gen.: %a@." Timer.print t;
   res
 

compiler/lib-wasm/generate.mli

@@ -27,7 +27,6 @@ val f :
   -> live_vars:int array
   -> in_cps:Effects.in_cps
   -> deadcode_sentinal:Code.Var.t
-  -> debug:Parse_bytecode.Debug.t
   -> Wasm_ast.var * (string list * (string * Javascript.expression) list)
 
 val add_start_function : context:Code_generation.context -> Wasm_ast.var -> unit

compiler/lib/code.ml

@@ -384,7 +384,7 @@ type expr =
       }
   | Block of int * Var.t array * array_or_not * mutability
   | Field of Var.t * int * field_type
-  | Closure of Var.t list * cont
+  | Closure of Var.t list * cont * Parse_info.t option
   | Constant of constant
   | Prim of prim * prim_arg list
   | Special of special
@@ -538,7 +538,7 @@ module Print = struct
         Format.fprintf f "}"
     | Field (x, i, Non_float) -> Format.fprintf f "%a[%d]" Var.print x i
     | Field (x, i, Float) -> Format.fprintf f "FLOAT{%a[%d]}" Var.print x i
-    | Closure (l, c) -> Format.fprintf f "fun(%a){%a}" var_list l cont c
+    | Closure (l, c, _) -> Format.fprintf f "fun(%a){%a}" var_list l cont c
     | Constant c -> Format.fprintf f "CONST{%a}" constant c
     | Prim (p, l) -> prim f p l
     | Special s -> special f s
@@ -597,10 +597,10 @@ let fold_closures p f accu =
     (fun _ block accu ->
       List.fold_left block.body ~init:accu ~f:(fun accu i ->
           match i with
-          | Let (x, Closure (params, cont)) -> f (Some x) params cont accu
+          | Let (x, Closure (params, cont, cloc)) -> f (Some x) params cont cloc accu
           | _ -> accu))
     p.blocks
-    (f None [] (p.start, []) accu)
+    (f None [] (p.start, []) None accu)
 
 (****)
 
@@ -756,16 +756,16 @@ let fold_closures_innermost_first { start; blocks; _ } f 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 (x, Closure (params, cont, cloc)) ->
                 let accu = visit blocks (fst cont) f accu in
-                f (Some x) params cont accu
+                f (Some x) params cont cloc accu
             | _ -> accu))
       pc
       blocks
       accu
   in
   let accu = visit blocks start f accu in
-  f None [] (start, []) accu
+  f None [] (start, []) None accu
 
 let fold_closures_outermost_first { start; blocks; _ } f accu =
   let rec visit blocks pc f accu =
@@ -775,15 +775,15 @@ let fold_closures_outermost_first { start; blocks; _ } f 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
+            | Let (x, Closure (params, cont, cloc)) ->
+                let accu = f (Some x) params cont cloc accu in
                 visit blocks (fst cont) f accu
             | _ -> accu))
       pc
       blocks
       accu
   in
-  let accu = f None [] (start, []) accu in
+  let accu = f None [] (start, []) None accu in
   visit blocks start f accu
 
 let eq p1 p2 =
@@ -819,7 +819,7 @@ let invariant { blocks; start; _ } =
       | Apply _ -> ()
       | Block (_, _, _, _) -> ()
       | Field (_, _, _) -> ()
-      | Closure (l, cont) ->
+      | Closure (l, cont, _) ->
           List.iter l ~f:define;
           check_cont cont
       | Constant _ -> ()

compiler/lib/code.mli

@@ -204,7 +204,7 @@ type expr =
       }
   | Block of int * Var.t array * array_or_not * mutability
   | Field of Var.t * int * field_type
-  | Closure of Var.t list * cont
+  | Closure of Var.t list * cont * Parse_info.t option
   | Constant of constant
   | Prim of prim * prim_arg list
   | Special of special
@@ -266,24 +266,34 @@ type 'c fold_blocs = block Addr.Map.t -> Addr.t -> (Addr.t -> 'c -> 'c) -> 'c ->
 type fold_blocs_poly = { fold : 'a. 'a fold_blocs } [@@unboxed]
 
 val fold_closures :
-  program -> (Var.t option -> Var.t list -> cont -> 'd -> 'd) -> 'd -> 'd
+     program
+  -> (Var.t option -> Var.t list -> cont -> Parse_info.t option -> 'd -> 'd)
+  -> 'd
+  -> 'd
 (** [fold_closures p f init] folds [f] over all closures in the program [p],
     starting from the initial value [init]. For each closure, [f] is called
     with the following arguments: the closure name (enclosed in
     {!Stdlib.Some}), its parameter list, the address and parameter instantiation
-    of its first block, and the current accumulator. In addition, [f] is called
-    on the initial block [p.start], with [None] as the closure name.
-    All closures in all blocks of [p] are included in the fold, not only the
-    ones reachable from [p.start]. *)
+    of its first block, the optional closure location and the current accumulator.
+    In addition, [f] is called on the initial block [p.start], with
+    [None] as the closure name.  All closures in all blocks of [p] are
+    included in the fold, not only the ones reachable from
+    [p.start]. *)
 
 val fold_closures_innermost_first :
-  program -> (Var.t option -> Var.t list -> cont -> 'd -> 'd) -> 'd -> 'd
+     program
+  -> (Var.t option -> Var.t list -> cont -> Parse_info.t option -> 'd -> 'd)
+  -> 'd
+  -> 'd
 (** Similar to {!fold_closures}, but applies the fold function to the
     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
+     program
+  -> (Var.t option -> Var.t list -> cont -> Parse_info.t option -> '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. *)

compiler/lib/deadcode.ml

@@ -72,7 +72,7 @@ and mark_expr st e =
       List.iter args ~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
+  | Closure (_, (pc, _), _) -> mark_reachable st pc
   | Special _ -> ()
   | Prim (_, l) ->
       List.iter l ~f:(fun x ->
@@ -140,7 +140,8 @@ let filter_cont blocks st (pc, args) =
 
 let filter_closure blocks st i =
   match i with
-  | Let (x, Closure (l, cont)) -> Let (x, Closure (l, filter_cont blocks st cont))
+  | Let (x, Closure (l, cont, gloc)) ->
+      Let (x, Closure (l, filter_cont blocks st cont, gloc))
   | _ -> i
 
 let filter_live_last blocks st l =