Create fewer regions, according to typechecking (#59)

Instead of creating one at each expression in local mode, regions are now only created where typechecking indicates they should go:

  - Function bodies (except for local-returning functions)
  - Loop bodies
  - Toplevel expressions (module-level let, certain cases of class initialisation)
  - Inlined overapplications, where the inlined function returns local but the whole application does not

The last one of these is tricky, and fixes a bug: a call `f a b c` may be in heap mode, but actually correspond to an overapplication `(f a) b c`, where the one-argument function `f` returns a local closure. Previously, this local closure would leak, and the newly inserted region fixes this.

This fix requires tracking the mode of an application all the way through to Cmm, since the `caml_applyN` stubs are split into local and non-local versions according to whether this region should be present in the slow-path.

A couple of optimisations are updated to correctly handle regions:

  - All default arguments are now global so that the worker-wrapper transformation is always valid.
  - The local functions transformation in simplif now understands tailcalls within regions.

Finally, there is an optimisation in Lambda to prevent unnecessary regions being created.

Co-authored-by: Stephen Dolan <[email protected]>

Author: lpw25

asmcomp/cmm.ml

@@ -272,34 +272,34 @@ let iter_shallow_tail f = function
   | Cop _ ->
       false
 
-let rec map_tail f = function
+let map_shallow_tail f = function
   | Clet(id, exp, body) ->
-      Clet(id, exp, map_tail f body)
+      Clet(id, exp, f body)
   | Clet_mut(id, kind, exp, body) ->
-      Clet_mut(id, kind, exp, map_tail f body)
+      Clet_mut(id, kind, exp, f body)
   | Cphantom_let(id, exp, body) ->
-      Cphantom_let (id, exp, map_tail f body)
+      Cphantom_let (id, exp, f body)
   | Cifthenelse(cond, ifso_dbg, ifso, ifnot_dbg, ifnot, dbg) ->
       Cifthenelse
         (
           cond,
-          ifso_dbg, map_tail f ifso,
-          ifnot_dbg, map_tail f ifnot,
+          ifso_dbg, f ifso,
+          ifnot_dbg, f ifnot,
           dbg
         )
   | Csequence(e1, e2) ->
-      Csequence(e1, map_tail f e2)
+      Csequence(e1, f e2)
   | Cswitch(e, tbl, el, dbg') ->
-      Cswitch(e, tbl, Array.map (fun (e, dbg) -> map_tail f e, dbg) el, dbg')
+      Cswitch(e, tbl, Array.map (fun (e, dbg) -> f e, dbg) el, dbg')
   | Ccatch(rec_flag, handlers, body) ->
-      let map_h (n, ids, handler, dbg) = (n, ids, map_tail f handler, dbg) in
-      Ccatch(rec_flag, List.map map_h handlers, map_tail f body)
+      let map_h (n, ids, handler, dbg) = (n, ids, f handler, dbg) in
+      Ccatch(rec_flag, List.map map_h handlers, f body)
   | Ctrywith(e1, id, e2, dbg) ->
-      Ctrywith(map_tail f e1, id, map_tail f e2, dbg)
+      Ctrywith(f e1, id, f e2, dbg)
   | Cregion e ->
-      Cregion(map_tail f e)
+      Cregion(f e)
   | Ctail e ->
-      Ctail(map_tail f e)
+      Ctail(f e)
   | Cexit _ | Cop (Craise _, _, _) as cmm ->
       cmm
   | Cconst_int _
@@ -309,8 +309,22 @@ let rec map_tail f = function
   | Cvar _
   | Cassign _
   | Ctuple _
-  | Cop _ as c ->
-      f c
+  | Cop _ as cmm -> cmm
+
+let map_tail f =
+  let rec loop = function
+    | Cconst_int _
+    | Cconst_natint _
+    | Cconst_float _
+    | Cconst_symbol _
+    | Cvar _
+    | Cassign _
+    | Ctuple _
+    | Cop _ as c ->
+        f c
+    | cmm -> map_shallow_tail loop cmm
+  in
+  loop
 
 let iter_shallow f = function
   | Clet (_id, e1, e2) ->

asmcomp/cmm.mli

@@ -242,11 +242,15 @@ val iter_shallow_tail: (expression -> unit) -> expression -> bool
       considered to be in tail position (because their result become
       the final result for the expression).  *)
 
+val map_shallow_tail: (expression -> expression) -> expression -> expression
+  (** Apply the transformation to those immediate sub-expressions of an
+      expression that are in tail position, using the same definition of "tail"
+      as [iter_shallow_tail] *)
+
 val map_tail: (expression -> expression) -> expression -> expression
   (** Apply the transformation to an expression, trying to push it
-      to all inner sub-expressions that can produce the final result.
-      Same disclaimer as for [iter_shallow_tail] about the notion
-      of "tail" sub-expression. *)
+      to all inner sub-expressions that can produce the final result,
+      by recursively applying map_shallow_tail *)
 
 val iter_shallow: (expression -> unit) -> expression -> unit
   (** Apply the transformation to each immediate sub-expression. *)

asmcomp/cmm_helpers.ml

@@ -807,12 +807,16 @@ let lookup_label obj lab dbg =
     let table = Cop (Cload (Word_val, Mutable), [obj], dbg) in
     addr_array_ref table lab dbg)
 
-let call_cached_method obj tag cache pos args apos dbg =
+let send_function_name n (mode : Lambda.alloc_mode) =
+  let suff = match mode with Alloc_heap -> "" | Alloc_local -> "L" in
+  "caml_send" ^ Int.to_string n ^ suff
+
+let call_cached_method obj tag cache pos args (apos,mode) dbg =
   let arity = List.length args in
   let cache = array_indexing log2_size_addr cache pos dbg in
-  Compilenv.need_send_fun arity;
+  Compilenv.need_send_fun arity mode;
   Cop(Capply(typ_val, apos),
-      Cconst_symbol("caml_send" ^ Int.to_string arity, dbg) ::
+      Cconst_symbol(send_function_name arity mode, dbg) ::
         obj :: tag :: cache :: args,
       dbg)
 
@@ -859,10 +863,13 @@ let make_checkbound dbg = function
       Cop(Ccheckbound, args, dbg)
 
 (* Record application and currying functions *)
-
-let apply_function_sym n =
+let apply_function_name (n, (mode : Lambda.alloc_mode)) =
+  let suff = match mode with Alloc_heap -> "" | Alloc_local -> "L" in
+  "caml_apply" ^ Int.to_string n ^ suff
+let apply_function_sym n mode =
   assert (n > 0);
-  Compilenv.need_apply_fun n; "caml_apply" ^ Int.to_string n
+  Compilenv.need_apply_fun n mode;
+  apply_function_name (n, mode)
 let curry_function_sym ar =
   Compilenv.need_curry_fun ar;
   match ar with
@@ -1713,10 +1720,10 @@ let ptr_offset ptr offset dbg =
   then ptr
   else Cop(Caddv, [ptr; Cconst_int(offset * size_addr, dbg)], dbg)
 
-let direct_apply lbl args pos dbg =
+let direct_apply lbl args (pos, _mode) dbg =
   Cop(Capply(typ_val, pos), Cconst_symbol (lbl, dbg) :: args, dbg)
 
-let generic_apply mut clos args pos dbg =
+let generic_apply mut clos args (pos, mode) dbg =
   match args with
   | [arg] ->
       bind "fun" clos (fun clos ->
@@ -1725,23 +1732,23 @@ let generic_apply mut clos args pos dbg =
   | _ ->
       let arity = List.length args in
       let cargs =
-        Cconst_symbol(apply_function_sym arity, dbg) :: args @ [clos]
+        Cconst_symbol(apply_function_sym arity mode, dbg) :: args @ [clos]
       in
       Cop(Capply(typ_val, pos), cargs, dbg)
 
-let send kind met obj args apos dbg =
+let send kind met obj args akind dbg =
   let call_met obj args clos =
     (* met is never a simple expression, so it never gets turned into an
        Immutable load *)
-    generic_apply Asttypes.Mutable clos (obj :: args) apos dbg
+    generic_apply Asttypes.Mutable clos (obj :: args) akind dbg
   in
   bind "obj" obj (fun obj ->
       match (kind : Lambda.meth_kind), args with
         Self, _ ->
           bind "met" (lookup_label obj met dbg)
             (call_met obj args)
       | Cached, cache :: pos :: args ->
-          call_cached_method obj met cache pos args apos dbg
+          call_cached_method obj met cache pos args akind dbg
       | _ ->
           bind "met" (lookup_tag obj met dbg)
             (call_met obj args))
@@ -1807,23 +1814,51 @@ let cache_public_method meths tag cache dbg =
     Csequence(Cop (Cstore (Word_int, Assignment), [cache; Cvar tagged], dbg),
               Cvar tagged)))))
 
