Correctly stack debuginfo for inlined body in classic mode (#1152)

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -274,7 +274,7 @@ module Inlining = struct
   let make_inlined_body acc ~callee ~region_inlined_into ~params ~args
       ~my_closure ~my_region ~my_depth ~body ~free_names_of_body
       ~exn_continuation ~return_continuation ~apply_exn_continuation
-      ~apply_return_continuation ~apply_depth =
+      ~apply_return_continuation ~apply_depth ~apply_dbg =
     let rec_info =
       match apply_depth with
       | None -> Rec_info_expr.initial
@@ -303,10 +303,19 @@ module Inlining = struct
       in
       acc, Expr.apply_renaming body renaming
     in
-    Inlining_helpers.make_inlined_body ~callee ~region_inlined_into ~params
-      ~args ~my_closure ~my_region ~my_depth ~rec_info ~body:(acc, body)
-      ~exn_continuation ~return_continuation ~apply_exn_continuation
-      ~apply_return_continuation ~bind_params ~bind_depth ~apply_renaming
+    let acc, body =
+      Inlining_helpers.make_inlined_body ~callee ~region_inlined_into ~params
+        ~args ~my_closure ~my_region ~my_depth ~rec_info ~body:(acc, body)
+        ~exn_continuation ~return_continuation ~apply_exn_continuation
+        ~apply_return_continuation ~bind_params ~bind_depth ~apply_renaming
+    in
+    Let_with_acc.create acc
+      (Bound_pattern.singleton
+         (VB.create (Variable.create "inlined_dbg") Name_mode.normal))
+      (Named.create_prim
+         (Nullary (Enter_inlined_apply { dbg = apply_dbg }))
+         Debuginfo.none)
+      ~body
 
   let wrap_inlined_body_for_exn_extra_args acc ~extra_args
       ~apply_exn_continuation ~apply_return_continuation ~result_arity
@@ -327,6 +336,7 @@ module Inlining = struct
       ~make_inlined_body ~apply_cont_create ~let_cont_create
 
   let inline acc ~apply ~apply_depth ~func_desc:code =
+    let apply_dbg = Apply.dbg apply in
     let callee = Apply.callee apply in
     let region_inlined_into = Apply.region apply in
     let args = Apply.args apply in
@@ -357,7 +367,7 @@ module Inlining = struct
           make_inlined_body ~callee ~region_inlined_into
             ~params:(Bound_parameters.vars params)
             ~args ~my_closure ~my_region ~my_depth ~body ~free_names_of_body
-            ~exn_continuation ~return_continuation ~apply_depth
+            ~exn_continuation ~return_continuation ~apply_depth ~apply_dbg
         in
         let acc = Acc.with_free_names Name_occurrences.empty acc in
         let acc = Acc.increment_metrics cost_metrics acc in

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -500,7 +500,7 @@ let init_or_assign env (ia : Flambda_primitive.Init_or_assign.t) :
 let nullop _env (op : Flambda_primitive.nullary_primitive) : Fexpr.nullop =
   match op with
   | Begin_region -> Begin_region
-  | Invalid _ | Optimised_out _ | Probe_is_enabled _ ->
+  | Invalid _ | Optimised_out _ | Probe_is_enabled _ | Enter_inlined_apply _ ->
     Misc.fatal_errorf "TODO: Nullary primitive: %a" Flambda_primitive.print
       (Flambda_primitive.Nullary op)
 

middle_end/flambda2/simplify/simplify_nullary_primitive.ml

@@ -40,3 +40,12 @@ let simplify_nullary_primitive dacc original_prim (prim : P.nullary_primitive)
     let ty = T.any_region in
     let dacc = DA.add_variable dacc result_var ty in
     Simplify_primitive_result.create named ~try_reify:false dacc
+  | Enter_inlined_apply { dbg } ->
+    let dacc =
+      DA.map_denv dacc ~f:(fun denv ->
+          DE.set_inlined_debuginfo denv (DE.add_inlined_debuginfo denv dbg))
+    in
+    let named = Named.create_simple Simple.const_unit in
+    let ty = T.this_tagged_immediate Targetint_31_63.zero in
+    let dacc = DA.add_variable dacc result_var ty in
+    Simplify_primitive_result.create named ~try_reify:false dacc

middle_end/flambda2/terms/code_size.ml

@@ -309,6 +309,7 @@ let nullary_prim_size prim =
   | Optimised_out _ -> 0
   | Probe_is_enabled { name = _ } -> 4
   | Begin_region -> 1
+  | Enter_inlined_apply _ -> 0
 
 let unary_prim_size prim =
   match (prim : Flambda_primitive.unary_primitive) with

middle_end/flambda2/terms/flambda_primitive.ml

@@ -575,9 +575,12 @@ type nullary_primitive =
   | Optimised_out of K.t
   | Probe_is_enabled of { name : string }
   | Begin_region
+  | Enter_inlined_apply of { dbg : Debuginfo.t }
 
 let nullary_primitive_eligible_for_cse = function
-  | Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region -> false
+  | Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region
+  | Enter_inlined_apply _ ->
+    false
 
 let compare_nullary_primitive p1 p2 =
   match p1, p2 with
@@ -586,12 +589,23 @@ let compare_nullary_primitive p1 p2 =
   | Probe_is_enabled { name = name1 }, Probe_is_enabled { name = name2 } ->
     String.compare name1 name2
   | Begin_region, Begin_region -> 0
-  | Invalid _, (Optimised_out _ | Probe_is_enabled _ | Begin_region) -> -1
-  | Optimised_out _, (Probe_is_enabled _ | Begin_region) -> -1
+  | Enter_inlined_apply { dbg = dbg1 }, Enter_inlined_apply { dbg = dbg2 } ->
+    Debuginfo.compare dbg1 dbg2
+  | ( Invalid _,
+      ( Optimised_out _ | Probe_is_enabled _ | Begin_region
+      | Enter_inlined_apply _ ) ) ->
+    -1
+  | Optimised_out _, (Probe_is_enabled _ | Begin_region | Enter_inlined_apply _)
+    ->
+    -1
   | Optimised_out _, Invalid _ -> 1
-  | Probe_is_enabled _, Begin_region -> -1
+  | Probe_is_enabled _, (Begin_region | Enter_inlined_apply _) -> -1
   | Probe_is_enabled _, (Invalid _ | Optimised_out _) -> 1
+  | Begin_region, Enter_inlined_apply _ -> -1
   | Begin_region, (Invalid _ | Optimised_out _ | Probe_is_enabled _) -> 1
+  | ( Enter_inlined_apply _,
+      (Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region) ) ->
+    1
 
 let equal_nullary_primitive p1 p2 = compare_nullary_primitive p1 p2 = 0
 
@@ -606,13 +620,17 @@ let print_nullary_primitive ppf p =
   | Probe_is_enabled { name } ->
     Format.fprintf ppf "@[<hov 1>(Probe_is_enabled@ %s)@]" name
   | Begin_region -> Format.pp_print_string ppf "Begin_region"
+  | Enter_inlined_apply { dbg } ->
+    Format.fprintf ppf "@[<hov 1>(Enter_inlined_apply@ %a)@]"
+      Debuginfo.print_compact dbg
 
 let result_kind_of_nullary_primitive p : result_kind =
   match p with
   | Invalid k -> Singleton k
   | Optimised_out k -> Singleton k
   | Probe_is_enabled _ -> Singleton K.naked_immediate
   | Begin_region -> Singleton K.region
+  | Enter_inlined_apply _ -> Unit
 
 let effects_and_coeffects_of_begin_region : Effects_and_coeffects.t =
   (* Ensure these don't get moved, but allow them to be deleted. *)
@@ -627,10 +645,16 @@ let effects_and_coeffects_of_nullary_primitive p : Effects_and_coeffects.t =
        moved around. *)
     Arbitrary_effects, Has_coeffects, Strict
   | Begin_region -> effects_and_coeffects_of_begin_region
+  | Enter_inlined_apply _ ->
+    (* This doesn't really have effects, but without effects, these primitives
+       get deleted during lambda_to_flambda. *)
+    Arbitrary_effects, Has_coeffects, Strict
 
 let nullary_classify_for_printing p =
   match p with
-  | Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region -> Neither
+  | Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region
+  | Enter_inlined_apply _ ->
+    Neither
 
 type unary_primitive =
   | Duplicate_block of { kind : Duplicate_block_kind.t }
@@ -1666,7 +1690,9 @@ let equal t1 t2 = compare t1 t2 = 0
 
 let free_names t =
   match t with
-  | Nullary (Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region) ->
+  | Nullary
+      ( Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region
+      | Enter_inlined_apply _ ) ->
     Name_occurrences.empty
   | Unary (prim, x0) ->
     Name_occurrences.union
@@ -1691,7 +1717,9 @@ let free_names t =
 let apply_renaming t renaming =
   let apply simple = Simple.apply_renaming simple renaming in
   match t with
-  | Nullary (Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region) ->
+  | Nullary
+      ( Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region
+      | Enter_inlined_apply _ ) ->
     t
   | Unary (prim, x0) ->
     let prim' = apply_renaming_unary_primitive prim renaming in
@@ -1719,7 +1747,9 @@ let apply_renaming t renaming =
 
 let ids_for_export t =
   match t with
-  | Nullary (Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region) ->
+  | Nullary
+      ( Invalid _ | Optimised_out _ | Probe_is_enabled _ | Begin_region
+      | Enter_inlined_apply _ ) ->
     Ids_for_export.empty
   | Unary (prim, x0) ->
     Ids_for_export.union

middle_end/flambda2/terms/flambda_primitive.mli

@@ -208,6 +208,9 @@ type nullary_primitive =
       (** Starting delimiter of local allocation region, returning a region
           name. For regions for the "try" part of a "try...with", use
           [Begin_try_region] (below) instead. *)
+  | Enter_inlined_apply of { dbg : Debuginfo.t }
+      (** Used in classic mode to denote the start of an inlined function body.
+          This is then used in to_cmm to correctly add inlined debuginfo. *)
 
 (** Untagged binary integer arithmetic operations.
 

middle_end/flambda2/to_cmm/to_cmm_env.ml

@@ -34,7 +34,8 @@ type cont =
   | Inline of
       { handler_params : Bound_parameters.t;
         handler_params_occurrences : Num_occurrences.t Variable.Map.t;
-        handler_body : Flambda.Expr.t
+        handler_body : Flambda.Expr.t;
+        handler_body_inlined_debuginfo : Debuginfo.t
       }
 
 type extra_info = Untag of Cmm.expression
@@ -130,6 +131,8 @@ type t =
        This is relative to the flambda expression being currently translated,
        i.e. either the unit initialization code, or the body of a function. This
        information is reset when entering a new function. *)
+    inlined_debuginfo : Debuginfo.t;
+    (* Debuginfo corresponding to inlined functions. *)
     return_continuation : Continuation.t;
     (* The (non-exceptional) return continuation of the current context (used to
        determine which calls are tail-calls). *)
@@ -162,30 +165,7 @@ type translation_result =
     expr : expr_with_info
   }
 
