Generic functions for unboxed types (oxcaml#1104)

Author: Ekdohibs

backend/cmm_helpers.ml

@@ -369,6 +369,8 @@ val call_cached_method :
   expression ->
   expression ->
   expression list ->
+  machtype list ->
+  machtype ->
   Clambda.apply_kind ->
   Debuginfo.t ->
   expression
@@ -408,11 +410,13 @@ val opaque : expression -> Debuginfo.t -> expression
 
 (** Get the symbol for the generic application with [n] arguments, and ensure
     its presence in the set of defined symbols *)
-val apply_function_sym : int -> Lambda.alloc_mode -> string
+val apply_function_sym :
+  machtype list -> machtype -> Lambda.alloc_mode -> string
 
 (** Get the symbol for the generic currying or tuplifying wrapper with [n]
     arguments, and ensure its presence in the set of defined symbols. *)
-val curry_function_sym : Clambda.arity -> string
+val curry_function_sym :
+  Lambda.function_kind -> machtype list -> machtype -> string
 
 (** Bigarrays *)
 
@@ -755,6 +759,8 @@ val generic_apply :
   Asttypes.mutable_flag ->
   expression ->
   expression list ->
+  machtype list ->
+  machtype ->
   Clambda.apply_kind ->
   Debuginfo.t ->
   expression
@@ -774,6 +780,8 @@ val send :
   expression ->
   expression ->
   expression list ->
+  machtype list ->
+  machtype ->
   Clambda.apply_kind ->
   Debuginfo.t ->
   expression
@@ -1107,6 +1115,7 @@ val indirect_call :
   Lambda.region_close ->
   Lambda.alloc_mode ->
   expression ->
+  machtype list ->
   expression list ->
   expression
 
@@ -1118,6 +1127,7 @@ val indirect_full_call :
   Lambda.region_close ->
   Lambda.alloc_mode ->
   expression ->
+  machtype list ->
   expression list ->
   expression
 
@@ -1197,3 +1207,5 @@ val transl_attrib : Lambda.check_attribute -> Cmm.codegen_option list
 val make_symbol : ?compilation_unit:Compilation_unit.t -> string -> string
 
 val kind_of_layout : Lambda.layout -> value_kind
+
+val machtype_of_layout : Lambda.layout -> machtype

backend/cmm_helpers.mli

@@ -465,7 +465,7 @@ let rec transl env e =
                                    ~startenv:(startenv - pos) ~is_last dbg ::
                 transl_fundecls (pos + 3) rem
               | arity ->
-                Cconst_symbol (curry_function_sym arity, dbg) ::
+                Cconst_symbol (curry_function_sym (fst arity) (List.init (snd arity) (fun _ -> typ_val)) typ_val, dbg) ::
                 alloc_closure_info ~arity
                                    ~startenv:(startenv - pos) ~is_last dbg ::
                 Cconst_symbol (f.label, dbg) ::
@@ -496,12 +496,16 @@ let rec transl env e =
   | Ugeneric_apply(clos, args, kind, dbg) ->
       let clos = transl env clos in
       let args = List.map (transl env) args in
-      generic_apply (mut_from_env env clos) clos args kind dbg
+      let args_type = List.map (fun _ -> typ_val) args in
+      let return = typ_val in
+      generic_apply (mut_from_env env clos) clos args args_type return kind dbg
   | Usend(kind, met, obj, args, pos, dbg) ->
       let met = transl env met in
       let obj = transl env obj in
       let args = List.map (transl env) args in
-      send kind met obj args pos dbg
+      let args_type = List.map (fun _ -> typ_val) args in
+      let return = typ_val in
+      send kind met obj args args_type return pos dbg
   | Ulet(str, kind, id, exp, body) ->
       transl_let env str kind id exp (fun env -> transl env body)
   | Uphantom_let (var, defining_expr, body) ->

backend/cmmgen.ml

@@ -41,11 +41,11 @@ type export_info_raw =
   | Flambda1_raw of Export_info.t
   | Flambda2_raw of Flambda2_cmx.Flambda_cmx_format.raw option
 
-type apply_fn := int * Lambda.alloc_mode
+type apply_fn := Cmm.machtype list * Cmm.machtype * Lambda.alloc_mode
 
 (* Curry/apply/send functions *)
 type generic_fns =
-  { curry_fun: Clambda.arity list;
+  { curry_fun: (Lambda.function_kind * Cmm.machtype list * Cmm.machtype) list;
     apply_fun: apply_fn list;
     send_fun: apply_fn list }
 

file_formats/cmx_format.mli

@@ -323,24 +323,24 @@ let approx_env () = !merged_environment
 
 (* Record that a currying function or application function is needed *)
 
-let need_curry_fun arity =
+let need_curry_fun kind arity result =
   let fns = current_unit.ui_generic_fns in
-  if not (List.mem arity fns.curry_fun) then
+  if not (List.mem (kind, arity, result) fns.curry_fun) then
     current_unit.ui_generic_fns <-
-      { fns with curry_fun = arity :: fns.curry_fun }
+      { fns with curry_fun = (kind, arity, result) :: fns.curry_fun }
 
-let need_apply_fun n mode =
-  assert(n > 0);
+let need_apply_fun arity result mode =
+  assert(List.compare_length_with arity 0 > 0);
   let fns = current_unit.ui_generic_fns in
-  if not (List.mem (n,mode) fns.apply_fun) then
+  if not (List.mem (arity, result, mode) fns.apply_fun) then
     current_unit.ui_generic_fns <-
-      { fns with apply_fun = (n,mode) :: fns.apply_fun }
+      { fns with apply_fun = (arity, result, mode) :: fns.apply_fun }
 
-let need_send_fun n mode =
+let need_send_fun arity result mode =
   let fns = current_unit.ui_generic_fns in
-  if not (List.mem (n,mode) fns.send_fun) then
+  if not (List.mem (arity, result, mode) fns.send_fun) then
     current_unit.ui_generic_fns <-
-      { fns with send_fun = (n,mode) :: fns.send_fun }
+      { fns with send_fun = (arity, result, mode) :: fns.send_fun }
 
 (* Write the description of the current unit *)
 

middle_end/compilenv.ml

@@ -68,9 +68,9 @@ val get_unit_export_info
 val flambda2_set_export_info : Flambda2_cmx.Flambda_cmx_format.t -> unit
         (* Set the export information for the current unit (Flambda 2 only). *)
 
-val need_curry_fun: Clambda.arity -> unit
-val need_apply_fun: int -> Lambda.alloc_mode -> unit
-val need_send_fun: int -> Lambda.alloc_mode -> unit
+val need_curry_fun: Lambda.function_kind -> Cmm.machtype list -> Cmm.machtype -> unit
+val need_apply_fun: Cmm.machtype list -> Cmm.machtype -> Lambda.alloc_mode -> unit
+val need_send_fun: Cmm.machtype list -> Cmm.machtype -> Lambda.alloc_mode -> unit
         (* Record the need of a currying (resp. application,
            message sending) function with the given arity *)
 

middle_end/compilenv.mli

@@ -108,9 +108,10 @@ let translate_apply0 env res apply =
         Ece.all ))
   | Function { function_call = Indirect_unknown_arity; alloc_mode } ->
     fail_if_probe apply;
