Track modes in Lambda.lfunction and onwards (ocaml-flambda#33)

Extends Lambda.function_kind to track modes of curried functions.
The supported modes have a heap-returning prefix, followed by a
local-returning suffix, both of which may be empty. Functions which
do not fit this pattern cannot be expressed as a single Lfunction,
and must be explicitly curried in Lambda and Clambda. (See the new
Lambda.check_lfunction for the exact invariants)

Most of the changes are straightforward plumbing of the new info.
Some delicate changes are needed to preserve the new invariants, in:

  - simplif.ml: optimisations which combine Lfunctions
  - translcore.ml: transl_curried_function and friends
  - closure.ml: optimisations for partial and over-application
  - cmm_helpers.ml: mode-aware variants of caml_curry

Caveat: build_apply in translcore is wrong

Author: stedolan

asmcomp/cmm_helpers.ml

@@ -55,6 +55,7 @@ let black_block_header tag sz = Nativeint.logor (block_header tag sz) caml_black
 let local_block_header tag sz = Nativeint.logor (block_header tag sz) caml_local
 let white_closure_header sz = block_header Obj.closure_tag sz
 let black_closure_header sz = black_block_header Obj.closure_tag sz
+let local_closure_header sz = local_block_header Obj.closure_tag sz
 let infix_header ofs = block_header Obj.infix_tag ofs
 let float_header = block_header Obj.double_tag (size_float / size_addr)
 let floatarray_header len =
@@ -76,6 +77,11 @@ let pos_arity_in_closinfo = 8 * size_addr - 8
        (* arity = the top 8 bits of the closinfo word *)
 
 let closure_info ~arity ~startenv =
+  let arity =
+    match arity with
+    | Lambda.Tupled, n -> -n
+    | Lambda.Curried _, n -> n
+  in
   assert (-128 <= arity && arity <= 127);
   assert (0 <= startenv && startenv < 1 lsl (pos_arity_in_closinfo - 1));
   Nativeint.(add (shift_left (of_int arity) pos_arity_in_closinfo)
@@ -84,7 +90,10 @@ let closure_info ~arity ~startenv =
 
 let alloc_float_header dbg = Cconst_natint (float_header, dbg)
 let alloc_floatarray_header len dbg = Cconst_natint (floatarray_header len, dbg)
-let alloc_closure_header sz dbg = Cconst_natint (white_closure_header sz, dbg)
+let alloc_closure_header ~mode sz dbg =
+  match (mode : Lambda.alloc_mode) with
+  | Alloc_heap -> Cconst_natint (white_closure_header sz, dbg)
+  | Alloc_local -> Cconst_natint (local_closure_header sz, dbg)
 let alloc_infix_header ofs dbg = Cconst_natint (infix_header ofs, dbg)
 let alloc_closure_info ~arity ~startenv dbg =
   Cconst_natint (closure_info ~arity ~startenv, dbg)
@@ -838,12 +847,16 @@ let make_checkbound dbg = function
 (* Record application and currying functions *)
 
 let apply_function_sym n =
+  assert (n > 0);
   Compilenv.need_apply_fun n; "caml_apply" ^ Int.to_string n
-let curry_function_sym n =
-  Compilenv.need_curry_fun n;
-  if n >= 0
-  then "caml_curry" ^ Int.to_string n
-  else "caml_tuplify" ^ Int.to_string (-n)
+let curry_function_sym ar =
+  Compilenv.need_curry_fun ar;
+  match ar with
+  | Lambda.Curried {nlocal}, n ->
+     "caml_curry" ^ Int.to_string n ^
+       (if nlocal > 0 then "L" ^ Int.to_string nlocal else "")
+  | Lambda.Tupled, n ->
+     "caml_tuplify" ^ Int.to_string n
 
 (* Big arrays *)
 
@@ -1969,7 +1982,7 @@ let tuplify_function arity =
 *)
 
 let max_arity_optimized = 15
-let final_curry_function arity =
+let final_curry_function ~nlocal ~arity =
   let dbg = placeholder_dbg in
   let last_arg = V.create_local "arg" in
   let last_clos = V.create_local "clos" in
@@ -1998,7 +2011,9 @@ let final_curry_function arity =
                  newclos (n-1))
     end in
   let fun_name =