-let create offsets functions_info ~trans_prim ~return_continuation
-    ~exn_continuation =
-  { return_continuation;
-    exn_continuation;
-    offsets;
-    functions_info;
-    trans_prim;
-    stages = [];
-    bindings = Variable.Map.empty;
-    inline_once_aliases = Variable.Map.empty;
-    vars_extra = Variable.Map.empty;
-    vars = Variable.Map.empty;
-    conts = Continuation.Map.empty;
-    exn_handlers = Continuation.Set.singleton exn_continuation;
-    exn_conts_extra_args = Continuation.Map.empty
-  }
-
-let enter_function_body env ~return_continuation ~exn_continuation =
-  create env.offsets env.functions_info ~trans_prim:env.trans_prim
-    ~return_continuation ~exn_continuation
-
-let return_continuation env = env.return_continuation
-
-let exn_continuation env = env.exn_continuation
+(* Printing *)
 
 let print_extra_info ppf = function
   | Untag e -> Format.fprintf ppf "Untag(%a)" Printcmm.expression e
@@ -249,6 +229,46 @@ let print ppf t =
   Format.fprintf ppf "@[<hov 1>(@[<hov 1>(stages %a)@]@ )@]" print_stages
     t.stages
 
+(* Creation *)
+
+let create offsets functions_info ~trans_prim ~return_continuation
+    ~exn_continuation =
+  { return_continuation;
+    exn_continuation;
+    offsets;
+    functions_info;
+    trans_prim;
+    inlined_debuginfo = Debuginfo.none;
+    stages = [];
+    bindings = Variable.Map.empty;
+    inline_once_aliases = Variable.Map.empty;
+    vars_extra = Variable.Map.empty;
+    vars = Variable.Map.empty;
+    conts = Continuation.Map.empty;
+    exn_handlers = Continuation.Set.singleton exn_continuation;
+    exn_conts_extra_args = Continuation.Map.empty
+  }
+
+let enter_function_body env ~return_continuation ~exn_continuation =
+  create env.offsets env.functions_info ~trans_prim:env.trans_prim
+    ~return_continuation ~exn_continuation
+
+(* Debuginfo *)
+
+let enter_inlined_apply t dbg =
+  let inlined_debuginfo = Debuginfo.inline t.inlined_debuginfo dbg in
+  { t with inlined_debuginfo }
+
+let set_inlined_debuginfo t inlined_debuginfo = { t with inlined_debuginfo }
+
+let add_inlined_debuginfo t dbg = Debuginfo.inline t.inlined_debuginfo dbg
+
+(* Continuations *)
+
+let return_continuation env = env.return_continuation
+
+let exn_continuation env = env.exn_continuation
+
 (* Code and closures *)
 
 let get_code_metadata env code_id =
