Add modes on parameters and a framework for attributes on them

middle_end/closure/closure.ml

@@ -868,7 +868,7 @@ let warning_if_forced_inlined ~loc ~attribute warning =
     Location.prerr_warning (Debuginfo.Scoped_location.to_location loc)
       (Warnings.Inlining_impossible warning)
 
-let fail_if_probe ~probe msg =
+let fail_if_probe ~(probe : Lambda.probe) msg =
   match probe with
   | None -> ()
   | Some {name} ->
@@ -1080,8 +1080,29 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
           List.fold_left (fun kinds (arg, _, kind) -> V.Map.add arg kind kinds)
             kinds first_args
         in
+        let new_clos_mode, kind =
+          (* If the closure has a local suffix, and we've supplied
+             enough args to hit it, then the closure must be local
+             (because the args or closure might be). *)
+          let nparams = List.length params_layout in
+          assert (nparams >= nlocal);
+          let heap_params = nparams - nlocal in
+          if nargs <= heap_params then
+            alloc_heap, Curried {nlocal}
+          else
+            let supplied_local_args = nargs - heap_params in
+            alloc_local, Curried {nlocal = nlocal - supplied_local_args}
+        in
+        if is_local_mode clos_mode then assert (is_local_mode new_clos_mode);
+        (* CR ncourant: mode = new_clos_mode is incorrect; but the modes will
+           not be used for anything, so it is fine here. *)
         let final_args =
-          List.map (fun kind -> V.create_local "arg", kind) rem_layouts
+          List.map (fun kind -> {
+                name = V.create_local "arg";
+                layout = kind;
+                attributes = Lambda.default_param_attribute;
+                mode = new_clos_mode
+              }) rem_layouts
         in
         let rec iter args body =
           match args with
@@ -1092,25 +1113,11 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
         in
         let internal_args =
           (List.map (fun (arg1, _arg2, _) -> Lvar arg1) first_args)
-          @ (List.map (fun (arg, _) -> Lvar arg ) final_args)
+          @ (List.map (fun p -> Lvar p.name ) final_args)
         in
         let funct_var = V.create_local "funct" in
         let fenv = V.Map.add funct_var fapprox fenv in
         let kinds = V.Map.add funct_var Lambda.layout_function kinds in
-        let new_clos_mode, kind =
-          (* If the closure has a local suffix, and we've supplied
-             enough args to hit it, then the closure must be local
-             (because the args or closure might be). *)
-          let nparams = List.length params_layout in
-          assert (nparams >= nlocal);
-          let heap_params = nparams - nlocal in
-          if nargs <= heap_params then
-            alloc_heap, Curried {nlocal}
-          else
-            let supplied_local_args = nargs - heap_params in
-            alloc_local, Curried {nlocal = nlocal - supplied_local_args}
-        in
-        if is_local_mode clos_mode then assert (is_local_mode new_clos_mode);
         let ret_mode = if fundesc.fun_region then alloc_heap else alloc_local in
         let (new_fun, approx) =
           close { backend; fenv; cenv; mutable_vars; kinds; catch_env }
@@ -1525,7 +1532,7 @@ and close_functions { backend; fenv; cenv; mutable_vars; kinds; catch_env } fun_
               {fun_label = label;
                fun_arity = {
                  function_kind = kind ;
-                 params_layout = List.map snd params ;
+                 params_layout = List.map (fun p -> p.layout) params ;
                  return_layout = return
                };
                fun_closed = initially_closed;
@@ -1534,7 +1541,8 @@ and close_functions { backend; fenv; cenv; mutable_vars; kinds; catch_env } fun_
                fun_poll = attr.poll;
                fun_region = region} in
             let dbg = Debuginfo.from_location loc in
-            (id, params, return, body, mode, attrib, fundesc, dbg)
+            (id, List.map (fun (p : Lambda.lparam) -> let No_attributes = p.attributes in (p.name, p.layout, p.mode)) params,
+             return, body, mode, attrib, fundesc, dbg)
         | (_, _) -> fatal_error "Closure.close_functions")
       fun_defs in
   (* Build an approximate fenv for compiling the functions *)
