flambda-backend: Add region in Call_kind for C calls (#2180)

lambda/lambda.ml

@@ -1443,6 +1443,25 @@ let mod_field ?(read_semantics=Reads_agree) pos =
 let mod_setfield pos =
   Psetfield (pos, Pointer, Root_initialization)
 
+let alloc_mode_of_primitive_description (p : Primitive.description) =
+  if not Config.stack_allocation then
+    if p.prim_alloc then Some alloc_heap else None
+  else
+    match p.prim_native_repr_res with
+    | (Prim_local | Prim_poly), _ ->
+      (* For primitives that might allocate locally, [p.prim_alloc] just says
+         whether [caml_c_call] is required, without telling us anything
+         about local allocation.  (However if [p.prim_alloc = false] we
+         do actually know that the primitive does not allocate on the heap.) *)
+      Some alloc_local
+    | Prim_global, _ ->
+      (* For primitives that definitely do not allocate locally,
+         [p.prim_alloc = false] actually tells us that the primitive does
+         not allocate at all. *)
+      if p.prim_alloc then Some alloc_heap else None
+
+(* Changes to this function may also require changes in Flambda 2 (e.g.
+   closure_conversion.ml). *)
 let primitive_may_allocate : primitive -> alloc_mode option = function
   | Pbytes_to_string | Pbytes_of_string
   | Parray_to_iarray | Parray_of_iarray
@@ -1458,12 +1477,7 @@ let primitive_may_allocate : primitive -> alloc_mode option = function
   | Psetufloatfield _ -> None
   | Pduprecord _ -> Some alloc_heap
   | Pmake_unboxed_product _ | Punboxed_product_field _ -> None
-  | Pccall p ->
-     if not p.prim_alloc then None
-     else begin match p.prim_native_repr_res with
-       | (Prim_local|Prim_poly), _ -> Some alloc_local
-       | Prim_global, _ -> Some alloc_heap
-     end
+  | Pccall p -> alloc_mode_of_primitive_description p
   | Praise _ -> None
   | Psequor | Psequand | Pnot
   | Pnegint | Paddint | Psubint | Pmulint

lambda/lambda.mli

@@ -752,6 +752,10 @@ val primitive_may_allocate : primitive -> alloc_mode option
       revised.
   *)
 
+val alloc_mode_of_primitive_description :
+  Primitive.description -> alloc_mode option
+  (** Like [primitive_may_allocate], for [external] calls. *)
+
 (***********************)
 (* For static failures *)
 (***********************)

testsuite/tests/typing-local/alloc.heap.reference

@@ -23,6 +23,7 @@
       longarray: Allocation
   floatgenarray: Allocation
     longfgarray: Allocation
+    maniparray0: Allocation
      maniparray: Allocation
     manipfarray: Allocation
             ref: Allocation

testsuite/tests/typing-local/alloc.ml

@@ -9,6 +9,19 @@
       reference = "${test_source_directory}/alloc.heap.reference"
  *)
 
+(* First test to ensure that noalloc externals that locally allocate
+   don't cause a crash in the middle end (originally seen on
+   flambda-backend PR2180). *)
+
+(* This will never be called, caml_alloc_dummy is just chosen as a primitive that
+   exists in the bytecode runtime too *)
+external foo : unit -> ('a[@local_opt]) =
+  "caml_alloc_dummy" "caml_alloc_dummy" [@@noalloc]
+
+let foo () = foo ()
+
+(* Remaining tests *)
+
 type t = int
 
 type smallrecord = {  a : t; b : t; c : t }
@@ -264,6 +277,16 @@ external array_blit :
 external array_fill :
   local_ 'a array -> int -> int -> 'a -> unit = "caml_array_fill"
 
+let maniparray0 =
+  let l = [42] in
+  fun arr ->
+    (* This function should only locally allocate in the C runtime function
+       for doing the array allocation, and not in the OCaml code, in order
+       to ensure that locally-allocating C calls hold onto regions. *)
+    let x = local_array 6 l in
+    assert (x = arr);
+    ()
+
 let maniparray arr =              (* arr = 1,2,3,1,2,3 *)
   let x = local_array 2 [2] in    (* 2,2 *)
   let x = array_append x x in     (* 2,2,2,2 *)
@@ -472,6 +495,7 @@ let () =
   run "longarray" makelongarray 42;
   run "floatgenarray" makeshortarray 42.;
   run "longfgarray" makelongarray 42.;
+  run "maniparray0" maniparray0 [| [42]; [42]; [42]; [42]; [42]; [42] |];
   run "maniparray" maniparray [| [1]; [2]; [3]; [1]; [2]; [3] |];
   run "manipfarray" manipfarray [| 1.; 2.; 3.; 1.; 2.; 3. |];
   run "ref" makeref 42;

testsuite/tests/typing-local/alloc.stack.reference

@@ -23,6 +23,7 @@
       longarray: No Allocation
   floatgenarray: No Allocation
     longfgarray: No Allocation
+    maniparray0: No Allocation
      maniparray: No Allocation
     manipfarray: No Allocation
             ref: No Allocation