DWARF inlined frames: Debuginfo.t and Flambda 2

Author: mshinwell

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -322,8 +322,8 @@ module Inlining = struct
           ~are_rebuilding_terms decision;
         res
 
-  let make_inlined_body acc ~callee ~region_inlined_into ~params ~args
-      ~my_closure ~my_region ~my_depth ~body ~free_names_of_body
+  let make_inlined_body acc ~callee ~called_code_id ~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_dbg =
     let rec_info =
@@ -355,16 +355,20 @@ module Inlining = struct
       acc, Expr.apply_renaming body renaming
     in
     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
+      Inlining_helpers.make_inlined_body ~callee ~called_code_id
+        ~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 inlined_debuginfo =
+      Inlined_debuginfo.create ~called_code_id ~apply_dbg
     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 }))
+         (Nullary (Enter_inlined_apply { dbg = inlined_debuginfo }))
          Debuginfo.none)
       ~body
 
@@ -421,7 +425,8 @@ module Inlining = struct
           | Known free_names -> free_names
         in
         let make_inlined_body =
-          make_inlined_body ~callee ~region_inlined_into
+          make_inlined_body ~callee ~called_code_id:(Code.code_id code)
+            ~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 ~apply_dbg

middle_end/flambda2/simplify/env/downwards_env.ml

@@ -31,7 +31,7 @@ type get_imported_code = unit -> Exported_code.t
 type t =
   { round : int;
     typing_env : TE.t;
-    inlined_debuginfo : Debuginfo.t;
+    inlined_debuginfo : Inlined_debuginfo.t;
     can_inline : bool;
     inlining_state : Inlining_state.t;
     propagating_float_consts : bool;
@@ -49,11 +49,6 @@ type t =
     loopify_state : Loopify_state.t
   }
 
-let print_debuginfo ppf dbg =
-  if Debuginfo.is_none dbg
-  then Format.pp_print_string ppf "None"
-  else Debuginfo.print_compact ppf dbg
-
 let [@ocamlformat "disable"] print ppf { round; typing_env;
                 inlined_debuginfo; can_inline;
                 inlining_state; propagating_float_consts;
@@ -84,7 +79,7 @@ let [@ocamlformat "disable"] print ppf { round; typing_env;
       )@]"
     round
     TE.print typing_env
-    print_debuginfo inlined_debuginfo
+    Inlined_debuginfo.print inlined_debuginfo
     can_inline
     Inlining_state.print inlining_state
     propagating_float_consts
@@ -112,7 +107,7 @@ let create ~round ~(resolver : resolver)
   in
   { round;
     typing_env;
-    inlined_debuginfo = Debuginfo.none;
+    inlined_debuginfo = Inlined_debuginfo.none;
     can_inline = true;
     inlining_state = Inlining_state.default ~round;
     propagating_float_consts;
@@ -193,7 +188,7 @@ let enter_set_of_closures
     } =
   { round;
     typing_env = TE.closure_env typing_env;
-    inlined_debuginfo = Debuginfo.none;
+    inlined_debuginfo = Inlined_debuginfo.none;
     can_inline;
     inlining_state;
     propagating_float_consts;
@@ -461,12 +456,6 @@ let define_code t ~code_id ~code =
   let all_code = Code_id.Map.add code_id code t.all_code in
   { t with typing_env; all_code }
 
-let set_inlined_debuginfo t dbg = { t with inlined_debuginfo = dbg }
-
-let get_inlined_debuginfo t = t.inlined_debuginfo
-
-let add_inlined_debuginfo t dbg = Debuginfo.inline t.inlined_debuginfo dbg
-
 let cse t = t.cse
 
 let comparison_results t = t.comparison_results
@@ -519,6 +508,14 @@ let set_inlining_arguments arguments t =
     inlining_state = Inlining_state.with_arguments arguments t.inlining_state
   }
 
+let set_inlined_debuginfo t ~from =
+  { t with inlined_debuginfo = from.inlined_debuginfo }
+
+let set_inlined_debuginfo' t inlined_debuginfo = { t with inlined_debuginfo }
+
+let add_inlined_debuginfo t dbg =
+  Inlined_debuginfo.rewrite t.inlined_debuginfo dbg
+
 let enter_inlined_apply ~called_code ~apply ~was_inline_always t =
   let arguments =
     Inlining_state.arguments t.inlining_state
@@ -543,8 +540,12 @@ let enter_inlined_apply ~called_code ~apply ~was_inline_always t =
         in
         Inlining_state.increment_depth t.inlining_state ~by)
   in
