flambda-backend: Closing over argument crosses modes (#2305)

Author: riaqn

testsuite/tests/typing-local/crossing.ml

@@ -449,4 +449,13 @@ let foo () =
   ()
 [%%expect{|
 val foo : unit -> unit = <fun>
-|}]
+|}]
+
+(* Values crosses modes when pattern-matching, an implication is that, closing
+   over [local int] won't force the closure to be [local]. *)
+let foo (local_ x : int) =
+  let bar y = x + y in
+  ref bar
+[%%expect{|
+val foo : local_ int -> (int -> int) ref = <fun>
+|}]

testsuite/tests/typing-local/local.ml

@@ -2783,6 +2783,13 @@ let f () = ((fun x -> (fun y -> "") [@extension.curry])
 val f : unit -> local_ string -> (string -> string) = <fun>
 |}];;
 
+(* mode crossing - the inner closure is [global] despite closing over [local_
+int] *)
+let f () = ((fun x y -> x + y) : (local_ int -> (int -> int)));;
+[%%expect{|
+val f : unit -> local_ int -> (int -> int) = <fun>
+|}];;
+
 (* Illegal: the expected mode is global *)
 let f () = local_ ((fun x y -> x + y) : (_ -> _));;
 [%%expect{|

typing/typecore.ml

@@ -804,6 +804,18 @@ let mode_cross_to_min env ty mode =
   else
     Value.disallow_right mode
 
+(* cross the monadic fragment to max, and the comonadic fragment to min *)
+let alloc_mode_cross_to_max_min env ty {monadic; comonadic} =
+  let (monadic, comonadic)=
+    if mode_cross env ty then
+      Alloc.Monadic.max, Alloc.Comonadic.min
+    else
+      monadic, comonadic
+  in
+  let monadic = Alloc.Monadic.disallow_left monadic in
+  let comonadic = Alloc.Comonadic.disallow_right comonadic in
+  {monadic; comonadic}
+
 let expect_mode_cross env ty (expected_mode : expected_mode) =
   if mode_cross env ty then
     { expected_mode with
@@ -6676,6 +6688,9 @@ and type_function
                   Final_arg
                 | Some fun_alloc_mode ->
                   assert(not is_final_val_param);
+                  (* If the argument cross modes, it crosses to max on monadic
+                     axes, and min on comonadic axes. *)
+                  let arg_mode = alloc_mode_cross_to_max_min env ty_arg arg_mode in
                   begin match
                     Alloc.submode (Alloc.close_over arg_mode) fun_alloc_mode
                   with