-    ( C.indirect_call ~dbg Cmm.typ_val pos
+    let args_ty, ty = Cmm.(List.map (fun _ -> [| Val |]) args, [| Val |]) in
+    ( C.indirect_call ~dbg ty pos
         (Alloc_mode.For_types.to_lambda alloc_mode)
-        callee args,
+        callee args_ty args,
       env,
       res,
       Ece.all )
@@ -129,9 +130,14 @@ let translate_apply0 env res apply =
         return_arity |> Flambda_arity.With_subkinds.to_arity
         |> C.machtype_of_return_arity
       in
+      let args_ty =
+        List.map
+          (fun k -> C.machtype_of_kind (Flambda_kind.With_subkind.kind k))
+          (Flambda_arity.With_subkinds.to_list param_arity)
+      in
       ( C.indirect_full_call ~dbg ty pos
           (Alloc_mode.For_types.to_lambda alloc_mode)
-          callee args,
+          callee args_ty args,
         env,
         res,
         Ece.all )
@@ -175,10 +181,14 @@ let translate_apply0 env res apply =
       Ece.all )
   | Call_kind.Method { kind; obj; alloc_mode } ->
     fail_if_probe apply;
+    let args_ty, ty = Cmm.(List.map (fun _ -> [| Val |]) args, [| Val |]) in
     let obj, env, res, _ = C.simple ~dbg env res obj in
     let kind = Call_kind.Method_kind.to_lambda kind in
     let alloc_mode = Alloc_mode.For_types.to_lambda alloc_mode in