-let region e =
-  (* [Cregion e] is equivalent to [e] if [e] contains no local allocs *)
-  let rec check_local_allocs = function
-    | Cregion _ ->
-       (* Local allocations within a nested region do not affect this region.
-          Note this prevents O(n^2) behaviour with many nested regions. *)
-       ()
+let has_local_allocs e =
+  let rec loop = function
+    | Cregion e ->
+        (* Local allocations within a nested region do not affect this region,
+           except inside a Ctail block *)
+        loop_until_tail e
     | Cop (Calloc Alloc_local, _, _)
     | Cop ((Cextcall _ | Capply _), _, _) ->
-       raise Exit
+        raise Exit
+    | e ->
+        iter_shallow loop e
+  and loop_until_tail = function
+    | Ctail e -> loop e
+    | Cregion _ -> ()
+    | e -> ignore (iter_shallow_tail loop_until_tail e)
+  in
+  match loop e with
+  | () -> false
+  | exception Exit -> true
+
+let remove_region_tail e =
+  let rec has_tail = function
+    | Ctail _
+    | Cop(Capply(_, Apply_tail), _, _) -> raise Exit
+    | Cregion _ -> ()
+    | e -> ignore (iter_shallow_tail has_tail e)
+  in
+  let rec remove_tail = function
+    | Ctail e -> e
+    | Cop(Capply(mach, Apply_tail), args, dbg) ->
+       Cop(Capply(mach, Apply_nontail), args, dbg)
+    | Cregion _ as e -> e
     | e ->
