flambda-backend: Fix Translcore to look for allocations in exclaves (#1495)

Author: lukemaurer

lambda/translcore.ml

@@ -145,8 +145,8 @@ let may_allocate_in_region lam =
   let rec loop_region lam =
     shallow_iter ~tail:(function
       | Lexclave body -> loop body
-      | _ -> ()
-    ) ~non_tail:(fun _ -> ()) lam
+      | lam -> loop_region lam
+    ) ~non_tail:(fun lam -> loop_region lam) lam
   and loop = function
     | Lvar _ | Lmutvar _ | Lconst _ -> ()
 

testsuite/tests/typing-local/regions.ml

@@ -1,4 +1,5 @@
 (* TEST
+   include ocamlcommon
    * native *)
 
 external local_stack_offset : unit -> int = "caml_local_stack_offset"
@@ -246,4 +247,30 @@ let () =
   create_and_ignore ();
   check_empty "mode-crossed region"
 
+let[@inline never] allocate_in_exclave a =
+  (* This needs to be a [while] so that we're allowed to have an [exclave_] *)
+  while
+    let pair = local_ opaque_identity (a, a) in
+    let b, _ = pair in
+    let b : int = b in
+    exclave_ (
+      (* This should allocate in the function's region - in particular, the
+         function needs to have one *)
+      let pair = local_ opaque_identity (b, b) in
+      let c, _ = pair in
+      c = 42)
+  do
+    ()
+  done
+
+let () =
+  let () =
+    (* Temporarily disable this test for flambda1 until it's fixed there *)
+    if not Config.flambda then
+      allocate_in_exclave 1
+  in
+  (* If [allocate_in_exclave] had its region elided, the allocation will have
+     happened in our region instead *)
+  check_empty "allocation from exclave"
+
 let () = Gc.compact ()

testsuite/tests/typing-local/regions.reference

@@ -27,3 +27,4 @@
 dynamic/partial overapply: OK
          method overapply: OK
       mode-crossed region: OK
+  allocation from exclave: OK