flambda-backend: Add type-directed disambiguation for local_ expr e…

…xpressions (#2303)

* Add regression test

* Fix

testsuite/tests/typing-local/local.ml

@@ -2806,11 +2806,11 @@ let foo () =
   let _bar : int -> int -> int = local_ (fun x y -> x + y) in
   ()
 [%%expect{|
-Line 2, characters 33-58:
+Line 2, characters 40-58:
 2 |   let _bar : int -> int -> int = local_ (fun x y -> x + y) in
-                                     ^^^^^^^^^^^^^^^^^^^^^^^^^
-Error: This expression has type int -> local_ (int -> int)
-       but an expression was expected of type int -> (int -> int)
+                                            ^^^^^^^^^^^^^^^^^^
+Error: This function or one of its parameters escape their region
+       when it is partially applied.
 |}];;
 
 (* test that [function] checks all its branches either for local_ or the
@@ -2859,3 +2859,33 @@ Line 1, characters 16-22:
                     ^^^^^^
 Error: The locality axis has already been specified.
 |}]
+
+(* type-directed disambiguation *)
+
+module M = struct
+  type t = M_constructor
+end
+
+let foo (local_ _ : M.t) = ();;
+let foo_f (local_ _ : M.t -> unit) = ();;
+[%%expect{|
+module M : sig type t = M_constructor end
+val foo : local_ M.t -> unit = <fun>
+val foo_f : local_ (M.t -> unit) -> unit = <fun>
+|}]
+
+let () = foo M_constructor
+[%%expect{|
+|}]
+
+let () = foo_f (fun M_constructor -> ())
+[%%expect{|
+|}]
+
+let () = foo (local_ M_constructor)
+[%%expect{|
+|}]
+
+let () = foo_f (local_ (fun M_constructor -> ()))
+[%%expect{|
+|}]

typing/typecore.ml

@@ -4610,6 +4610,10 @@ let rec is_inferred sexp =
   match Jane_syntax.Expression.of_ast sexp with
   | Some (jexp, _attrs) -> is_inferred_jane_syntax jexp
   | None      -> match sexp.pexp_desc with
+  | Pexp_apply
+      ({ pexp_desc = Pexp_extension({txt; _}, _) },
+       [Nolabel, exp]) when txt = Jane_syntax.Mode_expr.extension_name ->
+      is_inferred exp
   | Pexp_ident _ | Pexp_apply _ | Pexp_field _ | Pexp_constraint _
   | Pexp_coerce _ | Pexp_send _ | Pexp_new _ -> true
   | Pexp_sequence (_, e) | Pexp_open (_, e) -> is_inferred e