Changes to arity in clambda (ocaml-flambda#1106)

backend/cmm_helpers.ml

@@ -114,10 +114,7 @@ let caml_int64_ops = "caml_int64_ops"
 let pos_arity_in_closinfo = (8 * size_addr) - 8
 (* arity = the top 8 bits of the closinfo word *)
 
-let closure_info ~arity ~startenv ~is_last =
-  let arity =
-    match arity with Lambda.Tupled, n -> -n | Lambda.Curried _, n -> n
-  in
+let pack_closure_info ~arity ~startenv ~is_last =
   assert (-128 <= arity && arity <= 127);
   assert (0 <= startenv && startenv < 1 lsl (pos_arity_in_closinfo - 2));
   Nativeint.(
@@ -129,6 +126,19 @@ let closure_info ~arity ~startenv ~is_last =
             (pos_arity_in_closinfo - 1))
          (add (shift_left (of_int startenv) 1) 1n)))
 
+let closure_info' ~arity ~startenv ~is_last =
+  let arity =
+    match arity with
+    | Lambda.Tupled, l -> -List.length l
+    | Lambda.Curried _, l -> List.length l
+  in
+  pack_closure_info ~arity ~startenv ~is_last
+
+let closure_info ~(arity : Clambda.arity) ~startenv ~is_last =
+  closure_info'
+    ~arity:(arity.function_kind, arity.params_layout)
+    ~startenv ~is_last
+
 let alloc_float_header mode dbg =
   match mode with
   | Lambda.Alloc_heap -> Cconst_natint (float_header, dbg)
@@ -2684,7 +2694,6 @@ let intermediate_curry_functions ~nlocal ~arity result =
       let mode : Lambda.alloc_mode =
         if num >= narity - nlocal then Lambda.alloc_local else Lambda.alloc_heap
       in
-      let curried n = Lambda.Curried { nlocal = min nlocal n }, n in
       let has_nary = curry_clos_has_nary_application ~narity (num + 1) in
       let function_slot_size = if has_nary then 3 else 2 in
       Cfunction
@@ -2697,11 +2706,14 @@ let intermediate_curry_functions ~nlocal ~arity result =
                     (function_slot_size + machtype_stored_size arg_type + 1)
                     (dbg ());
                   Cconst_symbol (name1 ^ "_" ^ Int.to_string (num + 1), dbg ());
