Merging of Debuginfo.t across CSEd occurrences (#1767)

middle_end/flambda2/simplify/env/downwards_acc.ml

@@ -29,13 +29,14 @@ type t =
     demoted_exn_handlers : Continuation.Set.t;
     code_ids_to_remember : Code_id.Set.t;
     code_ids_to_never_delete : Code_id.Set.t;
-    slot_offsets : Slot_offsets.t Code_id.Map.t
+    slot_offsets : Slot_offsets.t Code_id.Map.t;
+    debuginfo_rewrites : Debuginfo.t Simple.Map.t
   }
 
 let [@ocamlformat "disable"] print ppf
       { denv; continuation_uses_env; shareable_constants; used_value_slots;
         lifted_constants; flow_acc; demoted_exn_handlers; code_ids_to_remember;
-        code_ids_to_never_delete; slot_offsets } =
+        code_ids_to_never_delete; slot_offsets; debuginfo_rewrites } =
   Format.fprintf ppf "@[<hov 1>(\
       @[<hov 1>(denv@ %a)@]@ \
       @[<hov 1>(continuation_uses_env@ %a)@]@ \
@@ -46,7 +47,8 @@ let [@ocamlformat "disable"] print ppf
       @[<hov 1>(demoted_exn_handlers@ %a)@]@ \
       @[<hov 1>(code_ids_to_remember@ %a)@]@ \
       @[<hov 1>(code_ids_to_never_delete@ %a)@]@ \
-      @[<hov 1>(slot_offsets@ %a)@]\
+      @[<hov 1>(slot_offsets@ %a)@ \
+      @[<hov 1>(debuginfo_rewrites@ %a)@]\
       )@]"
     DE.print denv
     CUE.print continuation_uses_env
@@ -58,6 +60,7 @@ let [@ocamlformat "disable"] print ppf
     Code_id.Set.print code_ids_to_remember
     Code_id.Set.print code_ids_to_never_delete
     (Code_id.Map.print Slot_offsets.print) slot_offsets
+    (Simple.Map.print Debuginfo.print_compact) debuginfo_rewrites
 
 let create denv continuation_uses_env =
   { denv;
@@ -69,7 +72,8 @@ let create denv continuation_uses_env =
     flow_acc = Flow.Acc.empty ();
     demoted_exn_handlers = Continuation.Set.empty;
     code_ids_to_remember = Code_id.Set.empty;
-    code_ids_to_never_delete = Code_id.Set.empty
+    code_ids_to_never_delete = Code_id.Set.empty;
+    debuginfo_rewrites = Simple.Map.empty
   }
 
 let denv t = t.denv
@@ -208,3 +212,17 @@ let demoted_exn_handlers t = t.demoted_exn_handlers
 let slot_offsets t = t.slot_offsets
 
 let with_slot_offsets t ~slot_offsets = { t with slot_offsets }
+
+let find_debuginfo_rewrite t ~bound_to =
+  Simple.Map.find_opt bound_to t.debuginfo_rewrites
+
+let merge_debuginfo_rewrite t ~bound_to dbg =
+  let dbg =
+    match find_debuginfo_rewrite t ~bound_to with
+    | None -> dbg
+    | Some earlier_dbg -> Debuginfo.merge ~into:earlier_dbg dbg
+  in
+  { t with
+    debuginfo_rewrites =
+      Simple.Map.add (* or replace *) bound_to dbg t.debuginfo_rewrites
+  }

middle_end/flambda2/simplify/env/downwards_acc.mli

@@ -103,3 +103,7 @@ val are_rebuilding_terms : t -> Are_rebuilding_terms.t
 val slot_offsets : t -> Slot_offsets.t Code_id.Map.t
 
 val with_slot_offsets : t -> slot_offsets:Slot_offsets.t Code_id.Map.t -> t
+
+val merge_debuginfo_rewrite : t -> bound_to:Simple.t -> Debuginfo.t -> t
+
+val find_debuginfo_rewrite : t -> bound_to:Simple.t -> Debuginfo.t option

middle_end/flambda2/simplify/simplify_let_expr.ml

@@ -74,6 +74,30 @@ let rebuild_let simplify_named_result removed_operations ~rewrite_id
   let bindings =
     Simplify_named_result.bindings_to_place simplify_named_result
   in
+  let bindings =
+    List.map
+      (fun ({ Expr_builder.let_bound;
+              simplified_defining_expr;
+              original_defining_expr = _
+            } as binding) ->
+        match simplified_defining_expr.named with
+        | Prim (prim, _dbg) -> (
+          match Bound_pattern.must_be_singleton_opt let_bound with
+          | None -> binding
+          | Some bound_var -> (
+            let simple = Simple.var (Bound_var.var bound_var) in
+            match
+              DA.find_debuginfo_rewrite (UA.creation_dacc uacc) ~bound_to:simple
+            with
+            | None -> binding
+            | Some dbg ->
+              { binding with
+                simplified_defining_expr =
+                  Simplified_named.create (Named.create_prim prim dbg)
+              }))
+        | Simple _ | Set_of_closures _ | Rec_info _ -> binding)
+      bindings
+  in
   let uacc, bindings =
     let Flow_types.Mutable_unboxing_result.{ let_rewrites; _ } =
       UA.mutable_unboxing_result uacc