@@ -1584,7 +1592,7 @@ and close_functions { backend; fenv; cenv; mutable_vars; kinds; catch_env } fun_
     in
     let kinds_body =
       List.fold_right
-        (fun (id, kind) kinds -> V.Map.add id kind kinds)
+        (fun (id, kind, _) kinds -> V.Map.add id kind kinds)
         params (V.Map.add env_param Lambda.layout_function kinds_rec)
     in
     let (ubody, approx) =
@@ -1602,13 +1610,13 @@ and close_functions { backend; fenv; cenv; mutable_vars; kinds; catch_env } fun_
     let fun_params =
       if !useless_env
       then params
-      else params @ [env_param, Lambda.layout_function]
+      else params @ [env_param, Lambda.layout_function, alloc_heap]
     in
     let f =
       {
         label  = fundesc.fun_label;
         arity  = fundesc.fun_arity;
-        params = List.map (fun (var, _) -> VP.create var) fun_params;
+        params = List.map (fun (var, _, _) -> VP.create var) fun_params;
         body   = ubody;
         dbg;
         env = Some env_param;
@@ -1621,7 +1629,7 @@ and close_functions { backend; fenv; cenv; mutable_vars; kinds; catch_env } fun_
        their wrapper functions) to be inlined *)
     let n =
       List.fold_left
-        (fun n (id, _) -> n + if V.name id = "*opt*" then 8 else 1)
+        (fun n (id, _, _) -> n + if V.name id = "*opt*" then 8 else 1)
         0
         fun_params
     in
@@ -1637,7 +1645,7 @@ and close_functions { backend; fenv; cenv; mutable_vars; kinds; catch_env } fun_
       | Never_inline -> min_int
       | Unroll _ -> assert false
     in
-    let fun_params = List.map (fun (var, _) -> VP.create var) fun_params in
+    let fun_params = List.map (fun (var, _, _) -> VP.create var) fun_params in
     if lambda_smaller ubody threshold
     then fundesc.fun_inline <- Some(fun_params, ubody);
 

middle_end/flambda/closure_conversion.ml

@@ -226,6 +226,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
     (* CR-soon mshinwell: some of this is now very similar to the let rec case
        below *)
     let set_of_closures_var = Variable.create Names.set_of_closures in
+    let params = List.map (fun (p : Lambda.lparam) -> let No_attributes = p.attributes in (p.name, p.layout)) params in
     let set_of_closures =
       let decl =
         Function_decl.create ~let_rec_ident:None ~closure_bound_var ~kind ~mode
@@ -279,6 +280,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
               let debug_info = Debuginfo.from_location loc in
               Variable.create_with_same_name_as_ident ~debug_info let_rec_ident
             in
+            let params = List.map (fun (p : Lambda.lparam) -> let No_attributes = p.attributes in (p.name, p.layout)) params in
             let function_declaration =
               Function_decl.create ~let_rec_ident:(Some let_rec_ident)
                 ~closure_bound_var ~kind ~mode ~region
