Fix simplification of never_return applications (#2142)

Author: Ekdohibs

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 =

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

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