flambda-backend: Add modes on parameters and a framework for attributes on them (#1257)

Author: Ekdohibs

bytecomp/bytegen.ml

@@ -632,7 +632,7 @@ let rec comp_expr env exp sz cont =
       let lbl = new_label() in
       let fv = Ident.Set.elements(free_variables exp) in
       let to_compile =
-        { params = List.map fst params; body = body; label = lbl;
+        { params = List.map (fun p -> p.name) params; body = body; label = lbl;
           free_vars = fv; num_defs = 1; rec_vars = []; rec_pos = 0 } in
       Stack.push to_compile functions_to_compile;
       comp_args env (List.map (fun n -> Lvar n) fv) sz
@@ -655,7 +655,7 @@ let rec comp_expr env exp sz cont =
           | (_id, Lfunction{params; body}) :: rem ->
               let lbl = new_label() in
               let to_compile =
-                { params = List.map fst params; body = body; label = lbl;
+                { params = List.map (fun p -> p.name) params; body = body; label = lbl;
                   free_vars = fv; num_defs = ndecl; rec_vars = rec_idents;
                   rec_pos = pos} in
               Stack.push to_compile functions_to_compile;

lambda/lambda.ml

@@ -537,6 +537,15 @@ type function_attribute = {
 
 type scoped_location = Debuginfo.Scoped_location.t
 
+type parameter_attribute = No_attributes
+
+type lparam = {
+  name : Ident.t;
+  layout : layout;
+  attributes : parameter_attribute;
+  mode : alloc_mode
+}
+
 type lambda =
     Lvar of Ident.t
   | Lmutvar of Ident.t
@@ -568,7 +577,7 @@ type lambda =
 
 and lfunction =
   { kind: function_kind;
-    params: (Ident.t * layout) list;
+    params: lparam list;
     return: layout;
     body: lambda;
     attr: function_attribute; (* specified with [@inline] attribute *)
@@ -715,6 +724,8 @@ let default_function_attribute = {
 let default_stub_attribute =
   { default_function_attribute with stub = true; check = Ignore_assert_all Zero_alloc }
 
+let default_param_attribute = No_attributes
+
 (* Build sharing keys *)
 (*
    Those keys are later compared with Stdlib.compare.
@@ -901,7 +912,7 @@ let rec free_variables = function
       free_variables_list (free_variables fn) args
   | Lfunction{body; params} ->
       Ident.Set.diff (free_variables body)
-        (Ident.Set.of_list (List.map fst params))
+        (Ident.Set.of_list (List.map (fun p -> p.name) params))
   | Llet(_, _k, id, arg, body)
   | Lmutlet(_k, id, arg, body) ->
       Ident.Set.union
@@ -1073,6 +1084,12 @@ let subst update_env ?(freshen_bound_variables = false) s input_lam =
         ((id', rhs) :: ids' , l)
       ) ids ([], l)
   in
+  let bind_params params l =
+    List.fold_right (fun p (params', l) ->
+        let name', l = bind p.name l in
+        ({ p with name = name' } :: params' , l)
+      ) params ([], l)
+  in
   let rec subst s l lam =
     match lam with
     | Lvar id as lam ->
@@ -1097,7 +1114,7 @@ let subst update_env ?(freshen_bound_variables = false) s input_lam =
         Lapply{ap with ap_func = subst s l ap.ap_func;
                       ap_args = subst_list s l ap.ap_args}
     | Lfunction lf ->
-        let params, l' = bind_many lf.params l in
+        let params, l' = bind_params lf.params l in
         Lfunction {lf with params; body = subst s l' lf.body}
     | Llet(str, k, id, arg, body) ->
         let id, l' = bind id l in

lambda/lambda.mli

@@ -414,6 +414,14 @@ type function_attribute = {
   tmc_candidate: bool;
 }
 
+type parameter_attribute = No_attributes
+
+type lparam = {
+  name : Ident.t;
+  layout : layout;
+  attributes : parameter_attribute;
+  mode : alloc_mode
+}
 
 type scoped_location = Debuginfo.Scoped_location.t
 
@@ -451,7 +459,7 @@ type lambda =
 
 and lfunction = private
   { kind: function_kind;
-    params: (Ident.t * layout) list;
+    params: lparam list;
     return: layout;
     body: lambda;
     attr: function_attribute; (* specified with [@inline] attribute *)
@@ -568,7 +576,7 @@ val name_lambda_list: (lambda * layout) list -> (lambda list -> lambda) -> lambd
 
 val lfunction :
   kind:function_kind ->
-  params:(Ident.t * layout) list ->
+  params:lparam list ->
   return:layout ->
   body:lambda ->
   attr:function_attribute -> (* specified with [@inline] attribute *)
@@ -653,6 +661,7 @@ val swap_float_comparison : float_comparison -> float_comparison
 
 val default_function_attribute : function_attribute
 val default_stub_attribute : function_attribute
+val default_param_attribute : parameter_attribute
 
 val find_exact_application :
   function_kind -> arity:int -> lambda list -> lambda list option

lambda/printlambda.ml

@@ -667,7 +667,7 @@ let apply_specialised_attribute ppf = function
   | Always_specialise -> fprintf ppf " always_specialise"
   | Never_specialise -> fprintf ppf " never_specialise"
 
-let apply_probe ppf = function
+let apply_probe ppf : probe -> unit = function
   | None -> ()
   | Some {name} -> fprintf ppf " (probe %s)" name
 
@@ -702,16 +702,22 @@ let rec lam ppf = function
         match kind with
         | Curried {nlocal} ->
             fprintf ppf "@ {nlocal = %d}" nlocal;
-            List.iter (fun (param, k) ->
-                fprintf ppf "@ %a%a" Ident.print param layout k) params
+            List.iter (fun (p : Lambda.lparam) ->
+                (* Make sure we change this once there are attributes *)
+                let No_attributes = p.attributes in
+                fprintf ppf "@ %a%s%a"
+                  Ident.print p.name (alloc_kind p.mode) layout p.layout) params
         | Tupled ->
             fprintf ppf " (";
             let first = ref true in
             List.iter
-              (fun (param, k) ->
-                if !first then first := false else fprintf ppf ",@ ";
-                Ident.print ppf param;
-                layout ppf k)
+              (fun (p : Lambda.lparam) ->
+                 (* Make sure we change this once there are attributes *)
+                 let No_attributes = p.attributes in
+                 if !first then first := false else fprintf ppf ",@ ";
+                 Ident.print ppf p.name;
+                 Format.fprintf ppf "%s" (alloc_kind p.mode);
+                 layout ppf p.layout)
               params;
             fprintf ppf ")" in
       let rmode = if region then alloc_heap else alloc_local in

lambda/simplif.ml

@@ -373,8 +373,9 @@ let exact_application {kind; params; _} args =
   Lambda.find_exact_application kind ~arity args
 
 let beta_reduce params body args =
-  List.fold_left2 (fun l (param, kind) arg -> Llet(Strict, kind, param, arg, l))
-                  body params args
+  List.fold_left2
+    (fun l param arg -> Llet(Strict, param.layout, param.name, arg, l))
+    body params args
 
 (* Simplification of lets *)
 
@@ -778,7 +779,8 @@ let split_default_wrapper ~id:fun_id ~kind ~params ~return ~body
     | Llet(Strict, k, id,
            (Lifthenelse(Lprim (Pisint _, [Lvar optparam], _), _, _, _) as def),
            rest) when
-        Ident.name optparam = "*opt*" && List.mem_assoc optparam params
+        Ident.name optparam = "*opt*" &&
+        List.exists (fun p -> Ident.same p.name optparam) params
           && not (List.mem_assoc optparam map)
       ->
         let wrapper_body, inner = aux ((optparam, id) :: map) add_region rest in
@@ -793,14 +795,19 @@ let split_default_wrapper ~id:fun_id ~kind ~params ~return ~body
         List.iter (fun (id, _) -> if Ident.Set.mem id fv then raise Exit) map;
 
         let inner_id = Ident.create_local (Ident.name fun_id ^ "_inner") in
-        let map_param p layout =
+        let map_param (p : Lambda.lparam) =
           try
             (* If the param is optional, then it must be a value *)
-            List.assoc p map, Lambda.layout_field
+            {
+              name = List.assoc p.name map;
+              layout = Lambda.layout_field;
+              attributes = Lambda.default_param_attribute;
+              mode = p.mode
+            }
           with
-            Not_found -> p, layout
+            Not_found -> p
         in
-        let args = List.map (fun (p, layout) -> Lvar (fst (map_param p layout))) params in
+        let args = List.map (fun p -> Lvar (map_param p).name) params in
         let wrapper_body =
           Lapply {
             ap_func = Lvar inner_id;
@@ -815,11 +822,13 @@ let split_default_wrapper ~id:fun_id ~kind ~params ~return ~body
             ap_probe=None;
           }
         in
-        let inner_params = List.map (fun (param, layout) -> map_param param layout) params in
-        let new_ids = List.map (fun (param, layout) -> (Ident.rename param, layout)) inner_params in
+        let inner_params = List.map map_param params in
+        let new_ids =
+          List.map (fun p -> { p with name = Ident.rename p.name }) inner_params
+        in
         let subst =
-          List.fold_left2 (fun s (id, _) (new_id, _) ->
-            Ident.Map.add id new_id s
+          List.fold_left2 (fun s p new_p ->
+            Ident.Map.add p.name new_p.name s
           ) Ident.Map.empty inner_params new_ids
         in
         let body = Lambda.rename subst body in
@@ -1005,9 +1014,13 @@ let simplify_local_functions lam =
       | lam ->
           Lambda.shallow_map ~tail:rewrite ~non_tail:rewrite lam
     in
+    let new_params lf =
+      List.map
+        (fun p -> (p.name, p.layout)) lf.params
+    in
     List.fold_right
       (fun (st, lf) lam ->
-         Lstaticcatch (lam, (st, lf.params), rewrite lf.body, lf.return)
+         Lstaticcatch (lam, (st, new_params lf), rewrite lf.body, lf.return)
       )
       (LamTbl.find_all static lam0)
       lam

lambda/simplif.mli

@@ -32,7 +32,7 @@ val simplify_lambda: lambda -> lambda
 val split_default_wrapper
    : id:Ident.t
   -> kind:function_kind
-  -> params:(Ident.t * Lambda.layout) list
+  -> params:Lambda.lparam list
   -> return:Lambda.layout
   -> body:lambda
   -> attr:function_attribute

lambda/tmc.ml

@@ -62,7 +62,11 @@ and offset = Offset of lambda
 let offset_code (Offset t) = t
 
 let add_dst_params ({var; offset} : Ident.t destination) params =
-  (var, Lambda.layout_block) :: (offset, Lambda.layout_int) :: params
+  { name = var ; layout = Lambda.layout_block ;
+    attributes = Lambda.default_param_attribute ; mode = alloc_heap } ::
+  { name = offset ; layout = Lambda.layout_int ;
+    attributes = Lambda.default_param_attribute ; mode = alloc_heap } ::
+  params
 
 let add_dst_args ({var; offset} : offset destination) args =
   Lvar var :: offset_code offset :: args

lambda/transl_list_comprehension.ml

@@ -245,7 +245,17 @@ let rec translate_bindings
           ~kind:(Curried { nlocal = 2 })
           (* Only the accumulator is local, but since the function itself is
              local, [nlocal] has to be equal to the number of parameters *)
-          ~params:[element, element_kind; inner_acc, Pvalue Pgenval]
+          ~params:[
+            {name = element;
+             layout = element_kind;
+             attributes = Lambda.default_param_attribute;
+             (* CR ncourant: check *)
+             mode = alloc_heap};
+            {name = inner_acc;
+             layout = Pvalue Pgenval;
+             attributes = Lambda.default_param_attribute;
+             mode = alloc_local}
+          ]
           ~return:(Pvalue Pgenval)
           ~attr:default_function_attribute
           ~loc