diff --git a/middle_end/flambda2/simplify/expr_builder.ml b/middle_end/flambda2/simplify/expr_builder.ml index 63d931e6fcc..66f10bfeffd 100644 --- a/middle_end/flambda2/simplify/expr_builder.ml +++ b/middle_end/flambda2/simplify/expr_builder.ml @@ -830,9 +830,14 @@ let rewrite_fixed_arity_apply uacc ~use_id arity apply = in uacc, RE.create_apply (UA.are_rebuilding_terms uacc) apply in - match Apply.continuation apply with - | Never_returns -> make_apply apply - | Return cont -> + match use_id, Apply.continuation apply with + | _, Never_returns -> make_apply apply + | None, Return _ -> + Misc.fatal_errorf + "Expr_builder.rewrite_fixed_arity_apply: got no use_id for the return continuation but the apply \ + could return:@ %a@." + Apply.print apply + | Some use_id, Return cont -> rewrite_fixed_arity_continuation uacc cont ~use_id arity ~around:(fun uacc return_cont -> let exn_cont = diff --git a/middle_end/flambda2/simplify/expr_builder.mli b/middle_end/flambda2/simplify/expr_builder.mli index 1746bb8b512..3797a304fd7 100644 --- a/middle_end/flambda2/simplify/expr_builder.mli +++ b/middle_end/flambda2/simplify/expr_builder.mli @@ -138,7 +138,7 @@ val rewrite_switch_arm : val rewrite_fixed_arity_apply : Upwards_acc.t -> - use_id:Apply_cont_rewrite_id.t -> + use_id:Apply_cont_rewrite_id.t option -> [`Unarized] Flambda_arity.t -> Apply.t -> Upwards_acc.t * Rebuilt_expr.t diff --git a/middle_end/flambda2/simplify/simplify_apply_expr.ml b/middle_end/flambda2/simplify/simplify_apply_expr.ml index 51fc93f2335..716d31b3c92 100644 --- a/middle_end/flambda2/simplify/simplify_apply_expr.ml +++ b/middle_end/flambda2/simplify/simplify_apply_expr.ml @@ -179,15 +179,7 @@ let rebuild_non_inlined_direct_full_application apply ~use_id ~exn_cont_use_id * in *) let apply = if erase_callee then Apply.erase_callee apply else apply in let uacc, expr = - match use_id with - | None -> - let uacc = - UA.add_free_names uacc (Apply.free_names apply) - |> UA.notify_added ~code_size:(Code_size.apply apply) - in - uacc, RE.create_apply (UA.are_rebuilding_terms uacc) apply - | Some use_id -> - EB.rewrite_fixed_arity_apply uacc ~use_id result_arity apply + EB.rewrite_fixed_arity_apply uacc ~use_id result_arity apply in after_rebuild expr uacc @@ -861,12 +853,6 @@ let rebuild_function_call_where_callee's_type_unavailable apply call_kind let simplify_function_call_where_callee's_type_unavailable dacc apply (call : Call_kind.Function_call.t) ~apply_alloc_mode ~down_to_up = fail_if_probe apply; - let cont = - match Apply.continuation apply with - | Never_returns -> - Misc.fatal_error "cannot simplify an application that never returns" - | Return continuation -> continuation - in let denv = DA.denv dacc in if Are_rebuilding_terms.are_rebuilding (DE.are_rebuilding_terms denv) then @@ -875,6 +861,18 @@ let simplify_function_call_where_callee's_type_unavailable dacc apply ~tracker:(DE.inlining_history_tracker denv) ~apply (); let env_at_use = denv in + let dacc, use_id = + match Apply.continuation apply with + | Never_returns -> dacc, None + | Return continuation -> + let dacc, use_id = + DA.record_continuation_use dacc continuation + (Non_inlinable { escaping = true }) + ~env_at_use + ~arg_types:(T.unknown_types_from_arity (Apply.return_arity apply)) + in + dacc, Some use_id + in let dacc, exn_cont_use_id = DA.record_continuation_use dacc (Exn_continuation.exn_handler (Apply.exn_continuation apply)) @@ -884,12 +882,6 @@ let simplify_function_call_where_callee's_type_unavailable dacc apply (T.unknown_types_from_arity (Exn_continuation.arity (Apply.exn_continuation apply))) in - let dacc, use_id = - DA.record_continuation_use dacc cont - (Non_inlinable { escaping = true }) - ~env_at_use - ~arg_types:(T.unknown_types_from_arity (Apply.return_arity apply)) - in let call_kind = match call with | Indirect_unknown_arity -> @@ -903,8 +895,7 @@ let simplify_function_call_where_callee's_type_unavailable dacc apply Call_kind.indirect_function_call_known_arity apply_alloc_mode in let dacc = - record_free_names_of_apply_as_used ~use_id:(Some use_id) ~exn_cont_use_id - dacc apply + record_free_names_of_apply_as_used ~use_id ~exn_cont_use_id dacc apply in down_to_up dacc ~rebuild: @@ -1049,7 +1040,8 @@ let rebuild_method_call apply ~use_id ~exn_cont_use_id uacc ~after_rebuild = apply in let uacc, expr = - EB.rewrite_fixed_arity_apply uacc ~use_id (Apply.return_arity apply) apply + EB.rewrite_fixed_arity_apply uacc ~use_id:(Some use_id) + (Apply.return_arity apply) apply in after_rebuild expr uacc @@ -1108,15 +1100,7 @@ let rebuild_c_call apply ~use_id ~exn_cont_use_id ~return_arity uacc apply in let uacc, expr = - match use_id with - | Some use_id -> - EB.rewrite_fixed_arity_apply uacc ~use_id return_arity apply - | None -> - let uacc = - UA.add_free_names uacc (Apply.free_names apply) - |> UA.notify_added ~code_size:(Code_size.apply apply) - in - uacc, RE.create_apply (UA.are_rebuilding_terms uacc) apply + EB.rewrite_fixed_arity_apply uacc ~use_id return_arity apply in after_rebuild expr uacc