correctly propagate param_modes in partial applications

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -2024,8 +2024,8 @@ let close_let_rec acc env ~function_declarations
       named ~body
 
 let wrap_partial_application acc env apply_continuation (apply : IR.apply)
-    approx ~provided ~missing_arity ~first_complex_local_param
-    ~contains_no_escaping_local_allocs =
+    approx ~provided ~missing_arity ~missing_param_modes
+    ~first_complex_local_param ~contains_no_escaping_local_allocs =
   (* In case of partial application, creates a wrapping function from scratch to
      allow inlining and lifting *)
   let wrapper_id = Ident.create_local ("partial_" ^ Ident.name apply.func) in
@@ -2037,13 +2037,13 @@ let wrap_partial_application acc env apply_continuation (apply : IR.apply)
   let num_provided = List.length provided in
   let params =
     List.mapi
-      (fun n kind : Function_decl.param ->
+      (fun n (kind, mode) : Function_decl.param ->
         { name = Ident.create_local ("param" ^ string_of_int (num_provided + n));
           kind;
           attributes = Lambda.default_param_attribute;
-          mode = failwith "use param_modes"
+          mode = Alloc_mode.For_types.to_lambda mode
         })
-      (Flambda_arity.to_list missing_arity)
+      (List.combine (Flambda_arity.to_list missing_arity) missing_param_modes)
   in
   let return_continuation = Continuation.create ~sort:Return () in
   let exn_continuation =
@@ -2229,7 +2229,8 @@ type call_args_split =
   | Exact of IR.simple list
   | Partial_app of
       { provided : IR.simple list;
-        missing_arity : Flambda_arity.t
+        missing_arity : Flambda_arity.t;
+        missing_param_modes : Alloc_mode.For_types.t list
       }
   | Over_app of
       { full : IR.simple list;
@@ -2247,6 +2248,7 @@ let close_apply acc env (apply : IR.apply) : Expr_with_acc.t =
       Some
         ( Code_metadata.params_arity metadata,
           Code_metadata.is_tupled metadata,
+          Code_metadata.param_modes metadata,
           Code_metadata.first_complex_local_param metadata,
           Code_metadata.contains_no_escaping_local_allocs metadata )
     | Value_unknown -> None
@@ -2264,6 +2266,7 @@ let close_apply acc env (apply : IR.apply) : Expr_with_acc.t =
   | Some
       ( arity,
         is_tupled,
+        param_modes,
         first_complex_local_param,
         contains_no_escaping_local_allocs ) -> (
     let acc, args_with_arities = find_simples_and_arity acc env apply.args in
@@ -2288,9 +2291,11 @@ let close_apply acc env (apply : IR.apply) : Expr_with_acc.t =
         else if args_l < arity_l
         then
           let _provided_arity, missing_arity = cut args_l arity in
+          let _provided_modes, missing_param_modes = cut args_l param_modes in
           Partial_app
             { provided = args;
-              missing_arity = Flambda_arity.create missing_arity
+              missing_arity = Flambda_arity.create missing_arity;
+              missing_param_modes
             }
         else
           let full, remaining = cut arity_l args in
@@ -2313,7 +2318,7 @@ let close_apply acc env (apply : IR.apply) : Expr_with_acc.t =
       close_exact_or_unknown_apply acc env
         { apply with args; continuation = apply.continuation }
         (Some approx) ~replace_region:None
-    | Partial_app { provided; missing_arity } ->
+    | Partial_app { provided; missing_arity; missing_param_modes } ->
       (match apply.inlined with
       | Always_inlined | Unroll _ ->
         Location.prerr_warning
@@ -2323,7 +2328,7 @@ let close_apply acc env (apply : IR.apply) : Expr_with_acc.t =
                inlined_attribute_on_partial_application_msg Inlined))
       | Never_inlined | Hint_inlined | Default_inlined -> ());
       wrap_partial_application acc env apply.continuation apply approx ~provided
-        ~missing_arity ~first_complex_local_param
+        ~missing_arity ~missing_param_modes ~first_complex_local_param
         ~contains_no_escaping_local_allocs
     | Over_app { full; remaining; remaining_arity } ->
       let full_args_call apply_continuation ~region acc =

middle_end/flambda2/simplify/simplify_apply_expr.ml

@@ -322,7 +322,8 @@ let simplify_direct_full_application ~simplify_expr dacc apply function_type
 
 let simplify_direct_partial_application ~simplify_expr dacc apply
     ~callee's_code_id ~callee's_code_metadata ~callee's_function_slot
-    ~param_arity ~result_arity ~recursive ~down_to_up ~coming_from_indirect
+    ~param_arity ~param_modes ~result_arity ~recursive ~down_to_up
+    ~coming_from_indirect
     ~(closure_alloc_mode_from_type : Alloc_mode.For_types.t) ~current_region
     ~first_complex_local_param =
   (* Partial-applications are converted in full applications. Let's assume that
@@ -426,11 +427,8 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
         remaining_param_arity
       |> Bound_parameters.create
     in
-    (* CR ncourant: use correct mode *)
-    let remaining_params_alloc_modes =
-      List.map
-        (fun _ -> Alloc_mode.For_allocations.as_type new_closure_alloc_mode)
-        remaining_param_arity
+    let _, remaining_params_alloc_modes =
+      Misc.Stdlib.List.split_at (List.length args) param_modes
     in
     let call_kind =
       Call_kind.direct_function_call callee's_code_id apply_alloc_mode
@@ -778,8 +776,10 @@ let simplify_direct_function_call ~simplify_expr dacc apply
             Apply.print apply;
         simplify_direct_partial_application ~simplify_expr dacc apply
           ~callee's_code_id ~callee's_code_metadata ~callee's_function_slot
-          ~param_arity:params_arity ~result_arity ~recursive ~down_to_up
-          ~coming_from_indirect ~closure_alloc_mode_from_type ~current_region
+          ~param_arity:params_arity
+          ~param_modes:(Code_metadata.param_modes callee's_code_metadata)
+          ~result_arity ~recursive ~down_to_up ~coming_from_indirect
+          ~closure_alloc_mode_from_type ~current_region
           ~first_complex_local_param:
             (Code_metadata.first_complex_local_param callee's_code_metadata))
       else