-       iter_shallow check_local_allocs e
+       map_shallow_tail remove_tail e
   in
-  match check_local_allocs e with
+  match has_tail e with
   | () -> e
-  | exception Exit -> Cregion e
+  | exception Exit -> remove_tail e
 
+let region e =
+  (* [Cregion e] is equivalent to [e] if [e] contains no local allocs *)
+  if has_local_allocs e then
+    Cregion e
+  else
+    remove_region_tail e
 
 (* CR mshinwell: These will be filled in by later pull requests. *)
 let placeholder_dbg () = Debuginfo.none
@@ -1840,7 +1875,7 @@ let placeholder_fun_dbg ~human_name:_ = Debuginfo.none
            (app closN-1.code aN closN-1))))
 *)
 
-let apply_function_body arity =
+let apply_function_body (arity, (mode : Lambda.alloc_mode)) =
   let dbg = placeholder_dbg in
   let arg = Array.make arity (V.create_local "arg") in
   for i = 1 to arity - 1 do arg.(i) <- V.create_local "arg" done;
@@ -1875,13 +1910,15 @@ let apply_function_body arity =
        :: List.map (fun s -> Cvar s) all_args,
        dbg ()),
    dbg (),
-   app_fun clos 0,
+   (match mode with
+    | Alloc_heap -> Cregion (app_fun clos 0)
+    | Alloc_local -> app_fun clos 0),
    dbg ()))
 
-let send_function arity =
+let send_function (arity, mode) =
   let dbg = placeholder_dbg in
   let cconst_int i = Cconst_int (i, dbg ()) in