@@ -623,6 +625,7 @@ and close_functions t external_env function_declarations : Flambda.named =
          Misc.fatal_error "Closure_conversion: Tupled Alloc_local function found"
     in
     let params = List.mapi (fun i (v, kind) ->
+        (* CR ncourant: actually now we have the alloc_mode in lambda, propagate it *)
       let alloc_mode =
         if i < nheap then Lambda.alloc_heap else Lambda.alloc_local
       in
@@ -699,6 +702,7 @@ and close_let_bound_expression t ?let_rec_ident let_bound_var env
       let debug_info = Debuginfo.from_location loc in
       Variable.rename ~debug_info let_bound_var
     in
+    let params = List.map (fun (p : Lambda.lparam) -> let No_attributes = p.attributes in (p.name, p.layout)) params in
     let decl =
       Function_decl.create ~let_rec_ident ~closure_bound_var ~kind ~mode ~region
         ~params ~body ~attr ~loc ~return_layout:return

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1337,6 +1337,11 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
   let loc = Function_decl.loc decl in
   let dbg = Debuginfo.from_location loc in
   let params = Function_decl.params decl in
+  let param_modes =
+    List.map
+      (fun (p : Function_decl.param) -> Alloc_mode.For_types.from_lambda p.mode)
+      params
+  in
   let return = Function_decl.return decl in
   let return_continuation = Function_decl.return_continuation decl in
   let acc, exn_continuation =
@@ -1456,8 +1461,8 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
   in
   let closure_env =
     List.fold_right
-      (fun (id, kind) env ->
-        let env, _var = Env.add_var_like env id User_visible kind in
+      (fun (p : Function_decl.param) env ->
+        let env, _var = Env.add_var_like env p.name User_visible p.kind in
         env)
       params closure_env
   in
@@ -1482,11 +1487,12 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
   (* CR-someday pchambart: eta-expansion wrappers for primitives are not marked
      as stubs but certainly should be. *)
   let stub = Function_decl.stub decl in
-  let param_vars =
-    List.map (fun (id, kind) -> fst (Env.find_var closure_env id), kind) params
-  in
   let params =
-    List.map (fun (var, kind) -> BP.create var kind) param_vars
+    List.map
+      (fun (p : Function_decl.param) ->
+        let var = fst (Env.find_var closure_env p.name) in
+        BP.create var p.kind)
+      params
     |> Bound_parameters.create
   in
   let acc = Acc.with_seen_a_function acc false in
@@ -1606,6 +1612,7 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
   let code =
     Code.create code_id ~params_and_body
       ~free_names_of_params_and_body:(Acc.free_names acc) ~params_arity
+      ~param_modes
       ~first_complex_local_param:(Function_decl.first_complex_local_param decl)
       ~result_arity:return ~result_types:Unknown
       ~contains_no_escaping_local_allocs:
@@ -1706,7 +1713,7 @@ let close_functions acc external_env ~current_region function_declarations =
     List.fold_left
       (fun approx_map decl ->
         (* The only fields of metadata which are used for this pass are
-           params_arity, is_tupled, first_complex_local_param,
+           params_arity, param_modes, is_tupled, first_complex_local_param,
            contains_no_escaping_local_allocs, and result_arity. We try to
            populate the different fields as much as possible, but put dummy
            values when they are not yet computed or simply too expensive to
@@ -1715,7 +1722,14 @@ let close_functions acc external_env ~current_region function_declarations =
         let code_id = Function_slot.Map.find function_slot function_code_ids in
         let params = Function_decl.params decl in
         let params_arity =
-          List.map (fun (_, kind) -> kind) params |> Flambda_arity.create
+          List.map (fun (p : Function_decl.param) -> p.kind) params
+          |> Flambda_arity.create
+        in
+        let param_modes =
+          List.map
+            (fun (p : Function_decl.param) ->
+              Alloc_mode.For_types.from_lambda p.mode)
+            params
         in
         let result_arity = Function_decl.return decl in
         let poll_attribute =
@@ -1735,7 +1749,7 @@ let close_functions acc external_env ~current_region function_declarations =
           Code_metadata.create code_id ~params_arity
             ~first_complex_local_param:
               (Function_decl.first_complex_local_param decl)
-            ~result_arity ~result_types:Unknown
+            ~param_modes ~result_arity ~result_types:Unknown
             ~contains_no_escaping_local_allocs:
               (Function_decl.contains_no_escaping_local_allocs decl)
             ~stub:(Function_decl.stub decl) ~inline:Never_inline ~check
@@ -2017,8 +2031,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 ~result_arity ~first_complex_local_param
-    ~contains_no_escaping_local_allocs =
+    approx ~provided ~missing_arity ~missing_param_modes ~result_arity
+    ~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
@@ -2030,16 +2044,21 @@ 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_with_subkind ->
-        ( Ident.create_local ("param" ^ string_of_int (num_provided + n)),
-          kind_with_subkind ))
-      (Flambda_arity.to_list missing_arity)
+      (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 = Alloc_mode.For_types.to_lambda mode
+        })
+      (List.combine (Flambda_arity.to_list missing_arity) missing_param_modes)
   in
   let return_continuation = Continuation.create ~sort:Return () in
   let exn_continuation =
     IR.{ exn_handler = Continuation.create (); extra_args = [] }
   in
-  let all_args = provided @ List.map (fun (a, _) -> IR.Var a) params in
+  let all_args =
+    provided @ List.map (fun (p : Function_decl.param) -> IR.Var p.name) params
+  in
   let result_mode =
     if contains_no_escaping_local_allocs
     then Lambda.alloc_heap
@@ -2218,6 +2237,7 @@ type call_args_split =
   | Partial_app of
       { provided : IR.simple list;
         missing_arity : Flambda_arity.t;
+        missing_param_modes : Alloc_mode.For_types.t list;
         result_arity : Flambda_arity.t
       }
   | Over_app of
@@ -2237,6 +2257,7 @@ let close_apply acc env (apply : IR.apply) : Expr_with_acc.t =
         ( Code_metadata.params_arity metadata,
           Code_metadata.result_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
@@ -2255,6 +2276,7 @@ let close_apply acc env (apply : IR.apply) : Expr_with_acc.t =
       ( params_arity,
         result_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
@@ -2279,9 +2301,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_param_modes;
               result_arity
             }
         else
@@ -2305,7 +2329,8 @@ 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; result_arity } ->
+    | Partial_app { provided; missing_arity; missing_param_modes; result_arity }
+      ->
       (match apply.inlined with
       | Always_inlined | Unroll _ ->
         Location.prerr_warning
@@ -2315,8 +2340,8 @@ 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 ~result_arity ~first_complex_local_param
-        ~contains_no_escaping_local_allocs
+        ~missing_arity ~missing_param_modes ~result_arity
+        ~first_complex_local_param ~contains_no_escaping_local_allocs
     | Over_app { full; remaining; remaining_arity } ->
       let full_args_call apply_continuation ~region acc =
         let mode =

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

@@ -672,11 +672,18 @@ end
 
 module Function_decls = struct
   module Function_decl = struct
+    type param =
+      { name : Ident.t;
+        kind : Flambda_kind.With_subkind.t;
+        attributes : Lambda.parameter_attribute;
+        mode : Lambda.alloc_mode
+      }
+
     type t =
       { let_rec_ident : Ident.t;
         function_slot : Function_slot.t;
         kind : Lambda.function_kind;
-        params : (Ident.t * Flambda_kind.With_subkind.t) list;
+        params : param list;
         return : Flambda_arity.t;
         return_continuation : Continuation.t;
         exn_continuation : IR.exn_continuation;
@@ -807,7 +814,7 @@ module Function_decls = struct
     set_diff
       (set_diff
          (all_free_idents function_decls)
-         (List.map fst (all_params function_decls)))
+         (List.map (fun p -> p.Function_decl.name) (all_params function_decls)))
       (let_rec_idents function_decls)
 
   let create function_decls alloc_mode =

middle_end/flambda2/from_lambda/closure_conversion_aux.mli

@@ -293,11 +293,18 @@ module Function_decls : sig
   module Function_decl : sig
     type t
 
+    type param =
+      { name : Ident.t;
+        kind : Flambda_kind.With_subkind.t;
+        attributes : Lambda.parameter_attribute;
+        mode : Lambda.alloc_mode
+      }
+
     val create :
       let_rec_ident:Ident.t option ->
       function_slot:Function_slot.t ->
       kind:Lambda.function_kind ->
-      params:(Ident.t * Flambda_kind.With_subkind.t) list ->
+      params:param list ->
       return:Flambda_arity.t ->
       return_continuation:Continuation.t ->
       exn_continuation:IR.exn_continuation ->
@@ -318,7 +325,7 @@ module Function_decls : sig
 
     val kind : t -> Lambda.function_kind
 
-    val params : t -> (Ident.t * Flambda_kind.With_subkind.t) list
+    val params : t -> param list
 
     val return : t -> Flambda_arity.t