-    C.send kind callee obj args (pos, alloc_mode) dbg, env, res, Ece.all
+    ( C.send kind callee obj args args_ty ty (pos, alloc_mode) dbg,
+      env,
+      res,
+      Ece.all )
 
 (* Function calls that have an exn continuation with extra arguments must be
    wrapped with assignments for the mutable variables used to pass the extra

middle_end/flambda2/to_cmm/to_cmm_expr.ml

@@ -34,8 +34,14 @@ type closure_code_pointers =
 
 let get_func_decl_params_arity t code_id =
   let info = Env.get_code_metadata t code_id in
-  let num_params =
-    Flambda_arity.With_subkinds.cardinal (Code_metadata.params_arity info)
+  let params_ty =
+    List.map
+      (fun k -> C.machtype_of_kind (Flambda_kind.With_subkind.kind k))
+      (Flambda_arity.With_subkinds.to_list (Code_metadata.params_arity info))
+  in
+  let result_ty =
+    C.machtype_of_return_arity
+      (Flambda_arity.With_subkinds.to_arity (Code_metadata.result_arity info))
   in
   let kind : Lambda.function_kind =
     if Code_metadata.is_tupled info
@@ -44,11 +50,11 @@ let get_func_decl_params_arity t code_id =
       Lambda.Curried { nlocal = Code_metadata.num_trailing_local_params info }
   in
   let closure_code_pointers =
-    match kind, num_params with
-    | Curried _, (0 | 1) -> Full_application_only
+    match kind, params_ty with
+    | Curried _, ([] | [_]) -> Full_application_only
     | (Curried _ | Tupled), _ -> Full_and_partial_application
   in
-  let arity = kind, num_params in
+  let arity = kind, params_ty, result_ty in
   arity, closure_code_pointers, Code_metadata.dbg info
 
 type for_static_sets =
@@ -147,11 +153,11 @@ end = struct
     | Function_slot { size; function_slot; last_function_slot } -> (
       let code_id = Function_slot.Map.find function_slot decls in
       let code_linkage_name = Code_id.linkage_name code_id in
-      let arity, closure_code_pointers, dbg =
+      let (kind, params_ty, result_ty), closure_code_pointers, dbg =
         get_func_decl_params_arity env code_id
       in
       let closure_info =
-        C.closure_info ~arity ~startenv:(startenv - slot_offset)
+        C.closure_info ~arity:(kind, List.length params_ty) ~startenv:(startenv - slot_offset)
           ~is_last:last_function_slot
       in
       let acc =
@@ -194,7 +200,8 @@ end = struct
           P.symbol_from_linkage_name ~dbg code_linkage_name
           :: P.int ~dbg closure_info
           :: P.symbol_from_linkage_name ~dbg
-               (Linkage_name.of_string (C.curry_function_sym arity))
+               (Linkage_name.of_string
+                  (C.curry_function_sym kind params_ty result_ty))
           :: acc
         in
         acc, slot_offset + size, env, res, Ece.pure, updates)

middle_end/flambda2/to_cmm/to_cmm_set_of_closures.ml

@@ -154,18 +154,34 @@ let print_global_table table =
 open Cmx_format
 open Cmxs_format
 
+let machtype_identifier t =
+  let open Cmm in
+  let char_of_component = function
+    | Val -> 'V' | Int -> 'I' | Float -> 'F' | Addr -> 'A'
+  in
+  String.of_seq (Seq.map char_of_component (Array.to_seq t))
+
+let unique_arity_identifier arity =
+  let open Cmm in
+  if List.for_all (function [|Val|] -> true | _ -> false) arity then
+    Int.to_string (List.length arity)
+  else
+    String.concat "_" (List.map machtype_identifier arity)
+
+let return_arity_identifier t =
+  let open Cmm in
+  match t with [|Val|] -> "" | _ -> "_R" ^ machtype_identifier t
+
 let print_generic_fns gfns =
   let pr_afuns _ fns =
     let mode = function Lambda.Alloc_heap -> "" | Lambda.Alloc_local -> "L" in
-    List.iter (fun (arity, m) -> printf " %d%s" arity (mode m)) fns
-  in
+    List.iter (fun (arity,result,m) -> printf " %s%s%s" (unique_arity_identifier arity) (return_arity_identifier result) (mode m)) fns in
   let pr_cfuns _ fns =
-    List.iter
-      (function
-        | Lambda.Curried { nlocal }, a -> printf " %dL%d" a nlocal
-        | Lambda.Tupled, a -> printf " -%d" a)
-      fns
-  in
+    List.iter (function
+      | (Lambda.Curried {nlocal}, arity, result) ->
+        printf " %s%sL%d" (unique_arity_identifier arity) (return_arity_identifier result) nlocal
+      | (Lambda.Tupled, arity, result) ->
+        printf " -%s%s" (unique_arity_identifier arity) (return_arity_identifier result)) fns in
   printf "Currying functions:%a\n" pr_cfuns gfns.curry_fun;
   printf "Apply functions:%a\n" pr_afuns gfns.apply_fun;
   printf "Send functions:%a\n" pr_afuns gfns.send_fun