-  let (args, clos', body) = apply_function_body (1+arity) in
+  let (args, clos', body) = apply_function_body (1+arity, mode) in
   let cache = V.create_local "cache"
   and obj = List.hd args
   and tag = V.create_local "tag" in
@@ -1915,7 +1952,7 @@ let send_function arity =
   in
   let body = Clet(VP.create clos', clos, body) in
   let cache = cache in
-  let fun_name = "caml_send" ^ Int.to_string arity in
+  let fun_name = send_function_name arity mode in
   let fun_args =
     [obj, typ_val; tag, typ_int; cache, typ_val]
     @ List.map (fun id -> (id, typ_val)) (List.tl args) in
@@ -1931,7 +1968,7 @@ let send_function arity =
 let apply_function arity =
   let (args, clos, body) = apply_function_body arity in
   let all_args = args @ [clos] in
-  let fun_name = "caml_apply" ^ Int.to_string arity in
+  let fun_name = apply_function_name arity in
   let fun_dbg = placeholder_fun_dbg ~human_name:fun_name in
   Cfunction
    {fun_name;
@@ -2129,26 +2166,27 @@ let curry_function = function
      assert (n > 0);
      intermediate_curry_functions ~nlocal ~arity:n 0
 
-module Int = Numbers.Int
+module ApplyFnSet =
+  Set.Make (struct type t = int * Lambda.alloc_mode let compare = compare end)
 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)
+let default_apply = ApplyFnSet.of_list [2,Alloc_heap; 3,Alloc_heap]
   (* These apply funs are always present in the main program because
      the run-time system needs them (cf. runtime/<arch>.S) . *)
 
 let generic_functions shared units =
   let (apply,send,curry) =
     List.fold_left
       (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 ApplyFnSet.add ui.ui_apply_fun apply,
+         List.fold_right ApplyFnSet.add ui.ui_send_fun send,
          List.fold_right AritySet.add ui.ui_curry_fun curry)
-      (Int.Set.empty,Int.Set.empty,AritySet.empty)
+      (ApplyFnSet.empty,ApplyFnSet.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
+  let apply = if shared then apply else ApplyFnSet.union apply default_apply in
+  let accu = ApplyFnSet.fold (fun nr accu -> apply_function nr :: accu) apply [] in
+  let accu = ApplyFnSet.fold (fun nr accu -> send_function nr :: accu) send accu in
   AritySet.fold (fun arity accu -> curry_function arity @ accu) curry accu
 
 (* Primitives *)

asmcomp/cmm_helpers.mli

@@ -281,7 +281,7 @@ val lookup_label : expression -> expression -> Debuginfo.t -> expression
     - args : the additional arguments to the method call *)
 val call_cached_method :
   expression -> expression -> expression -> expression -> expression list ->
-  Lambda.apply_position -> Debuginfo.t -> expression
+  Clambda.apply_kind -> Debuginfo.t -> expression
 
 (** Allocations *)
 
@@ -311,7 +311,7 @@ val check_bound :
 
 (** 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 -> string
+val apply_function_sym : int -> 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. *)
@@ -543,7 +543,7 @@ val ptr_offset : expression -> int -> Debuginfo.t -> expression
 
 (** Direct application of a function via a symbol *)
 val direct_apply :
-  string -> expression list -> Lambda.apply_position
+  string -> expression list -> Clambda.apply_kind
   -> Debuginfo.t -> expression
 
 (** Generic application of a function to one or several arguments.
@@ -553,7 +553,7 @@ val direct_apply :
     the currently defined closure. *)
 val generic_apply :
   Asttypes.mutable_flag -> expression -> expression list
-  -> Lambda.apply_position -> Debuginfo.t -> expression
+  -> Clambda.apply_kind -> Debuginfo.t -> expression
 
 (** Method call : [send kind met obj args dbg]
     - [met] is a method identifier, which can be a hashed variant or an index
@@ -564,7 +564,7 @@ val generic_apply :
     cache and cache position) *)
 val send :
   Lambda.meth_kind -> expression -> expression -> expression list
-  -> Lambda.apply_position -> Debuginfo.t -> expression
+  -> Clambda.apply_kind -> Debuginfo.t -> expression
 
 (** Construct [Cregion e], eliding some useless regions *)
 val region : expression -> expression

asmcomp/cmmgen.ml

@@ -171,6 +171,10 @@ let rec expr_size env = function
       | RHS_block blocksize -> RHS_infix { blocksize; offset }
       | RHS_nonrec -> RHS_nonrec
       | _ -> assert false)
+  | Uregion exp ->
+      expr_size env exp
+  | Utail _ ->
+      Misc.fatal_error "Utail in non-tail position"
   | _ -> RHS_nonrec
 
 (* Translate structured constants to Cmm data items *)
@@ -422,13 +426,13 @@ let rec transl env e =
       let ptr = transl env arg in
       let dbg = Debuginfo.none in
       ptr_offset ptr offset dbg
-  | Udirect_apply(lbl, args, pos, dbg) ->
+  | Udirect_apply(lbl, args, kind, dbg) ->
       let args = List.map (transl env) args in
-      direct_apply lbl args pos dbg
-  | Ugeneric_apply(clos, args, pos, dbg) ->
+      direct_apply lbl args kind dbg
+  | 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 pos dbg
+      generic_apply (mut_from_env env clos) clos args kind dbg
   | Usend(kind, met, obj, args, pos, dbg) ->
       let met = transl env met in
       let obj = transl env obj in