@@ -324,8 +344,14 @@ let add_jump_cont env k ~param_types =
 
 let add_inline_cont env k ~handler_params ~handler_params_occurrences
     ~handler_body =
+  let handler_body_inlined_debuginfo = env.inlined_debuginfo in
   let info =
-    Inline { handler_params; handler_body; handler_params_occurrences }
+    Inline
+      { handler_params;
+        handler_body;
+        handler_params_occurrences;
+        handler_body_inlined_debuginfo
+      }
   in
   let conts = Continuation.Map.add k info env.conts in
   { env with conts }

middle_end/flambda2/to_cmm/to_cmm_env.mli

@@ -93,6 +93,19 @@ val enter_function_body :
   exn_continuation:Continuation.t ->
   t
 
+(** {2 Debuginfo} *)
+
+(** Add the inlined debuginfo from the env to the debuginfo provided,
+    in order to get the correct debuginfo to attach. *)
+val add_inlined_debuginfo : t -> Debuginfo.t -> Debuginfo.t
+
+(** Adjust the inlined debuginfo in the env to represent the fact
+    that we entered the inlined body of a function. *)
+val enter_inlined_apply : t -> Debuginfo.t -> t
+
+(** Set the inlined debuginfo. *)
+val set_inlined_debuginfo : t -> Debuginfo.t -> t
+
 (** {2 Continuations} *)
 
 (** Returns the return continuation of the environment. *)
@@ -286,7 +299,8 @@ type cont = private
   | Inline of
       { handler_params : Bound_parameters.t;
         handler_params_occurrences : Num_occurrences.t Variable.Map.t;
-        handler_body : Flambda.Expr.t
+        handler_body : Flambda.Expr.t;
+        handler_body_inlined_debuginfo : Debuginfo.t
       }
 
 (** Record that the given continuation should be compiled to a jump, creating a