diff --git a/middle_end/flambda2/simplify/env/downwards_acc.ml b/middle_end/flambda2/simplify/env/downwards_acc.ml index aa6ab3cb19f..b9396da809a 100644 --- a/middle_end/flambda2/simplify/env/downwards_acc.ml +++ b/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 "@[(\ @[(denv@ %a)@]@ \ @[(continuation_uses_env@ %a)@]@ \ @@ -46,7 +47,8 @@ let [@ocamlformat "disable"] print ppf @[(demoted_exn_handlers@ %a)@]@ \ @[(code_ids_to_remember@ %a)@]@ \ @[(code_ids_to_never_delete@ %a)@]@ \ - @[(slot_offsets@ %a)@]\ + @[(slot_offsets@ %a)@ \ + @[(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 + } diff --git a/middle_end/flambda2/simplify/env/downwards_acc.mli b/middle_end/flambda2/simplify/env/downwards_acc.mli index 2548f536c51..a3176149622 100644 --- a/middle_end/flambda2/simplify/env/downwards_acc.mli +++ b/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 diff --git a/middle_end/flambda2/simplify/simplify_let_expr.ml b/middle_end/flambda2/simplify/simplify_let_expr.ml index 40516a877ea..8a0b6778bd8 100644 --- a/middle_end/flambda2/simplify/simplify_let_expr.ml +++ b/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 diff --git a/middle_end/flambda2/simplify/simplify_primitive.ml b/middle_end/flambda2/simplify/simplify_primitive.ml index 14ff4ad9002..05c036ab5a6 100644 --- a/middle_end/flambda2/simplify/simplify_primitive.ml +++ b/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 diff --git a/ocaml/lambda/debuginfo.ml b/ocaml/lambda/debuginfo.ml index e36d3c18207..4b19c2e1494 100644 --- a/ocaml/lambda/debuginfo.ml +++ b/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 } diff --git a/ocaml/lambda/debuginfo.mli b/ocaml/lambda/debuginfo.mli index 93d953efb91..a426bbf3e30 100644 --- a/ocaml/lambda/debuginfo.mli +++ b/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