stack_ support poly-returning primitives

testsuite/tests/typing-modes/stack.ml

@@ -299,3 +299,22 @@ Line 3, characters 2-5:
 Error: This value escapes its region.
   Hint: Cannot return a local value without an "exclave_" annotation.
 |}]
+
+(* all primitives with poly return mode are supported by stack_ *)
+let mk () =
+  let r = stack_ (ref "hello") in
+  r
+[%%expect{|
+Line 3, characters 2-3:
+3 |   r
+      ^
+Error: This value escapes its region.
+  Hint: Cannot return a local value without an "exclave_" annotation.
+|}]
+
+let mk () = exclave_
+  let r = stack_ (ref "hello") in
+  r
+[%%expect{|
+val mk : unit -> local_ string ref = <fun>
+|}]

typing/typecore.ml

@@ -6782,28 +6782,44 @@ and type_expect_
       let unsupported category =
         raise (Error (exp.exp_loc, env, Unsupported_stack_allocation category))
       in
+      let force_local ~context locality =
+        match Locality.submode Locality.local locality with
+        | Ok () -> ()
+        | Error _ -> raise (Error (e.pexp_loc, env,
+            Cannot_stack_allocate context))
+      in
+      let not_allocation () =
+        raise (Error (exp.exp_loc, env, Not_allocation))
+      in
       begin match exp.exp_desc with
       | Texp_function { alloc_mode; _} | Texp_tuple (_, alloc_mode)
       | Texp_construct (_, _, _, Some alloc_mode)
       | Texp_variant (_, Some (_, alloc_mode))
       | Texp_record {alloc_mode = Some alloc_mode; _}
       | Texp_array (_, _, _, alloc_mode)
       | Texp_field (_, _, _, Boxing (alloc_mode, _), _) ->
-        begin match Locality.submode Locality.local
-          (Alloc.proj (Comonadic Areality) alloc_mode.mode) with
-        | Ok () -> ()
-        | Error _ -> raise (Error (e.pexp_loc, env,
-            Cannot_stack_allocate alloc_mode.locality_context))
-        end
+          force_local ~context:alloc_mode.locality_context
+            (Alloc.proj (Comonadic Areality) alloc_mode.mode)
+      | Texp_apply ({exp_desc = Texp_ident (_, _, desc, _, _)}, _, _, _, _) ->
+          begin match desc.val_kind with
+          | Val_prim prim ->
+              let _, mode, _ = instance_prim prim desc.val_type in
+              begin match prim.prim_native_repr_res, mode with
+              (* if the locality of returned value of the primitive is poly, we
+                treat it as an allocation *)
+              | (Prim_poly, _), Some mode -> force_local ~context:None mode
+              | _ -> not_allocation ()
+              end;
+          | _ -> not_allocation ()
+          end
       | Texp_list_comprehension _ -> unsupported List_comprehension
       | Texp_array_comprehension _ -> unsupported Array_comprehension
       | Texp_new _ -> unsupported Object
       | Texp_override _ -> unsupported Object
       | Texp_lazy _ -> unsupported Lazy
       | Texp_object _ -> unsupported Object
       | Texp_pack _ -> unsupported Module
-      | _ ->
-        raise (Error (exp.exp_loc, env, Not_allocation))
+      | _ -> not_allocation ()
       end;
       submode ~loc ~env (Value.min_with (Comonadic Areality) Regionality.local)
         expected_mode;