-                  alloc_closure_info
-                    ~arity:(curried (if has_nary then narity - num - 1 else 1))
-                    ~startenv:
-                      (function_slot_size + machtype_non_scanned_size arg_type)
-                    (dbg ()) ~is_last:true ]
+                  Cconst_natint
+                    ( pack_closure_info
+                        ~arity:(if has_nary then narity - num - 1 else 1)
+                        ~startenv:
+                          (function_slot_size
+                          + machtype_non_scanned_size arg_type)
+                        ~is_last:true,
+                      dbg () ) ]
                 @ (if has_nary
                   then
                     [ Cconst_symbol
@@ -3520,7 +3532,7 @@ let fundecls_size fundecls =
     (fun (f : Clambda.ufunction) ->
       let indirect_call_code_pointer_size =
         match f.arity with
-        | Curried _, (0 | 1) ->
+        | { function_kind = Curried _; params_layout = [] | [_]; _ } ->
           0
           (* arity 1 does not need an indirect call handler. arity 0 cannot be
              indirect called *)
@@ -3557,17 +3569,17 @@ let emit_constant_closure ((_, global_symb) as symb) fundecls clos_vars cont =
       | (f2 : Clambda.ufunction) :: rem -> (
         let is_last = match rem with [] -> true | _ :: _ -> false in
         match f2.arity with
-        | (Curried _, (0 | 1)) as arity ->
+        | { function_kind = Curried _; params_layout = [] | [_]; _ } as arity ->
           (Cint (infix_header pos) :: closure_symbol f2)
           @ Csymbol_address f2.label
             :: Cint (closure_info ~arity ~startenv:(startenv - pos) ~is_last)
             :: emit_others (pos + 3) rem
         | arity ->
           (Cint (infix_header pos) :: closure_symbol f2)
           @ Csymbol_address
-              (curry_function_sym (fst arity)
-                 (List.init (snd arity) (fun _ -> typ_val))
-                 typ_val)
+              (curry_function_sym arity.function_kind
+                 (List.map machtype_of_layout arity.params_layout)
+                 (machtype_of_layout arity.return_layout))
             :: Cint (closure_info ~arity ~startenv:(startenv - pos) ~is_last)
             :: Csymbol_address f2.label
             :: emit_others (pos + 4) rem)
@@ -3578,15 +3590,15 @@ let emit_constant_closure ((_, global_symb) as symb) fundecls clos_vars cont =
     @ closure_symbol f1
     @
     match f1.arity with
-    | (Curried _, (0 | 1)) as arity ->
+    | { function_kind = Curried _; params_layout = [] | [_]; _ } as arity ->
       Csymbol_address f1.label
       :: Cint (closure_info ~arity ~startenv ~is_last)
       :: emit_others 3 remainder
     | arity ->
       Csymbol_address
-        (curry_function_sym (fst arity)
-           (List.init (snd arity) (fun _ -> typ_val))
-           typ_val)
+        (curry_function_sym arity.function_kind
+           (List.map machtype_of_layout arity.params_layout)
+           (machtype_of_layout arity.return_layout))
       :: Cint (closure_info ~arity ~startenv ~is_last)
       :: Csymbol_address f1.label :: emit_others 4 remainder)
 

backend/cmm_helpers.mli

@@ -63,11 +63,17 @@ val boxedintnat_header : nativeint
 val closure_info :
   arity:Clambda.arity -> startenv:int -> is_last:bool -> nativeint
 
+val closure_info' :
+  arity:Lambda.function_kind * 'a list ->
+  startenv:int ->
+  is_last:bool ->
+  nativeint
+
 (** Wrappers *)
 val alloc_infix_header : int -> Debuginfo.t -> expression
 
 val alloc_closure_info :
-  arity:Lambda.function_kind * int ->
+  arity:Clambda.arity ->
   startenv:int ->
   is_last:bool ->
   Debuginfo.t ->

backend/cmmgen.ml

@@ -459,15 +459,18 @@ let rec transl env e =
             let dbg = f.dbg in
             let without_header =
               match f.arity with
-              | Curried _, (1|0) as arity ->
+              | { function_kind = Curried _ ; params_layout = ([] | [_]) } as arity ->
                 Cconst_symbol (f.label, dbg) ::
                 alloc_closure_info ~arity
                                    ~startenv:(startenv - pos) ~is_last dbg ::
                 transl_fundecls (pos + 3) rem
               | arity ->
-                Cconst_symbol (
-                  curry_function_sym (fst arity)
-                    (List.init (snd arity) (fun _ -> typ_val)) typ_val, dbg) ::
+                Cconst_symbol
+                  (curry_function_sym
+                     arity.function_kind
+                     (List.map machtype_of_layout arity.params_layout)
+                     (machtype_of_layout arity.return_layout),
+                   dbg) ::
                 alloc_closure_info ~arity
                                    ~startenv:(startenv - pos) ~is_last dbg ::
                 Cconst_symbol (f.label, dbg) ::
@@ -488,26 +491,25 @@ let rec transl env e =
       let ptr = transl env arg in
       let dbg = Debuginfo.none in
       ptr_offset ptr offset dbg
-  | Udirect_apply(handler_code_sym, args, Some { name; }, _, dbg) ->
+  | Udirect_apply(handler_code_sym, args, Some { name; }, _, _, dbg) ->
       let args = List.map (transl env) args in
       return_unit dbg
         (Cop(Cprobe { name; handler_code_sym; }, args, dbg))
-  | Udirect_apply(lbl, args, None, kind, dbg) ->
+  | Udirect_apply(lbl, args, None, result_layout, kind, dbg) ->
       let args = List.map (transl env) args in
-      let return = typ_val in
-      direct_apply lbl return args kind dbg
-  | Ugeneric_apply(clos, args, kind, dbg) ->
+      direct_apply lbl (machtype_of_layout result_layout) args kind dbg
+  | Ugeneric_apply(clos, args, args_layout, result_layout, kind, dbg) ->
       let clos = transl env clos in
       let args = List.map (transl env) args in
-      let args_type = List.map (fun _ -> typ_val) args in
-      let return = typ_val in
+      let args_type = List.map machtype_of_layout args_layout in
+      let return = machtype_of_layout result_layout in
       generic_apply (mut_from_env env clos) clos args args_type return kind dbg
-  | Usend(kind, met, obj, args, pos, dbg) ->
+  | Usend(kind, met, obj, args, args_layout, result_layout, pos, dbg) ->
       let met = transl env met in
       let obj = transl env obj in
       let args = List.map (transl env) args in
-      let args_type = List.map (fun _ -> typ_val) args in
-      let return = typ_val in
+      let args_type = List.map machtype_of_layout args_layout in
+      let return = machtype_of_layout result_layout 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)
@@ -1487,8 +1489,15 @@ let transl_function f =
     else
       [ Reduce_code_size ]
   in
+  let params_layout =
+    if List.length f.params = List.length f.arity.params_layout then
+      f.arity.params_layout
+    else
+      f.arity.params_layout @ [Lambda.layout_function]
+  in
   Cfunction {fun_name = f.label;
-             fun_args = List.map (fun (id, _) -> (id, typ_val)) f.params;
+             fun_args = List.map2 (fun id ty -> (id, machtype_of_layout ty))
+                 f.params params_layout;
              fun_body = cmm_body;
              fun_codegen_options;
              fun_poll = f.poll;

middle_end/clambda.ml

@@ -20,7 +20,11 @@ open Asttypes
 open Lambda
 
 type function_label = string
-type arity = Lambda.function_kind * int
+type arity = {
+  function_kind : Lambda.function_kind ;
+  params_layout : Lambda.layout list ;
+  return_layout : Lambda.layout ;
+}
 type apply_kind = Lambda.region_close * Lambda.alloc_mode
 
 type ustructured_constant =
@@ -49,9 +53,9 @@ and ulambda =
     Uvar of Backend_var.t
   | Uconst of uconstant
   | Udirect_apply of
-      function_label * ulambda list * Lambda.probe * apply_kind * Debuginfo.t
+      function_label * ulambda list * Lambda.probe * Lambda.layout * apply_kind * Debuginfo.t
   | Ugeneric_apply of
-      ulambda * ulambda list * apply_kind * Debuginfo.t
+      ulambda * ulambda list * Lambda.layout list * Lambda.layout * apply_kind * Debuginfo.t
   | Uclosure of {
       functions : ufunction list ;
       not_scanned_slots : ulambda list ;
@@ -90,16 +94,15 @@ and ulambda =
   | Uassign of Backend_var.t * ulambda
   | Usend of
       meth_kind * ulambda * ulambda * ulambda list
-      * apply_kind * Debuginfo.t
+      * Lambda.layout list * Lambda.layout * apply_kind * Debuginfo.t
   | Uunreachable
   | Uregion of ulambda
   | Utail of ulambda
 
 and ufunction = {
   label  : function_label;
   arity  : arity;
-  params : (Backend_var.With_provenance.t * layout) list;
-  return : layout;
+  params : Backend_var.With_provenance.t list;
   body   : ulambda;
   dbg    : Debuginfo.t;
   env    : Backend_var.t option;

middle_end/clambda.mli

@@ -20,7 +20,11 @@ open Asttypes
 open Lambda
 
 type function_label = string
-type arity = Lambda.function_kind * int
+type arity = {
+  function_kind : Lambda.function_kind ;
+  params_layout : Lambda.layout list ;
+  return_layout : Lambda.layout ;
+}
 type apply_kind = Lambda.region_close * Lambda.alloc_mode
 
 type ustructured_constant =
@@ -60,9 +64,9 @@ and ulambda =
     Uvar of Backend_var.t
   | Uconst of uconstant
   | Udirect_apply of
-      function_label * ulambda list * Lambda.probe * apply_kind * Debuginfo.t
+      function_label * ulambda list * Lambda.probe * Lambda.layout * apply_kind * Debuginfo.t
   | Ugeneric_apply of
-      ulambda * ulambda list * apply_kind * Debuginfo.t
+      ulambda * ulambda list * Lambda.layout list * Lambda.layout * apply_kind * Debuginfo.t
   | Uclosure of {
       functions : ufunction list ;
       not_scanned_slots : ulambda list ;
@@ -101,16 +105,15 @@ and ulambda =
   | Uassign of Backend_var.t * ulambda
   | Usend of
       meth_kind * ulambda * ulambda * ulambda list
-      * apply_kind * Debuginfo.t
+      * Lambda.layout list * Lambda.layout * apply_kind * Debuginfo.t
   | Uunreachable
   | Uregion of ulambda
   | Utail of ulambda
 
 and ufunction = {
   label  : function_label;
   arity  : arity;
-  params : (Backend_var.With_provenance.t * layout) list;
-  return : layout;
+  params : Backend_var.With_provenance.t list;
   body   : ulambda;
   dbg    : Debuginfo.t;
   env    : Backend_var.t option;