Set inner continuations uses of let rec to non-inlinable

middle_end/flambda2/simplify/env/continuation_uses.ml

@@ -103,3 +103,6 @@ let get_typing_env_no_more_than_one_use t =
   | _ :: _ ->
     Misc.fatal_errorf "Only zero or one continuation use(s) expected:@ %a" print
       t
+
+let mark_non_inlinable t =
+  { t with uses = List.map U.mark_non_inlinable t.uses }

middle_end/flambda2/simplify/env/continuation_uses.mli

@@ -52,3 +52,5 @@ val get_typing_env_no_more_than_one_use :
   t -> Flambda2_types.Typing_env.t option
 
 val union : t -> t -> t
+
+val mark_non_inlinable : t -> t

middle_end/flambda2/simplify/env/continuation_uses_env.ml

@@ -71,3 +71,9 @@ let remove t cont =
   { continuation_uses = Continuation.Map.remove cont t.continuation_uses }
 
 let delete_continuation_uses = remove
+
+let mark_non_inlinable { continuation_uses } =
+  let continuation_uses =
+    Continuation.Map.map Continuation_uses.mark_non_inlinable continuation_uses
+  in
+  { continuation_uses }

middle_end/flambda2/simplify/env/continuation_uses_env.mli

@@ -27,3 +27,5 @@ val get_continuation_uses : t -> Continuation.t -> Continuation_uses.t option
 val remove : t -> Continuation.t -> t
 
 val union : t -> t -> t
+
+val mark_non_inlinable : t -> t

middle_end/flambda2/simplify/env/one_continuation_use.ml

@@ -40,3 +40,8 @@ let use_kind t = t.kind
 let arg_types t = t.arg_types
 
 let env_at_use t = t.env
+
+let mark_non_inlinable t =
+  match t.kind with
+  | Inlinable -> { t with kind = Non_inlinable { escaping = false } }
+  | Non_inlinable _ -> t

middle_end/flambda2/simplify/env/one_continuation_use.mli

@@ -35,3 +35,5 @@ val use_kind : t -> Continuation_use_kind.t
 val arg_types : t -> T.t list
 
 val env_at_use : t -> DE.t
+
+val mark_non_inlinable : t -> t

middle_end/flambda2/simplify/simplify_let_cont_expr.ml

@@ -833,12 +833,18 @@ let prepare_to_rebuild_one_recursive_let_cont_handler cont params
 
 let after_downwards_traversal_of_one_recursive_let_cont_handler cont
     unboxing_decisions ~down_to_up params ~original_cont_scope
-    ~cont_uses_in_body dacc ~rebuild:rebuild_handler =
+    ~cont_uses_in_body ~body_continuation_uses_env dacc ~rebuild:rebuild_handler
+    =
   let dacc = DA.map_data_flow dacc ~f:(Data_flow.exit_continuation cont) in
+  let handler_continuation_uses_env = DA.continuation_uses_env dacc in
+  let continuation_uses_env =
+    CUE.union body_continuation_uses_env
+      (CUE.mark_non_inlinable handler_continuation_uses_env)
+  in
   let arg_types_by_use_id =
     (* At this point all uses (in both the body and the handler) of [cont] are
        in [dacc]. *)
-    match CUE.get_continuation_uses (DA.continuation_uses_env dacc) cont with
+    match CUE.get_continuation_uses continuation_uses_env cont with
     | None ->
       ListLabels.map (Bound_parameters.to_list params) ~f:(fun _ ->
           Apply_cont_rewrite_id.Map.empty)
@@ -855,7 +861,7 @@ let after_downwards_traversal_of_one_recursive_let_cont_handler cont
     DA.map_data_flow dacc ~f:(fun data_flow ->
         Data_flow.add_extra_params_and_args cont extra_params_and_args data_flow)
   in
-  let cont_uses_env = CUE.remove (DA.continuation_uses_env dacc) cont in
+  let cont_uses_env = CUE.remove continuation_uses_env cont in
   let dacc = DA.with_continuation_uses_env dacc ~cont_uses_env in
   down_to_up dacc
     ~rebuild:
@@ -891,14 +897,16 @@ let simplify_recursive_let_cont_handlers ~simplify_expr ~denv_before_body
   let dacc =
     DA.map_denv dacc ~f:(fun denv -> DE.set_at_unit_toplevel_state denv false)
   in
+  let body_continuation_uses_env = DA.continuation_uses_env dacc in
   let cont_uses_in_body =
-    CUE.get_continuation_uses (DA.continuation_uses_env dacc) cont
+    CUE.get_continuation_uses body_continuation_uses_env cont
   in
   match cont_uses_in_body with
   | None ->
     let rebuild uacc ~after_rebuild = after_rebuild None uacc in
     down_to_up dacc ~rebuild
   | Some cont_uses_in_body ->
+    let dacc = DA.with_continuation_uses_env dacc ~cont_uses_env:CUE.empty in
     let arg_types_by_use_id_in_body =
       Continuation_uses.get_arg_types_by_use_id cont_uses_in_body
     in
@@ -933,7 +941,7 @@ let simplify_recursive_let_cont_handlers ~simplify_expr ~denv_before_body
       ~down_to_up:
         (after_downwards_traversal_of_one_recursive_let_cont_handler cont
            unboxing_decisions params ~original_cont_scope ~down_to_up
-           ~cont_uses_in_body)
+           ~cont_uses_in_body ~body_continuation_uses_env)
 
 let rebuild_recursive_let_cont_expr are_rebuilding_terms ~body
     ~free_names_of_body ~handlers =