Prevent fatal errors related to regions (#2096)

Author: Ekdohibs

middle_end/flambda2/simplify/inlining/inlining_transforms.ml

@@ -114,38 +114,50 @@ let inline dacc ~apply ~unroll_to ~was_inline_always function_decl =
     | Need_meet -> Rec_info_expr.unknown
     | Invalid -> (* CR vlaviron: ? *) Rec_info_expr.do_not_inline
   in
-  let denv =
-    DE.enter_inlined_apply ~called_code:code ~apply ~was_inline_always denv
-  in
-  let params_and_body = Code.params_and_body code in
-  Function_params_and_body.pattern_match params_and_body
-    ~f:(fun
-         ~return_continuation
-         ~exn_continuation
-         params
-         ~body
-         ~my_closure
-         ~is_my_closure_used:_
-         ~my_region
-         ~my_depth
-         ~free_names_of_body:_
-       ->
-      let make_inlined_body () =
-        make_inlined_body ~callee ~region_inlined_into ~unroll_to
-          ~params:(Bound_parameters.to_list params)
-          ~args ~my_closure ~my_region ~my_depth ~rec_info ~body
-          ~exn_continuation ~return_continuation
-      in
-      let expr =
-        match Exn_continuation.extra_args apply_exn_continuation with
-        | [] ->
-          make_inlined_body ()
-            ~apply_exn_continuation:
-              (Exn_continuation.exn_handler apply_exn_continuation)
-            ~apply_return_continuation
-        | extra_args ->
-          wrap_inlined_body_for_exn_extra_args ~extra_args
-            ~apply_exn_continuation ~apply_return_continuation
-            ~result_arity:(Code.result_arity code) ~make_inlined_body
-      in
-      DA.with_denv dacc denv, expr)
+  match region_inlined_into, Code.result_mode code with
+  | Heap, Alloc_local ->
+    (* The alloc_mode of the application and of the code are incompatible. This
+       should have been prevented by the typer; therefore we are in GADT-caused
+       unreachable code; we replace the inlined body by [Invalid]. *)
+    ( dacc,
+      Expr.create_invalid
+        (Flambda.Invalid.Calling_local_returning_closure_with_normal_apply apply)
+    )
+  | Local _, Alloc_heap (* This is allowed by subtyping *)
+  | Local _, Alloc_local
+  | Heap, Alloc_heap ->
+    let denv =
+      DE.enter_inlined_apply ~called_code:code ~apply ~was_inline_always denv
+    in
+    let params_and_body = Code.params_and_body code in
+    Function_params_and_body.pattern_match params_and_body
+      ~f:(fun
+           ~return_continuation
+           ~exn_continuation
+           params
+           ~body
+           ~my_closure
+           ~is_my_closure_used:_
+           ~my_region
+           ~my_depth
+           ~free_names_of_body:_
+         ->
+        let make_inlined_body () =
+          make_inlined_body ~callee ~region_inlined_into ~unroll_to
+            ~params:(Bound_parameters.to_list params)
+            ~args ~my_closure ~my_region ~my_depth ~rec_info ~body
+            ~exn_continuation ~return_continuation
+        in
+        let expr =
+          match Exn_continuation.extra_args apply_exn_continuation with
+          | [] ->
+            make_inlined_body ()
+              ~apply_exn_continuation:
+                (Exn_continuation.exn_handler apply_exn_continuation)
+              ~apply_return_continuation
+          | extra_args ->
+            wrap_inlined_body_for_exn_extra_args ~extra_args
+              ~apply_exn_continuation ~apply_return_continuation
+              ~result_arity:(Code.result_arity code) ~make_inlined_body
+        in
+        DA.with_denv dacc denv, expr)

middle_end/flambda2/simplify/simplify_apply_expr.ml

@@ -421,14 +421,28 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
   let new_closure_alloc_mode, first_complex_local_param =
     if num_non_unarized_args <= first_complex_local_param
     then
+      (* At this point, we *have* to allocate the closure on the heap, even if
+         the alloc_mode of the application was local. Indeed, consider a
+         three-argument function, of type [string -> string -> string ->
+         string], coerced to [string -> local_ t] where [type t = string ->
+         string -> string].
+
+         If we apply this function twice to single arguments, the first
+         application will have a local alloc_mode. However, the second
+         application has a heap alloc_mode, and contains a reference to the
+         partial closure made by the first application. Due to this, the first
+         application must have a closure allocated on the heap as well, even
+         though it was with a local alloc_mode. *)
       ( Alloc_mode.For_allocations.heap,
         first_complex_local_param - num_non_unarized_args )
     else
       match (apply_alloc_mode : Alloc_mode.For_allocations.t) with
       | Heap ->
-        Misc.fatal_errorf "Partial application of %a with wrong mode at %s"
-          Code_id.print callee's_code_id
-          (Debuginfo.to_string (Apply.dbg apply))
+        Misc.fatal_errorf
+          "Partial application of %a with wrong mode:@.apply = \
+           %a@callee's_code_metadata = %a@."
+          Code_id.print callee's_code_id Apply.print apply Code_metadata.print
+          callee's_code_metadata
       | Local _ -> apply_alloc_mode, 0
   in
   (match closure_alloc_mode_from_type with
@@ -637,7 +651,7 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
         applied_values
       |> Value_slot.Map.of_list
     in
-    ( Set_of_closures.create ~value_slots apply_alloc_mode function_decls,
+    ( Set_of_closures.create ~value_slots new_closure_alloc_mode function_decls,
       dacc,
       code_id,
       code )

middle_end/flambda2/terms/flambda.ml

@@ -1432,6 +1432,7 @@ module Invalid = struct
     | Apply_cont_of_unreachable_continuation of Continuation.t
     | Defining_expr_of_let of Bound_pattern.t * Named.t
     | Closure_type_was_invalid of Apply_expr.t
+    | Calling_local_returning_closure_with_normal_apply of Apply_expr.t
     | Zero_switch_arms
     | Code_not_rebuilt
     | To_cmm_dummy_body
@@ -1456,6 +1457,11 @@ module Invalid = struct
       Format.asprintf
         "@[<hov 1>(Closure_type_was_invalid@ @[<hov 1>(apply_expr@ %a)@])@]"
         Apply_expr.print apply_expr
+    | Calling_local_returning_closure_with_normal_apply apply_expr ->
+      Format.asprintf
+        "@[<hov 1>(Calling_local_returning_closure_with_normal_apply@ @[<hov \
+         1>(apply_expr@ %a)@])@]"
+        Apply_expr.print apply_expr
     | Zero_switch_arms -> "Zero_switch_arms"
     | Code_not_rebuilt -> "Code_not_rebuilt"
     | To_cmm_dummy_body -> "To_cmm_dummy_body"

middle_end/flambda2/terms/flambda.mli

@@ -106,6 +106,7 @@ module Invalid : sig
     | Apply_cont_of_unreachable_continuation of Continuation.t
     | Defining_expr_of_let of Bound_pattern.t * named
     | Closure_type_was_invalid of Apply_expr.t
+    | Calling_local_returning_closure_with_normal_apply of Apply_expr.t
     | Zero_switch_arms
     | Code_not_rebuilt
     | To_cmm_dummy_body

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

@@ -0,0 +1,14 @@
+(* TEST *)
+
+type (_, _) eq = Eq : ('a, 'a) eq
+
+let[@inline always] cast (type a b) (x : a) (Eq : (a, b) eq) : b = x
+
+let test (f : local_ 'a -> local_ 'b) (eq : (local_ 'a -> local_ 'b, 'a -> 'b) eq) (x : 'a) : 'b =
+  (cast f eq) x 
+
+type 'a box = Box of 'a
+
+let localf : (local_ 'a -> local_ 'a box) = fun x -> exclave_ (local_ (Box x))
+
+let g eq x = test localf eq x

ocaml/testsuite/tests/typing-local/local_gadt.reference

@@ -0,0 +1,12 @@
+(* TEST *)
+
+type 'a box = Box of 'a
+module X : sig
+  val f : int -> string -> local_ int box
+end = struct
+  let[@inline never] f x y = exclave_ (local_ (Box x))
+end
+
+let[@inline always] h x = x
+
+let g a = (h X.f) a

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

@@ -0,0 +1,20 @@
+(* TEST *)
+
+type 'a box = Box of 'a
+
+let[@inline never] f x y z =
+  match x, y, z with Box x, Box y, Box z -> x + y + z
+
+type t = int box -> int box -> int
+let[@inline always] h x = x
+let g = (h f :> int box -> local_ t)
+
+let[@inline never] go a b c =
+  let local_ g1 = (g a) in
+  let g2 = (g1 b) in
+  g2 c
+
+let[@inline never] test () =
+  Format.eprintf "%d@." (go (Box 1) (Box 2) (Box 3))
+
+let () = test ()

ocaml/testsuite/tests/typing-local/partial_app_bug.reference

@@ -0,0 +1 @@
+6