Treat Prim_poly primitives as Alloc_heap in Lambda

Author: mshinwell

ocaml/lambda/lambda.ml

@@ -1445,16 +1445,19 @@ let alloc_mode_of_primitive_description (p : Primitive.description) =
     if p.prim_alloc then Some alloc_heap else None
   else
     match p.prim_native_repr_res with
-    | (Prim_local | Prim_poly), _ ->
+    | Prim_local, _ ->
       (* 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, _ ->
+    | (Prim_global | Prim_poly), _ ->
       (* For primitives that definitely do not allocate locally,
          [p.prim_alloc = false] actually tells us that the primitive does
-         not allocate at all. *)
+         not allocate at all.
+
+         No external call that is [Prim_poly] may allocate locally.
+      *)
       if p.prim_alloc then Some alloc_heap else None
 
 (* Changes to this function may also require changes in Flambda 2 (e.g.

ocaml/lambda/printlambda.ml

@@ -88,6 +88,10 @@ let alloc_mode = function
   | Alloc_heap -> ""
   | Alloc_local -> "local"
 
+let alloc_mode' = function
+  | Alloc_heap -> "heap"
+  | Alloc_local -> "local"
+
 let boxed_integer_name = function
   | Pnativeint -> "nativeint"
   | Pint32 -> "int32"

ocaml/lambda/printlambda.mli

@@ -36,3 +36,4 @@ val print_bigarray :
   string -> bool -> Lambda.bigarray_kind -> formatter ->
   Lambda.bigarray_layout -> unit
 val check_attribute : formatter -> check_attribute -> unit
+val alloc_mode' : alloc_mode -> string

ocaml/lambda/translprim.ml

@@ -909,7 +909,11 @@ let lambda_of_prim prim_name prim loc args arg_exps =
 let check_primitive_arity loc p =
   let mode =
     match p.prim_native_repr_res with
-    | Prim_global, _ | Prim_poly, _ -> Some Mode.Locality.global
+    | Prim_global, _ | Prim_poly, _ ->
+      (* We assume all primitives are compiled to have the same arity for
+         different modes and types, so just pick one of the modes in the
+         [Prim_poly] case. *)
+      Some Mode.Locality.global
     | Prim_local, _ -> Some Mode.Locality.local
   in
   let prim = lookup_primitive loc mode Rc_normal p in
@@ -981,8 +985,42 @@ let transl_primitive loc p env ty ~poly_mode path =
          loc
      in
      let body = lambda_of_prim p.prim_name prim loc args None in
+     let alloc_mode = to_locality p.prim_native_repr_res in
+     let () =
+       (* CR mshinwell: Write a version of [primitive_may_allocate] that
+          works on the [prim] type. *)
+       match body with
+       | Lprim (prim, _, _) ->
+         (match Lambda.primitive_may_allocate prim with
+          | None ->
+            (* We don't check anything in this case; if the primitive doesn't
+               allocate, then after [Lambda] it will be translated to a term
+               not involving any region variables, meaning there would be
+               no concern about potentially unbound region variables. *)
+            ()
+          | Some lambda_alloc_mode ->
+            (* In this case we add a check to ensure the middle end has
+               the correct information as to whether a region was inserted
+               at this point. *)
+            match alloc_mode, lambda_alloc_mode with
+            | Alloc_heap, Alloc_heap
+            | Alloc_local, Alloc_local -> ()
+            | Alloc_local, Alloc_heap ->
+              (* This case is ok: the Lambda-derived information is more
+                 precise.  A region will be inserted, likely unused, and
+                 deleted by the middle end. *)
+              ()
+            | Alloc_heap, Alloc_local ->
+              Misc.fatal_errorf "Alloc mode incompatibility for:@ %a@ \
+                  (from to_locality, %s; from primitive_may_allocate, %s)"
+                Printlambda.lambda body
+                (Printlambda.alloc_mode' alloc_mode)
+                (Printlambda.alloc_mode' lambda_alloc_mode)
+         )
+       | _ -> ()
+     in
      let region =
-       match to_locality p.prim_native_repr_res with
+       match alloc_mode with
        | Alloc_heap -> true
        | Alloc_local -> false
      in

ocaml/testsuite/tests/typing-local/external.ml

@@ -0,0 +1,18 @@
+(* TEST
+   * flambda2
+   ** native
+*)
+
+module M : sig
+  val bits_of_float : float -> int64
+end = struct
+  external bits_of_float
+    :  (float[@local_opt])
+    -> (int64[@local_opt])
+    = "caml_int64_bits_of_float" "caml_int64_bits_of_float_unboxed"
+end
+
+let go_m f =
+  let i = M.bits_of_float f in
+  assert (i = 4L);
+  ()