Improved handling of calls inside values for inlining reports (#626)

Co-authored-by: Luke Maurer <Luke.Maurer@alumni.carleton.edu>

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -337,7 +337,7 @@ module Inlining = struct
             ~result_arity:(Code.result_arity code) ~make_inlined_body)
 end
 
-let close_c_call acc env ~let_bound_var
+let close_c_call acc env ~loc ~let_bound_var
     ({ prim_name;
        prim_arity;
        prim_alloc;
@@ -452,7 +452,8 @@ let close_c_call acc env ~let_bound_var
         Apply.create ~callee ~continuation:(Return return_continuation)
           exn_continuation ~args ~call_kind dbg ~inlined:Default_inlined
           ~inlining_state:(Inlining_state.default ~round:0)
-          ~probe_name:None ~relative_history:(Env.relative_history env)
+          ~probe_name:None
+          ~relative_history:(Env.relative_history_from_scoped ~loc env)
       in
       Expr_with_acc.create_apply acc apply
   in
@@ -567,7 +568,7 @@ let close_primitive acc env ~let_bound_var named (prim : Lambda.primitive) ~args
           IR.print_named named
       | Some exn_continuation -> exn_continuation
     in
-    close_c_call acc env ~let_bound_var prim ~args exn_continuation dbg k
+    close_c_call acc env ~loc ~let_bound_var prim ~args exn_continuation dbg k
   | Pgetglobal id, [] ->
     let is_predef_exn = Ident.is_predef id in
     if not (is_predef_exn || not (Ident.same id (Env.current_unit_id env)))
@@ -896,7 +897,8 @@ let close_exact_or_unknown_apply acc env
       (Debuginfo.from_location loc)
       ~inlined:inlined_call
       ~inlining_state:(Inlining_state.default ~round:0)
-      ~probe_name ~relative_history:(Env.relative_history env)
+      ~probe_name
+      ~relative_history:(Env.relative_history_from_scoped ~loc env)
   in
   if Flambda_features.classic_mode ()
   then
@@ -1632,7 +1634,8 @@ let wrap_over_application acc env full_call (apply : IR.apply) over_args
       Apply.create ~callee:(Simple.var returned_func) ~continuation
         apply_exn_continuation ~args ~call_kind apply_dbg ~inlined
         ~inlining_state:(Inlining_state.default ~round:0)
-        ~probe_name ~relative_history:(Env.relative_history env)
+        ~probe_name
+        ~relative_history:(Env.relative_history_from_scoped ~loc:apply.loc env)
     in
     match needs_region with
     | None -> Expr_with_acc.create_apply acc over_application

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

@@ -345,7 +345,7 @@ module Env = struct
     | Value_unknown -> t
     | approx -> add_value_approximation t alias approx
 
-  let use_path_to_root t path_to_root =
+  let set_path_to_root t path_to_root =
     if path_to_root = Debuginfo.Scoped_location.Loc_unknown
     then t
     else { t with path_to_root }
@@ -358,8 +358,9 @@ module Env = struct
   let inlining_history_tracker { inlining_history_tracker; _ } =
     inlining_history_tracker
 
-  let relative_history { inlining_history_tracker; _ } =
-    Inlining_history.Tracker.relative inlining_history_tracker
+  let relative_history_from_scoped ~loc { path_to_root; _ } =
+    Inlining_history.Relative.between_scoped_locations ~parent:path_to_root
+      ~child:loc
 end
 
 module Acc = struct

middle_end/flambda2/from_lambda/closure_conversion_aux.mli

@@ -158,15 +158,23 @@ module Env : sig
 
   val big_endian : t -> bool
 
-  val use_path_to_root : t -> Debuginfo.Scoped_location.t -> t
+  val set_path_to_root : t -> Debuginfo.Scoped_location.t -> t
 
   val path_to_root : t -> Debuginfo.Scoped_location.t
 
+  (* The inlining tracker is used to ensure that absolute histories are shared
+     between functions defined under the same scope. *)
   val use_inlining_history_tracker : t -> Inlining_history.Tracker.t -> t
 
   val inlining_history_tracker : t -> Inlining_history.Tracker.t
 
-  val relative_history : t -> Inlining_history.Relative.t
+  (* Relative paths are built directly from scoped locations.
+
+     This is fine because when we convert a function call we know that it was
+     never inlined beforehand and thus should inherit a path corresponding to
+     its true location in the source file. *)
+  val relative_history_from_scoped :
+    loc:Debuginfo.Scoped_location.t -> t -> Inlining_history.Relative.t
 end
 
 (** Used to pipe some data through closure conversion *)

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -1627,7 +1627,7 @@ and cps_function env ~fid ~stub ~(recursive : Recursive.t) ?free_idents
       ~name:(Ident.name fid)
   in
   let body acc ccenv =
-    let ccenv = CCenv.use_path_to_root ccenv loc in
+    let ccenv = CCenv.set_path_to_root ccenv loc in
     cps_tail acc new_env ccenv body body_cont body_exn_cont
   in
   Function_decl.create ~let_rec_ident:(Some fid) ~function_slot ~kind ~params

middle_end/flambda2/terms/inlining_history.ml

@@ -161,6 +161,43 @@ module Relative = struct
   let inline prev = Absolute.Inline { prev }
 
   let unknown prev = Absolute.Unknown { prev }
+
+  let between_scoped_locations ~(parent : Debuginfo.Scoped_location.t)
+      ~(child : Debuginfo.Scoped_location.t) =
+    let scopes_are_equal (a : Debuginfo.Scoped_location.scopes)
+        (b : Debuginfo.Scoped_location.scopes) =
+      match a, b with
+      | Cons a, Cons b ->
+        a.item = b.item && a.str = b.str && a.str_fun = b.str_fun
+        && a.name = b.name
+      | Empty, _ | _, Empty -> false
+    in
+    let rec aux ~parent ~child =
+      if scopes_are_equal parent child
+      then Absolute.Empty
+      else
+        match child with
+        | Debuginfo.Scoped_location.Empty -> Absolute.Empty
+        | Debuginfo.Scoped_location.Cons { item; name; prev; _ } -> (
+          let prev = aux ~parent ~child:prev in
+          match item with
+          | Sc_module_definition -> Absolute.Module { name; prev }
+          | Sc_class_definition -> Absolute.Class { name; prev }
+          | Sc_anonymous_function | Sc_method_definition | Sc_value_definition
+          | Sc_lazy | Sc_partial_or_eta_wrapper ->
+            prev)
+    in
+    let parent =
+      match parent with
+      | Loc_unknown -> Debuginfo.Scoped_location.empty_scopes
+      | Loc_known { scopes; _ } -> scopes
+    in
+    let child =
+      match child with
+      | Loc_unknown -> parent
+      | Loc_known { scopes; _ } -> scopes
+    in
+    aux ~parent ~child
 end
 
 let extend_absolute (compilation_unit, absolute) relative =
@@ -223,41 +260,14 @@ module Tracker = struct
 
   let fundecl_of_scoped_location ~name
       ~(path_to_root : Debuginfo.Scoped_location.t)
-      (scoped_location : Debuginfo.Scoped_location.t) ({ relative; _ } as t) =
-    let scopes_are_equal (a : Debuginfo.Scoped_location.scopes)
-        (b : Debuginfo.Scoped_location.scopes) =
-      match a, b with
-      | Cons a, Cons b ->
-        a.item = b.item && a.str = b.str && a.str_fun = b.str_fun
-        && a.name = b.name
-      | Empty, _ | _, Empty -> false
-    in
-    let rec aux ?root ~relative (scopes : Debuginfo.Scoped_location.scopes) =
-      if Option.map (scopes_are_equal scopes) root
-         |> Option.value ~default:false
-      then relative
-      else
-        match scopes with
-        | Empty -> relative
-        | Cons { item; name; prev; _ } -> (
-          let prev = aux ?root ~relative prev in
-          match item with
-          | Sc_module_definition -> Absolute.Module { name; prev }
-          | Sc_class_definition -> Absolute.Class { name; prev }
-          | Sc_anonymous_function | Sc_method_definition | Sc_value_definition
-          | Sc_lazy | Sc_partial_or_eta_wrapper ->
-            prev)
-    in
-    let root =
-      match path_to_root with
-      | Loc_unknown -> None
-      | Loc_known { scopes; _ } -> Some scopes
-    in
-    match scoped_location with
+      (loc : Debuginfo.Scoped_location.t) t =
+    match loc with
     | Loc_unknown -> unknown t
-    | Loc_known { scopes; _ } ->
-      fundecl ~function_relative_history:Relative.empty
-        ~dbg:(Debuginfo.from_location scoped_location)
-        ~name
-        { t with relative = aux ?root ~relative scopes }
+    | Loc_known _ ->
+      let relative =
+        Relative.between_scoped_locations ~parent:path_to_root ~child:loc
+      in
+      fundecl ~function_relative_history:relative
+        ~dbg:(Debuginfo.from_location loc)
+        ~name t
 end

middle_end/flambda2/terms/inlining_history.mli

@@ -107,6 +107,12 @@ module Relative : sig
   val compare : t -> t -> int
 
   val print : Format.formatter -> t -> unit
+
+  (* [between_scoped_location ~parent ~child] returns the relative path between
+     [Absolute.of_scoped_location parent] and [Absolute.of_scope_location
+     child] *)
+  val between_scoped_locations :
+    parent:Debuginfo.Scoped_location.t -> child:Debuginfo.Scoped_location.t -> t
 end
 
 module Tracker : sig