middle_end/flambda2/simplify/simplify_primitive.ml

@@ -36,7 +36,7 @@ let apply_cse dacc ~original_prim =
           Simple.print simple
       | canonical -> Some canonical))
 
-let try_cse dacc ~original_prim ~min_name_mode ~result_var : cse_result =
+let try_cse dacc dbg ~original_prim ~min_name_mode ~result_var : cse_result =
   (* CR-someday mshinwell: Use [meet] and [reify] for CSE? (discuss with
      lwhite) *)
   (* CR-someday mshinwell: Find example that suggested we needed to allow
@@ -47,6 +47,7 @@ let try_cse dacc ~original_prim ~min_name_mode ~result_var : cse_result =
     let result_var' = VB.var result_var in
     match apply_cse dacc ~original_prim with
     | Some replace_with ->
+      let dacc = DA.merge_debuginfo_rewrite dacc ~bound_to:replace_with dbg in
       let named = Named.create_simple replace_with in
       let ty = T.alias_type_of (P.result_kind' original_prim) replace_with in
       let dacc = DA.add_variable dacc result_var ty in
@@ -70,8 +71,11 @@ let try_cse dacc ~original_prim ~min_name_mode ~result_var : cse_result =
         | None -> dacc
         | Some eligible_prim ->
           let bound_to = Simple.var result_var' in
-          DA.map_denv dacc ~f:(fun denv ->
-              DE.add_cse denv eligible_prim ~bound_to)
+          let dacc =
+            DA.map_denv dacc ~f:(fun denv ->
+                DE.add_cse denv eligible_prim ~bound_to)
+          in
+          DA.merge_debuginfo_rewrite dacc ~bound_to dbg
       in
       Not_applied dacc
 
@@ -103,7 +107,7 @@ let simplify_primitive dacc (prim : P.t) dbg ~result_var =
     let original_prim : P.t =
       if orig_arg == arg then prim else Unary (unary_prim, arg)
     in
-    match try_cse dacc ~original_prim ~min_name_mode ~result_var with
+    match try_cse dacc dbg ~original_prim ~min_name_mode ~result_var with
     | Applied result -> result
     | Not_applied dacc ->
       Simplify_unary_primitive.simplify_unary_primitive dacc original_prim
@@ -122,7 +126,7 @@ let simplify_primitive dacc (prim : P.t) dbg ~result_var =
       then prim
       else Binary (binary_prim, arg1, arg2)
     in
-    match try_cse dacc ~original_prim ~min_name_mode ~result_var with
+    match try_cse dacc dbg ~original_prim ~min_name_mode ~result_var with
     | Applied result -> result
     | Not_applied dacc ->
       Simplify_binary_primitive.simplify_binary_primitive dacc original_prim
@@ -146,7 +150,7 @@ let simplify_primitive dacc (prim : P.t) dbg ~result_var =
       then prim
       else Ternary (ternary_prim, arg1, arg2, arg3)
     in
-    match try_cse dacc ~original_prim ~min_name_mode ~result_var with
+    match try_cse dacc dbg ~original_prim ~min_name_mode ~result_var with
     | Applied result -> result
     | Not_applied dacc ->
       Simplify_ternary_primitive.simplify_ternary_primitive dacc original_prim
@@ -172,7 +176,7 @@ let simplify_primitive dacc (prim : P.t) dbg ~result_var =
     let original_prim : P.t =
       Variadic (variadic_prim, List.map fst args_with_tys)
     in
-    match try_cse dacc ~original_prim ~min_name_mode ~result_var with
+    match try_cse dacc dbg ~original_prim ~min_name_mode ~result_var with
     | Applied result -> result
     | Not_applied dacc ->
       Simplify_variadic_primitive.simplify_variadic_primitive dacc original_prim

ocaml/lambda/debuginfo.ml

@@ -272,3 +272,6 @@ let print_compact ppf { dbg; } = print_compact ppf dbg
 let to_list { dbg; } = dbg
 
 let length { dbg; } = List.length dbg
+
+let merge ~into:{ dbg = dbg1 } { dbg = _dbg2; } =
+  { dbg = dbg1 }

ocaml/lambda/debuginfo.mli

@@ -97,3 +97,5 @@ val print_compact : Format.formatter -> t -> unit
 val to_list : t -> item list
 
 val length : t -> int
+
+val merge : into:t -> t -> t