diff --git a/testsuite/tests/typing-modes/stack.ml b/testsuite/tests/typing-modes/stack.ml index 7f74720d85..0487bcb14e 100644 --- a/testsuite/tests/typing-modes/stack.ml +++ b/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 = +|}] diff --git a/typing/typecore.ml b/typing/typecore.ml index 307a968fd7..c7f4ab3e75 100644 --- a/typing/typecore.ml +++ b/typing/typecore.ml @@ -6782,6 +6782,15 @@ 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) @@ -6789,12 +6798,20 @@ and type_expect_ | 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 @@ -6802,8 +6819,7 @@ and type_expect_ | 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;