-    "caml_curry" ^ Int.to_string arity ^ "_" ^ Int.to_string (arity-1)
+    "caml_curry" ^ Int.to_string arity
+    ^ (if nlocal > 0 then "L" ^ Int.to_string nlocal else "")
+    ^ "_" ^ Int.to_string (arity-1)
   in
   let fun_dbg = placeholder_fun_dbg ~human_name:fun_name in
   Cfunction
@@ -2009,34 +2024,38 @@ let final_curry_function arity =
     fun_dbg;
    }
 
-let rec intermediate_curry_functions arity num =
+let rec intermediate_curry_functions ~nlocal ~arity num =
   let dbg = placeholder_dbg in
   if num = arity - 1 then
-    [final_curry_function arity]
+    [final_curry_function ~nlocal ~arity]
   else begin
-    let name1 = "caml_curry" ^ Int.to_string arity in
+    let name1 = "caml_curry" ^ Int.to_string arity
+                ^ (if nlocal > 0 then "L" ^ Int.to_string nlocal else "") in
     let name2 = if num = 0 then name1 else name1 ^ "_" ^ Int.to_string num in
     let arg = V.create_local "arg" and clos = V.create_local "clos" in
     let fun_dbg = placeholder_fun_dbg ~human_name:name2 in
+    let mode : Lambda.alloc_mode =
+      if num >= arity - nlocal then Alloc_local else Alloc_heap in
+    let curried n : Clambda.arity = (Curried {nlocal=min nlocal n}, n) in
     Cfunction
      {fun_name = name2;
       fun_args = [VP.create arg, typ_val; VP.create clos, typ_val];
       fun_body =
          if arity - num > 2 && arity <= max_arity_optimized then
-           Cop(Calloc Alloc_heap,
-               [alloc_closure_header 5 (dbg ());
+           Cop(Calloc mode,
+               [alloc_closure_header ~mode 5 (dbg ());
                 Cconst_symbol(name1 ^ "_" ^ Int.to_string (num+1), dbg ());
-                alloc_closure_info ~arity:(arity - num - 1)
+                alloc_closure_info ~arity:(curried (arity - num - 1))
                                    ~startenv:3 (dbg ());
                 Cconst_symbol(name1 ^ "_" ^ Int.to_string (num+1) ^ "_app",
                   dbg ());
                 Cvar arg; Cvar clos],
                dbg ())
          else
-           Cop(Calloc Alloc_heap,
-                [alloc_closure_header 4 (dbg ());
+           Cop(Calloc mode,
+                [alloc_closure_header ~mode 4 (dbg ());
                  Cconst_symbol(name1 ^ "_" ^ Int.to_string (num+1), dbg ());
-                 alloc_closure_info ~arity:1 ~startenv:2 (dbg ());
+                 alloc_closure_info ~arity:(curried 1) ~startenv:2 (dbg ());
                  Cvar arg; Cvar clos],
                 dbg ());
       fun_codegen_options = [];
@@ -2082,19 +2101,21 @@ let rec intermediate_curry_functions arity num =
                fun_dbg;
               }
           in
-          cf :: intermediate_curry_functions arity (num+1)
+          cf :: intermediate_curry_functions ~nlocal ~arity (num+1)
        else
-          intermediate_curry_functions arity (num+1))
+          intermediate_curry_functions ~nlocal ~arity (num+1))
   end
 
-let curry_function arity =
-  assert(arity <> 0);
-  (* Functions with arity = 0 does not have a curry_function *)
-  if arity > 0
-  then intermediate_curry_functions arity 0
-  else [tuplify_function (-arity)]
+let curry_function = function
+  | Lambda.Tupled, n ->
+     assert (n > 0); [tuplify_function n]
+  | Lambda.Curried {nlocal}, n ->
+     assert (n > 0);
+     intermediate_curry_functions ~nlocal ~arity:n 0
 
 module Int = Numbers.Int
+module AritySet =
+  Set.Make (struct type t = Clambda.arity let compare = compare end)
 
 let default_apply = Int.Set.add 2 (Int.Set.add 3 Int.Set.empty)
   (* These apply funs are always present in the main program because
@@ -2106,13 +2127,13 @@ let generic_functions shared units =
       (fun (apply,send,curry) (ui : Cmx_format.unit_infos) ->
          List.fold_right Int.Set.add ui.ui_apply_fun apply,
          List.fold_right Int.Set.add ui.ui_send_fun send,
-         List.fold_right Int.Set.add ui.ui_curry_fun curry)
-      (Int.Set.empty,Int.Set.empty,Int.Set.empty)
+         List.fold_right AritySet.add ui.ui_curry_fun curry)
+      (Int.Set.empty,Int.Set.empty,AritySet.empty)
       units in
   let apply = if shared then apply else Int.Set.union apply default_apply in
   let accu = Int.Set.fold (fun n accu -> apply_function n :: accu) apply [] in
   let accu = Int.Set.fold (fun n accu -> send_function n :: accu) send accu in
-  Int.Set.fold (fun n accu -> curry_function n @ accu) curry accu
+  AritySet.fold (fun arity accu -> curry_function arity @ accu) curry accu
 
 (* Primitives *)
 
@@ -2713,7 +2734,7 @@ let fundecls_size fundecls =
     (fun (f : Clambda.ufunction) ->
        let indirect_call_code_pointer_size =
          match f.arity with
-         | 0 | 1 -> 0
+         | Curried _, (0 | 1) -> 0
            (* arity 1 does not need an indirect call handler.
               arity 0 cannot be indirect called *)
          | _ -> 1
@@ -2746,30 +2767,32 @@ let emit_constant_closure ((_, global_symb) as symb) fundecls clos_vars cont =
       let rec emit_others pos = function
           [] -> clos_vars @ cont
       | (f2 : Clambda.ufunction) :: rem ->
-          if f2.arity = 1 || f2.arity = 0 then
+          match f2.arity with
+          | Curried _, (0|1) as arity ->
             Cint(infix_header pos) ::
             (closure_symbol f2) @
             Csymbol_address f2.label ::
-            Cint(closure_info ~arity:f2.arity ~startenv:(startenv - pos)) ::
+            Cint(closure_info ~arity ~startenv:(startenv - pos)) ::
             emit_others (pos + 3) rem
-          else
+          | arity ->
             Cint(infix_header pos) ::
             (closure_symbol f2) @
             Csymbol_address(curry_function_sym f2.arity) ::
-            Cint(closure_info ~arity:f2.arity ~startenv:(startenv - pos)) ::
+            Cint(closure_info ~arity ~startenv:(startenv - pos)) ::
             Csymbol_address f2.label ::
             emit_others (pos + 4) rem in
       Cint(black_closure_header (fundecls_size fundecls
                                  + List.length clos_vars)) ::
       cdefine_symbol symb @
       (closure_symbol f1) @
-      if f1.arity = 1 || f1.arity = 0 then
+      match f1.arity with
+      | Curried _, (0|1) as arity ->
         Csymbol_address f1.label ::
-        Cint(closure_info ~arity:f1.arity ~startenv) ::
+        Cint(closure_info ~arity ~startenv) ::
         emit_others 3 remainder
-      else
+      | arity ->
         Csymbol_address(curry_function_sym f1.arity) ::
-        Cint(closure_info ~arity:f1.arity ~startenv) ::
+        Cint(closure_info ~arity ~startenv) ::
         Csymbol_address f1.label ::
         emit_others 4 remainder
 

asmcomp/cmm_helpers.mli

@@ -65,15 +65,18 @@ val boxedint64_header : nativeint
 val boxedintnat_header : nativeint
 
 (** Closure info for a closure of given arity and distance to environment *)
-val closure_info : arity:int -> startenv:int -> nativeint
+val closure_info : arity:Clambda.arity -> startenv:int -> nativeint
 
 (** Wrappers *)
+(* FIXME: these all need mode params *)
 val alloc_float_header : Debuginfo.t -> expression
 val alloc_floatarray_header : int -> Debuginfo.t -> expression
-val alloc_closure_header : int -> Debuginfo.t -> expression
+val alloc_closure_header :
+  mode:Lambda.alloc_mode -> int -> Debuginfo.t -> expression
 val alloc_infix_header : int -> Debuginfo.t -> expression
 val alloc_closure_info :
-      arity:int -> startenv:int -> Debuginfo.t -> expression
+      arity:(Lambda.function_kind * int) -> startenv:int ->
+      Debuginfo.t -> expression
 val alloc_boxedint32_header : Debuginfo.t -> expression
 val alloc_boxedint64_header : Debuginfo.t -> expression
 val alloc_boxedintnat_header : Debuginfo.t -> expression
@@ -326,10 +329,9 @@ val check_bound :
     ensure its presence in the set of defined symbols *)
 val apply_function_sym : int -> string
 
-(** If [n] is positive, get the symbol for the generic currying wrapper with
-    [n] arguments, and ensure its presence in the set of defined symbols.
-    Otherwise, do the same for the generic tuple wrapper with [-n] arguments. *)
-val curry_function_sym : int -> 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
 
 (** Bigarrays *)
 

asmcomp/cmmgen.ml

@@ -394,14 +394,15 @@ let rec transl env e =
             Cmmgen_state.add_function f;
             let dbg = f.dbg in
             let without_header =
-              if f.arity = 1 || f.arity = 0 then
+              match f.arity with
+              | Curried _, (1|0) as arity ->
                 Cconst_symbol (f.label, dbg) ::
-                alloc_closure_info ~arity:f.arity
+                alloc_closure_info ~arity
                                    ~startenv:(startenv - pos) dbg ::
                 transl_fundecls (pos + 3) rem
-              else
+              | arity ->
                 Cconst_symbol (curry_function_sym f.arity, dbg) ::
-                alloc_closure_info ~arity:f.arity
+                alloc_closure_info ~arity
                                    ~startenv:(startenv - pos) dbg ::
                 Cconst_symbol (f.label, dbg) ::
                 transl_fundecls (pos + 4) rem

file_formats/cmx_format.mli

@@ -41,7 +41,7 @@ type unit_infos =
     mutable ui_defines: string list;      (* Unit and sub-units implemented *)
     mutable ui_imports_cmi: crcs;         (* Interfaces imported *)
     mutable ui_imports_cmx: crcs;         (* Infos imported *)
-    mutable ui_curry_fun: int list;       (* Currying functions needed *)
+    mutable ui_curry_fun: Clambda.arity list; (* Currying functions needed *)
     mutable ui_apply_fun: int list;       (* Apply functions needed *)
     mutable ui_send_fun: int list;        (* Send functions needed *)
     mutable ui_export_info: export_info;

lambda/lambda.ml

@@ -261,7 +261,7 @@ type local_attribute =
   | Never_local (* [@local never] *)
   | Default_local (* [@local maybe] or no [@local] attribute *)
 
-type function_kind = Curried | Tupled
+type function_kind = Curried of {nlocal: int} | Tupled
 
 type let_kind = Strict | Alias | StrictOpt | Variable
 
@@ -317,7 +317,8 @@ and lfunction =
     body: lambda;
     attr: function_attribute; (* specified with [@inline] attribute *)
     loc: scoped_location;
-    mode: alloc_mode }
+    mode: alloc_mode;
+    ret_mode: alloc_mode; }
 
 and lambda_apply =
   { ap_func : lambda;
@@ -359,6 +360,29 @@ let const_unit = const_int 0
 
 let lambda_unit = Lconst const_unit
 
+let check_lfunction fn =
+  (* A curried function type with n parameters has n arrows. Of these,
+     the first [n-nlocal] have return mode Heap, while the remainder
+     have return mode Local, except possibly the final one.
+
+     That is, after supplying the first [n-nlocal] arguments, further
+     partial applications must be locally allocated.
+
+     A curried function with no local parameters or returns has kind
+     [Curried {nlocal=0}]. *)
+  let nparams = List.length fn.params in
+  begin match fn.mode, fn.ret_mode, fn.kind with
+  | Alloc_heap, _, Tupled -> ()
+  | Alloc_local, _, Tupled ->
+     (* Tupled optimisation does not apply to local functions *)
+     assert false
+  | mode, ret_mode, Curried {nlocal} ->
+     assert (0 <= nlocal);
+     assert (nlocal <= nparams);
+     if ret_mode = Alloc_local then assert (nlocal >= 1);
+     if mode = Alloc_local then assert (nlocal = nparams)
+  end
+
 let default_function_attribute = {
   inline = Default_inline;
   specialise = Default_specialise;
@@ -842,8 +866,9 @@ let shallow_map f = function
         ap_inlined;
         ap_specialised;
       }
-  | Lfunction { kind; params; return; body; attr; loc; mode } ->
-      Lfunction { kind; params; return; body = f body; attr; loc; mode }
+  | Lfunction { kind; params; return; body; attr; loc; mode; ret_mode } ->
+      Lfunction { kind; params; return; body = f body; attr; loc;
+                  mode; ret_mode }
   | Llet (str, k, v, e1, e2) ->
       Llet (str, k, v, f e1, f e2)
   | Lletrec (idel, e2) ->
@@ -957,11 +982,6 @@ let merge_inline_attributes attr1 attr2 =
     if attr1 = attr2 then Some attr1
     else None
 
-let function_is_curried func =
-  match func.kind with
-  | Curried -> true
-  | Tupled -> false
-
 let max_arity () =
   if !Clflags.native_code then 126 else max_int
   (* 126 = 127 (the maximal number of parameters supported in C--)

lambda/lambda.mli

@@ -236,7 +236,10 @@ type local_attribute =
   | Never_local (* [@local never] *)
   | Default_local (* [@local maybe] or no [@local] attribute *)
 
-type function_kind = Curried | Tupled
+type function_kind = Curried of {nlocal: int} | Tupled
+(* [nlocal] determines how many arguments may be partially applied
+   before the resulting closure must be locally allocated.
+   See [check_lfunction] for details *)
 
 type let_kind = Strict | Alias | StrictOpt | Variable
 (* Meaning of kinds for let x = e in e':
@@ -301,7 +304,8 @@ and lfunction =
     body: lambda;
     attr: function_attribute; (* specified with [@inline] attribute *)
     loc : scoped_location;
-    mode : alloc_mode }
+    mode : alloc_mode;
+    ret_mode : alloc_mode }
 
 and lambda_apply =
   { ap_func : lambda;
@@ -354,6 +358,7 @@ val make_key: lambda -> lambda option
 val const_unit: structured_constant
 val const_int : int -> structured_constant
 val lambda_unit: lambda
+val check_lfunction : lfunction -> unit
 val name_lambda: let_kind -> lambda -> (Ident.t -> lambda) -> lambda
 val name_lambda_list: lambda list -> (lambda list -> lambda) -> lambda
 
@@ -431,8 +436,6 @@ val swap_float_comparison : float_comparison -> float_comparison
 val default_function_attribute : function_attribute
 val default_stub_attribute : function_attribute
 
-val function_is_curried : lfunction -> bool
-
 val max_arity : unit -> int
   (** Maximal number of parameters for a function, or in other words,
       maximal length of the [params] list of a [lfunction] record.