+  let inlined_debuginfo =
+    Inlined_debuginfo.create ~called_code_id:(Code.code_id called_code)
+      ~apply_dbg:(Apply.dbg apply)
+  in
   { t with
-    inlined_debuginfo = Apply.dbg apply;
+    inlined_debuginfo;
     inlining_state;
     inlining_history_tracker =
       Inlining_history.Tracker.enter_inlined_apply

middle_end/flambda2/simplify/env/downwards_env.mli

@@ -147,11 +147,11 @@ val mem_code : t -> Code_id.t -> bool
 (** This function raises if the code ID is unbound. *)
 val find_code_exn : t -> Code_id.t -> Code_or_metadata.t
 
-val set_inlined_debuginfo : t -> Debuginfo.t -> t
+val set_inlined_debuginfo : t -> from:t -> t
 
-val add_inlined_debuginfo : t -> Debuginfo.t -> Debuginfo.t
+val set_inlined_debuginfo' : t -> Inlined_debuginfo.t -> t
 
-val get_inlined_debuginfo : t -> Debuginfo.t
+val add_inlined_debuginfo : t -> Debuginfo.t -> Debuginfo.t
 
 val round : t -> int
 

middle_end/flambda2/simplify/inlining/inlining_transforms.ml

@@ -20,9 +20,9 @@ module DE = Downwards_env
 module FT = Flambda2_types.Function_type
 module VB = Bound_var
 
-let make_inlined_body ~callee ~unroll_to ~params ~args ~my_closure ~my_region
-    ~my_depth ~rec_info ~body ~exn_continuation ~return_continuation
-    ~apply_exn_continuation ~apply_return_continuation =
+let make_inlined_body ~callee ~called_code_id ~unroll_to ~params ~args
+    ~my_closure ~my_region ~my_depth ~rec_info ~body ~exn_continuation
+    ~return_continuation ~apply_exn_continuation ~apply_return_continuation =
   let callee, rec_info =
     match callee with
     | None ->
@@ -68,10 +68,10 @@ let make_inlined_body ~callee ~unroll_to ~params ~args ~my_closure ~my_region
       ~body ~free_names_of_body:Unknown
     |> Expr.create_let
   in
-  Inlining_helpers.make_inlined_body ~callee ~params ~args ~my_closure
-    ~my_region ~my_depth ~rec_info ~body ~exn_continuation ~return_continuation
-    ~apply_exn_continuation ~apply_return_continuation ~bind_params ~bind_depth
-    ~apply_renaming:Expr.apply_renaming
+  Inlining_helpers.make_inlined_body ~callee ~called_code_id ~params ~args
+    ~my_closure ~my_region ~my_depth ~rec_info ~body ~exn_continuation
+    ~return_continuation ~apply_exn_continuation ~apply_return_continuation
+    ~bind_params ~bind_depth ~apply_renaming:Expr.apply_renaming
 
 let wrap_inlined_body_for_exn_extra_args ~extra_args ~apply_exn_continuation
     ~apply_return_continuation ~result_arity ~make_inlined_body =
@@ -143,7 +143,8 @@ let inline dacc ~apply ~unroll_to ~was_inline_always function_decl =
            ~free_names_of_body:_
          ->
         let make_inlined_body () =
-          make_inlined_body ~callee ~region_inlined_into ~unroll_to
+          make_inlined_body ~callee ~called_code_id:(Code.code_id code)
+            ~region_inlined_into ~unroll_to
             ~params:(Bound_parameters.to_list params)
             ~args ~my_closure ~my_region ~my_depth ~rec_info ~body
             ~exn_continuation ~return_continuation

middle_end/flambda2/simplify/join_points.ml

@@ -223,10 +223,7 @@ let compute_handler_env ?cut_after uses ~is_recursive ~env_at_fork
         (DE.inlining_history_tracker env_at_fork)
         handler_env
     in
-    let handler_env =
-      DE.set_inlined_debuginfo handler_env
-        (DE.get_inlined_debuginfo env_at_fork)
-    in
+    let handler_env = DE.set_inlined_debuginfo handler_env ~from:env_at_fork in
     let handler_env =
       DE.set_at_unit_toplevel_state handler_env
         (DE.at_unit_toplevel env_at_fork)

middle_end/flambda2/simplify/simplify_nullary_primitive.ml

@@ -42,8 +42,7 @@ let simplify_nullary_primitive dacc original_prim (prim : P.nullary_primitive)
     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))
