flambda-backend: Rc_close_at_apply should only close the current region early (#2736)

Co-authored-by: Nathanaëlle Courant <nathanaelle.courant@ocamlpro.com>

Author: mshinwell

lambda/lambda.mli

@@ -86,6 +86,26 @@ type region_close =
   | Rc_nontail        (* do not close region, must not TCO *)
   | Rc_close_at_apply (* close region and tail call *)
 
+(** Notes about applytail (as emitted by Printlambda) a.k.a. Rc_close_at_apply:
+
+    applytail  / Rc_close_at_apply means that a call occurs in tail position of
+    the nearest enclosing region, and should be compiled by closing that region
+    (and only that region) just before control is transferred to the call.
+
+    In the raw lambda generated by Translcore, this can occur only in tail
+    position of a function, so these do happen to be tail calls. But after
+    inlining (by Simplif or otherwise), this may no longer be the case.
+
+    In particular, this code pattern:
+
+      (function (region (.... (region ... (applytail f)))))
+
+    means that the applytail is in tail position of the inner region and closes
+    that one. It just so happens to be at the end of the outer function as well,
+    but it does not mean that it's a tail call of that function. (It's not a
+    tail call because the outer region needs to end there.)
+*)
+
 (* CR layouts v5: When we add more blocks of non-scannable values, consider
    whether some of the primitives specific to ufloat records
    ([Pmakeufloatblock], [Pufloatfield], and [Psetufloatfield]) can/should be

testsuite/tests/typing-local/close_at_apply.ml

@@ -0,0 +1,57 @@
+(* TEST
+   stack-allocation;
+   flambda2;
+   {
+     ocamlopt_flags="-Oclassic";
+     native;
+   }{
+     ocamlopt_flags="-O3";
+     native;
+   }
+*)
+
+external local_stack_offset : unit -> int = "caml_local_stack_offset"
+
+external ignore : ('a[@local_opt]) -> unit = "%ignore"
+external opaque_identity : ('a[@local_opt]) -> ('a[@local_opt]) = "%opaque"
+
+external ( +. )
+  :  (float[@local_opt])
+  -> (float[@local_opt])
+  -> (float[@local_opt])
+  = "%addfloat"
+
+let locally_allocate =
+  let f = 1.0 in
+  fun [@inline never] () -> ignore (opaque_identity (f +. 1.0) : float)
+;;
+
+let saved_stack_offset = ref 0
+
+let foo list =
+  (* The Simplif local function optimization will transform [run] and
+     [wrapper] into continuations. *)
+  let[@inline never][@local] run f = f () in
+  let[@inline never][@local] wrapper f =
+    match Sys.opaque_identity true with
+    | true ->
+      (* Make a local allocation *)
+      ignore (opaque_identity (opaque_identity 1.0 +. 1.0));
+      (* This tail call gets moved into non-tail position after the above
+         inlining transformation.  As such it's important that its
+         [Rc_close_at_apply] semantics is respected carefully: only the
+         current region, and not the parent one, must be closed (since [f]
+         is locally allocated in that parent region). *)
+      f ()
+    | false -> run (fun () -> f ()) [@nontail]
+  in
+  let f = fun _ ->
+    assert (local_stack_offset () = !saved_stack_offset);
+    locally_allocate ();
+    Stdlib.List.hd list
+  in
+  saved_stack_offset := local_stack_offset ();
+  wrapper f [@nontail]
+;;
+
+let _ : int = foo [ 1 ]