Add location after return statement

Author: vouillon

compiler/bin-wasm_of_ocaml/compile.ml

@@ -193,7 +193,7 @@ let build_js_runtime ~primitives ?runtime_arguments () =
       List.split_last
       @@ Driver.link_and_pack
            ~link:`Needed
-           [ Javascript.Return_statement (Some (EObj l)), N ]
+           [ Javascript.Return_statement (Some (EObj l), N), N ]
     with
     | Some x -> x
     | None -> assert false

compiler/lib/generate.ml

@@ -561,6 +561,7 @@ type state =
   ; dom : Structure.graph
   ; visited_blocks : Addr.Set.t ref
   ; ctx : Ctx.t
+  ; pc : Addr.t
   }
 
 module DTree = struct
@@ -681,7 +682,7 @@ let build_graph ctx pc =
   let visited_blocks = ref Addr.Set.empty in
   let structure = Structure.build_graph ctx.Ctx.blocks pc in
   let dom = Structure.dominator_tree structure in
-  { visited_blocks; structure; dom; ctx }
+  { visited_blocks; structure; dom; ctx; pc }
 
 (****)
 
@@ -835,7 +836,8 @@ let generate_apply_fun ctx { arity; exact; trampolined } =
     ( None
     , J.fun_
         (f :: params)
-        [ J.Return_statement (Some (apply_fun_raw ctx f' params' exact trampolined)), J.N
+        [ ( J.Return_statement (Some (apply_fun_raw ctx f' params' exact trampolined), J.N)
+          , J.N )
         ]
         J.N )
 
@@ -1282,7 +1284,7 @@ let rec translate_expr ctx queue loc x e level : _ * J.statement_list =
                 loc
             in
             let e =
-              J.EFun (Some f, J.fun_ args [ J.Return_statement (Some call), J.N ] J.N)
+              J.EFun (Some f, J.fun_ args [ J.Return_statement (Some call, J.N), J.N ] J.N)
             in
             e, const_p, queue
         | Extern "caml_alloc_dummy_function", _ -> assert false
@@ -1675,6 +1677,8 @@ and compile_decision_tree kind st scope_stack loc cx dtree ~fall_through =
         in
         ( never1 && never2
         , Js_simpl.if_statement
+            ~function_end:(fun () ->
+              source_location_ctx st.ctx ~force:After (After st.pc))
             e'
             loc
             (Js_simpl.block iftrue)
@@ -1758,15 +1762,25 @@ and compile_conditional st queue ~fall_through last scope_stack : _ * _ =
         let return_expr =
           if Var.equal st.ctx.deadcode_sentinal x then None else Some cx
         in
-        true, flush_all queue [ J.Return_statement return_expr, loc ]
+        let loc' =
+          match cx with
+          | ECall _ -> (
+              (* We usually don't have a good locations for tail
+                 calls, so use the end of the function instead *)
+              match source_location_ctx st.ctx ~force:After (After st.pc) with
+              | J.N -> loc
+              | loc -> loc)
+          | _ -> loc
+        in
+        true, flush_all queue [ J.Return_statement (return_expr, loc'), loc ]
     | Raise (x, k) ->
         let (_px, cx), queue = access_queue queue x in
         true, flush_all queue (throw_statement st.ctx cx k loc)
     | Stop ->
         let e_opt =
           if st.ctx.Ctx.should_export then Some (s_var Global_constant.exports) else None
         in
-        true, flush_all queue [ J.Return_statement e_opt, loc ]
+        true, flush_all queue [ J.Return_statement (e_opt, loc), loc ]
     | Branch cont -> compile_branch st queue cont scope_stack ~fall_through
     | Pushtrap (c1, x, e1) ->
         let never_body, body = compile_branch st [] c1 scope_stack ~fall_through in

compiler/lib/javascript.ml

@@ -347,7 +347,7 @@ and statement =
       * (statement * location)
   | Continue_statement of Label.t option
   | Break_statement of Label.t option
-  | Return_statement of expression option
+  | Return_statement of expression option * location
   | With_statement of expression * (statement * location)
   | Labelled_statement of Label.t * (statement * location)
   | Switch_statement of

compiler/lib/javascript.mli

@@ -265,7 +265,7 @@ and statement =
       * (statement * location)
   | Continue_statement of Label.t option
   | Break_statement of Label.t option
-  | Return_statement of expression option
+  | Return_statement of expression option * location
   | With_statement of expression * (statement * location)
   | Labelled_statement of Label.t * (statement * location)
   | Switch_statement of

compiler/lib/js_output.ml

@@ -509,13 +509,14 @@ struct
             PP.end_group f);
         PP.end_group f;
         (match b, consise with
-        | [ (Return_statement (Some e), loc) ], true ->
+        | [ (Return_statement (Some e, loc), loc') ], true ->
             (* Should not starts with '{' *)
             PP.start_group f 1;
             PP.break1 f;
-            output_debug_info f loc;
+            output_debug_info f loc';
             parenthesized_expression ~obj:true AssignementExpression f e;
-            PP.end_group f
+            PP.end_group f;
+            output_debug_info f loc
         | l, _ ->
             let b =
               match l with
@@ -1466,10 +1467,11 @@ struct
         let (Utf8 l) = nane_of_label s in
         PP.string f l;
         last_semi ()
-    | Return_statement e -> (
+    | Return_statement (e, loc) -> (
         match e with
         | None ->
             PP.string f "return";
+            output_debug_info f loc;
             last_semi ()
         | Some (EFun (i, ({ async = false; generator = false }, l, b, pc))) ->
             PP.start_group f 1;
@@ -1490,6 +1492,7 @@ struct
             function_body f b;
             output_debug_info f pc;
             PP.string f "}";
+            output_debug_info f loc;
             last_semi ();
             PP.end_group f
         | Some e ->
@@ -1498,6 +1501,7 @@ struct
             PP.non_breaking_space f;
             PP.start_group f 0;
             expression Expression f e;
+            output_debug_info f loc;
             last_semi ();
             PP.end_group f;
             PP.end_group f

compiler/lib/js_parser.mly

@@ -736,7 +736,7 @@ break_stmt:
  | T_BREAK l=label? sc { (Break_statement (l)) }
 
 return_stmt:
- | T_RETURN e=expr? sc { (Return_statement e) }
+ | T_RETURN e=expr? sc { (Return_statement (e, p $endpos(e))) }
 
 switch_stmt:
  | T_SWITCH "(" subject=expr ")" cb=case_block
@@ -1100,7 +1100,7 @@ async_arrow_function:
 (* was called consise body in spec *)
 arrow_body:
  | "{" b=function_body "}" { b, false }
- | e=assignment_expr_for_consise_body { [(Return_statement (Some e), p $symbolstartpos)], true }
+ | e=assignment_expr_for_consise_body { [(Return_statement (Some e, p $endpos), p $symbolstartpos)], true }
 
 (*----------------------------*)
 (* no in                    *)

compiler/lib/js_simpl.ml

@@ -110,13 +110,13 @@ exception Not_expression
 
 let rec expression_of_statement_list l =
   match l with
-  | (J.Return_statement (Some e), _) :: _ -> e
+  | (J.Return_statement (Some e, _), _) :: _ -> e
   | (J.Expression_statement e, _) :: rem -> J.ESeq (e, expression_of_statement_list rem)
   | _ -> raise Not_expression
 
 let expression_of_statement st =
   match fst st with
-  | J.Return_statement (Some e) -> e
+  | J.Return_statement (Some e, _) -> e
   | J.Block l -> expression_of_statement_list l
   | _ -> raise Not_expression
 
@@ -194,15 +194,16 @@ and depth_class_block b =
       | J.CEField _ -> acc
       | J.CEStaticBLock b -> depth_block b + 2)
 
-let rec if_statement_2 e loc iftrue truestop iffalse falsestop =
+let rec if_statement_2 ~function_end e loc iftrue truestop iffalse falsestop =
   let e = simplify_condition e in
   match fst iftrue, fst iffalse with
   (* Empty blocks *)
   | J.Block [], J.Block [] -> (
       match e with
       | J.EVar _ -> []
       | _ -> [ J.Expression_statement e, loc ])
-  | J.Block [], _ -> if_statement_2 (enot e) loc iffalse falsestop iftrue truestop
+  | J.Block [], _ ->
+      if_statement_2 ~function_end (enot e) loc iffalse falsestop iftrue truestop
   | _, J.Block [] -> [ J.If_statement (e, iftrue, None), loc ]
   | _ -> (
       try
@@ -226,7 +227,7 @@ let rec if_statement_2 e loc iftrue truestop iffalse falsestop =
         try
           let e1 = expression_of_statement iftrue in
           let e2 = expression_of_statement iffalse in
-          [ J.Return_statement (Some (J.ECond (e, e1, e2))), loc ]
+          [ J.Return_statement (Some (J.ECond (e, e1, e2)), function_end ()), loc ]
         with Not_expression ->
           let truestop, falsestop =
             if truestop && falsestop
@@ -247,32 +248,48 @@ let unopt b =
   | Some b -> b
   | None -> J.Block [], J.N
 
-let if_statement e loc iftrue truestop iffalse falsestop =
+let if_statement ~function_end e loc iftrue truestop iffalse falsestop =
   (*FIX: should be done at an earlier stage*)
   let e = simplify_condition e in
   match iftrue, iffalse with
   (* Shared statements *)
   | (J.If_statement (e', iftrue', iffalse'), _), _ when Poly.(iffalse = unopt iffalse') ->
-      if_statement_2 (J.EBin (J.And, e, e')) loc iftrue' truestop iffalse falsestop
+      if_statement_2
+        ~function_end
+        (J.EBin (J.And, e, e'))
+        loc
+        iftrue'
+        truestop
+        iffalse
+        falsestop
   | (J.If_statement (e', iftrue', iffalse'), _), _ when Poly.(iffalse = iftrue') ->
       if_statement_2
+        ~function_end
         (J.EBin (J.And, e, J.EUn (J.Not, e')))
         loc
         (unopt iffalse')
         truestop
         iffalse
         falsestop
   | _, (J.If_statement (e', iftrue', iffalse'), _) when Poly.(iftrue = iftrue') ->
-      if_statement_2 (J.EBin (J.Or, e, e')) loc iftrue truestop (unopt iffalse') falsestop
+      if_statement_2
+        ~function_end
+        (J.EBin (J.Or, e, e'))
+        loc
+        iftrue
+        truestop
+        (unopt iffalse')
+        falsestop
   | _, (J.If_statement (e', iftrue', iffalse'), _) when Poly.(iftrue = unopt iffalse') ->
       if_statement_2
+        ~function_end
         (J.EBin (J.Or, e, J.EUn (J.Not, e')))
         loc
         iftrue
         truestop
         iftrue'
         falsestop
-  | _ -> if_statement_2 e loc iftrue truestop iffalse falsestop
+  | _ -> if_statement_2 ~function_end e loc iftrue truestop iffalse falsestop
 
 let function_body b =
   (* We only check for a return at the end since it is by far the most
@@ -281,7 +298,7 @@ let function_body b =
     let rec check l =
       match l with
       | [] -> false
-      | [ (J.Return_statement None, _) ] -> true
+      | [ (J.Return_statement (None, _), _) ] -> true
       | _ :: r -> check r
     in
     check b
@@ -291,7 +308,7 @@ let function_body b =
     let rec remove acc l =
       match l with
       | [] -> acc
-      | [ (J.Return_statement None, _) ] -> acc
+      | [ (J.Return_statement (None, _), _) ] -> acc
       | i :: r -> remove (i :: acc) r
     in
     List.rev (remove [] b)

compiler/lib/js_simpl.mli

@@ -21,7 +21,8 @@
 open Javascript
 
 val if_statement :
-     expression
+     function_end:(unit -> location)
+  -> expression
   -> location
   -> statement * location
   -> bool

compiler/lib/js_traverse.ml

@@ -187,7 +187,7 @@ class map : mapper =
           ForAwaitOf_statement (e1, m#expression e2, (m#statement s, m#loc loc))
       | Continue_statement s -> Continue_statement s
       | Break_statement s -> Break_statement s
-      | Return_statement e -> Return_statement (m#expression_o e)
+      | Return_statement (e, loc) -> Return_statement (m#expression_o e, m#loc loc)
       | Labelled_statement (l, (s, loc)) ->
           Labelled_statement (l, (m#statement s, m#loc loc))
       | Throw_statement e -> Throw_statement (m#expression e)
@@ -549,7 +549,7 @@ class iter : iterator =
           m#statement s
       | Continue_statement _ -> ()
       | Break_statement _ -> ()
-      | Return_statement e -> m#expression_o e
+      | Return_statement (e, _) -> m#expression_o e
       | Labelled_statement (_, (s, _)) -> m#statement s
       | Throw_statement e -> m#expression e
       | Switch_statement (e, l, def, l') ->
@@ -1750,8 +1750,16 @@ class simpl =
           | If_statement (ENum n, _, iffalse) when Num.is_zero n -> opt_cons iffalse rem
           (* if (e1) return e2 else return e3 --> return e1 ? e2 : e3 *)
           | If_statement
-              (cond, (Return_statement (Some e1), _), Some (Return_statement (Some e2), _))
-            -> (Return_statement (Some (ECond (cond, e1, e2))), loc) :: rem
+              ( cond
+              , (Return_statement (Some e1, _), _)
+              , Some (Return_statement (Some e2, _), _) ) ->
+              ( Return_statement
+                  ( Some (ECond (cond, e1, e2))
+                  , U
+                    (*TODO: it would be better to use the location of the
+                      end of the function, but we can't easily get it. *) )
+              , loc )
+              :: rem
           (* if (e1) v1 = e2 else v1 = e3 --> v1 = e1 ? e2 : e3 *)
           | If_statement
               ( cond