+      DA.map_denv dacc ~f:(fun denv -> DE.set_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

middle_end/flambda2/simplify_shared/inlining_helpers.ml

@@ -17,10 +17,10 @@
 open! Flambda.Import
 module RC = Apply.Result_continuation
 
-let make_inlined_body ~callee ~region_inlined_into ~params ~args ~my_closure
-    ~my_region ~my_depth ~rec_info ~body ~exn_continuation ~return_continuation
-    ~apply_exn_continuation ~apply_return_continuation ~bind_params ~bind_depth
-    ~apply_renaming =
+let make_inlined_body ~callee ~called_code_id:_ ~region_inlined_into ~params
+    ~args ~my_closure ~my_region ~my_depth ~rec_info ~body ~exn_continuation
+    ~return_continuation ~apply_exn_continuation ~apply_return_continuation
+    ~bind_params ~bind_depth ~apply_renaming =
   let renaming = Renaming.empty in
   let renaming =
     match (apply_return_continuation : RC.t) with

middle_end/flambda2/simplify_shared/inlining_helpers.mli

@@ -16,6 +16,7 @@
 
 val make_inlined_body :
   callee:Simple.t option ->
+  called_code_id:Code_id.t ->
   region_inlined_into:Alloc_mode.For_allocations.t ->
   params:'param list ->
   args:Simple.List.t ->

middle_end/flambda2/term_basics/inlined_debuginfo.ml

@@ -0,0 +1,105 @@
+(******************************************************************************
+ *                             flambda-backend                                *
+ *                       Mark Shinwell, Jane Street                           *
+ * -------------------------------------------------------------------------- *
+ *                               MIT License                                  *
+ *                                                                            *
+ * Copyright (c) 2024 Jane Street Group LLC                                   *
+ * [email protected]                                         *
+ *                                                                            *
+ * Permission is hereby granted, free of charge, to any person obtaining a    *
+ * copy of this software and associated documentation files (the "Software"), *
+ * to deal in the Software without restriction, including without limitation  *
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
+ * and/or sell copies of the Software, and to permit persons to whom the      *
+ * Software is furnished to do so, subject to the following conditions:       *
+ *                                                                            *
+ * The above copyright notice and this permission notice shall be included    *
+ * in all copies or substantial portions of the Software.                     *
+ *                                                                            *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER *
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING    *
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER        *
+ * DEALINGS IN THE SOFTWARE.                                                  *
+ ******************************************************************************)
+
+type t =
+  { dbg : Debuginfo.t;
+    function_symbol : Linkage_name.t;
+    uid : string
+  }
+
+(* CR mshinwell/poechsel: if this becomes the flambda2 replacement for
+   [Debuginfo.t] we could maybe avoid making the uids in the same way, and add a
+   "below" constructor or something to identify different instances of
+   inlining. *)
+
+(* CR mshinwell/poechsel: maybe this could be integrated with the inlining
+   histories *)
+
+let print_debuginfo ppf dbg =
+  if Debuginfo.is_none dbg
+  then Format.pp_print_string ppf "None"
+  else Debuginfo.print_compact ppf dbg
+
+let print ppf { dbg; function_symbol; uid } =
+  Format.fprintf ppf
+    "@[<hov 1>(@[<hov 1>(dbg@ %a)@]@ @[<hov 1>(function_symbol@ %a)@]@ @[<hov \
+     1>(uid@ %a)@])@]"
+    print_debuginfo dbg Linkage_name.print function_symbol
+    Format.pp_print_string uid
+
+let none =
+  { dbg = Debuginfo.none;
+    function_symbol = Linkage_name.of_string "";
+    uid = ""
+  }
+
+let is_none t = Debuginfo.compare t.dbg Debuginfo.none = 0
+
+let inlining_counter = ref 0
+
+let create ~called_code_id ~apply_dbg =
+  let function_symbol = Code_id.linkage_name called_code_id in
+  let uid =
+    incr inlining_counter;
+    (* CR mshinwell: consider improving this *)
+    Hashtbl.hash (apply_dbg, function_symbol, !inlining_counter)
+    |> string_of_int
+  in
+  { dbg = apply_dbg; function_symbol; uid }
+
+let rewrite t dbg_from_inlined_body =
+  if is_none t
+  then dbg_from_inlined_body
+  else
+    let dbg_from_inlined_body =
+      (* uids of _all_ [Debuginfo.t] values in the body of the function being
+         inlined get freshened (consistently, for any given inlining).
+
+         See the doc comment in the .mli for further details. *)
+      Debuginfo.mapi_items dbg_from_inlined_body
+        ~f:(fun n ({ dinfo_function_symbol; _ } as item : Debuginfo.item) ->
+          let dinfo_function_symbol =
+            if n = 0
+            then Some (Linkage_name.to_string t.function_symbol)
+            else dinfo_function_symbol
+          in
+          Debuginfo.item_with_uid_and_function_symbol item
+            ~dinfo_uid:(Some t.uid) ~dinfo_function_symbol)
+    in
+    (* Note that [t.dbg] might not be a singleton -- however everything should
+       be ok in terms of each frame (except the first) being annotated with a
+       uid and function symbol in [t.dbg]. *)
+    Debuginfo.inline t.dbg dbg_from_inlined_body
+
+let compare t1 t2 =
+  let c = Debuginfo.compare t1.dbg t2.dbg in
+  if c <> 0
+  then c
+  else
+    let c = Linkage_name.compare t1.function_symbol t2.function_symbol in
+    if c <> 0 then c else String.compare t1.uid t2.uid