Fix function return layout bug

chamelon/compat.jst.ml

@@ -69,6 +69,8 @@ type texp_function_cases_identifier = {
   last_arg_sort : Jkind.Sort.t;
   last_arg_exp_extra : exp_extra option;
   last_arg_attributes : attributes;
+  env : Env.t;
+  ret_type : Types.type_expr;
 }
 
 type texp_function_body =
@@ -99,6 +101,8 @@ let texp_function_cases_identifier_defaults =
     last_arg_sort = Jkind.Sort.value;
     last_arg_exp_extra = None;
     last_arg_attributes = [];
+    env = Env.empty;
+    ret_type = Ctype.newvar (Jkind.any ~why:Dummy_jkind);
   }
 
 let texp_function_param_identifier_defaults =
@@ -159,6 +163,8 @@ let mkTexp_function ?(id = texp_function_defaults)
                 fc_cases = cases;
                 fc_param = param;
                 fc_partial = partial;
+                fc_env = id.env;
+                fc_ret_type = id.ret_type;
                 fc_arg_mode = id.last_arg_mode;
                 fc_arg_sort = id.last_arg_sort;
                 fc_exp_extra = id.last_arg_exp_extra;
@@ -263,6 +269,8 @@ let view_texp (e : expression_desc) =
                     last_arg_sort = cases.fc_arg_sort;
                     last_arg_exp_extra = cases.fc_exp_extra;
                     last_arg_attributes = cases.fc_attributes;
+                    env = cases.fc_env;
+                    ret_type = cases.fc_ret_type;
                   };
               }
       in

ocaml/lambda/translcore.ml

@@ -1293,14 +1293,9 @@ and transl_function_without_attributes
     match body with
     | Tfunction_body exp ->
         layout_exp return_sort exp
-    | Tfunction_cases { fc_cases = { c_rhs; _ } :: _ } ->
-        layout_exp return_sort c_rhs
-    | Tfunction_cases { fc_cases = [] } ->
-        (* ppxes can generate empty function cases, which compiles to
-           a function that always raises Match_failure. We try less
-           hard to calculate a detailed layout that the middle-end can
-           use for optimizations. *)
-        layout_of_sort loc return_sort
+    | Tfunction_cases cases ->
+        layout cases.fc_env cases.fc_loc return_sort cases.fc_ret_type
+
   in
   match
     transl_tupled_function ~scopes loc params body
@@ -2031,7 +2026,8 @@ and transl_letop ~scopes loc env let_ ands param param_sort case case_sort
                 { fc_cases = [case]; fc_param = param; fc_partial = partial;
                   fc_loc = ghost_loc; fc_exp_extra = None; fc_attributes = [];
                   fc_arg_mode = Mode.Alloc.disallow_right Mode.Alloc.legacy;
-                  fc_arg_sort = param_sort;
+                  fc_arg_sort = param_sort; fc_env = env;
+                  fc_ret_type = case.c_rhs.exp_type;
                 }))
     in
     let attr = function_attribute_disallowing_arity_fusion in

ocaml/testsuite/tests/syntactic-arity/flambda_kind.ml

@@ -0,0 +1,25 @@
+(* TEST
+   flambda2;
+   native;
+*)
+
+(* This is a regression test, see PR #2471 in ocaml-flambda/flambda-backend *)
+
+[@@@ocaml.flambda_o3]
+
+type _ value =
+  | Int : int value
+  | Float : float value
+
+let[@inline never] get (type a) : a value -> a = function
+  | Int -> 3
+  | Float -> 3.
+
+let[@inline] update (type a) (v : a value) (x : a) : a =
+  match v with
+  | Int -> x + 1
+  | Float -> x +. 1.
+
+let run x = update x (get x)
+
+let (_ : float) = run Float

ocaml/testsuite/tests/typing-layouts/basics_alpha.ml

@@ -376,9 +376,9 @@ let id5 : 'a void5 -> 'a void5 = function
 
 [%%expect{|
 type ('a : void) void5 = Void5 of 'a
-Line 4, characters 15-22:
+Lines 3-4, characters 33-22:
+3 | .................................function
 4 |   | Void5 x -> Void5 x
-                   ^^^^^^^
 Error: Non-value detected in [value_kind].
        Please report this error to the Jane Street compilers team.
        The layout of 'a is void, because

ocaml/testsuite/tests/typing-layouts/void_alpha.ml

@@ -211,14 +211,17 @@ type void_variant =
     }
 val r : '_weak2 list ref = {contents = []}
 val cons_r : '_weak2 -> unit = <fun>
-Lines 19-25, characters 5-23:
-19 | .....A ((cons_r 10; a1),
+Lines 17-35, characters 10-27:
+17 | ..........function
+18 |   | A (a1, a2, x, v, z, b1, b2) ->
+19 |      A ((cons_r 10; a1),
 20 |         (cons_r 8; {v = ((cons_r 9; a2).v)}),
 21 |         (cons_r 7; x),
-22 |         (cons_r 5; {v = ((cons_r 6; v).v)}),
-23 |         (cons_r 4; z),
-24 |         (cons_r 2; {v = ((cons_r 3; b1).v)}),
-25 |         (cons_r 1; b2))
+...
+32 |        v = (cons_r 5; {v = ((cons_r 6; v).v)});
+33 |        z = (cons_r 4; z);
+34 |        b1 = (cons_r 2; {v = ((cons_r 3; b1).v)});
+35 |        b2 = (cons_r 1; b2)}
 Error: Non-value detected in [value_kind].
        Please report this error to the Jane Street compilers team.
        The layout of t_void is void, because

ocaml/typing/printtyped.ml

@@ -367,7 +367,7 @@ and function_body i ppf (body : function_body) =
       expression (i+1) ppf e
   | Tfunction_cases
       { fc_cases; fc_loc; fc_exp_extra; fc_attributes; fc_arg_mode;
-        fc_arg_sort; fc_param = _; fc_partial = _; }
+        fc_arg_sort; fc_param = _; fc_partial = _; fc_env = _; fc_ret_type = _ }
     ->
       line i ppf "Tfunction_cases %a\n" fmt_location fc_loc;
       alloc_mode i ppf fc_arg_mode;

ocaml/typing/tast_iterator.ml

@@ -298,11 +298,12 @@ let function_body sub body =
   match body with
   | Tfunction_body body ->
       sub.expr sub body
-  | Tfunction_cases { fc_cases; fc_exp_extra; fc_loc; fc_attributes } ->
+  | Tfunction_cases { fc_cases; fc_exp_extra; fc_loc; fc_attributes; fc_env } ->
       List.iter (sub.case sub) fc_cases;
       Option.iter (extra sub) fc_exp_extra;
       sub.location sub fc_loc;
-      sub.attributes sub fc_attributes
+      sub.attributes sub fc_attributes;
+      sub.env sub fc_env
 
 let expr sub {exp_loc; exp_extra; exp_desc; exp_env; exp_attributes; _} =
   let extra x = extra sub x in

ocaml/typing/tast_mapper.ml

@@ -374,15 +374,16 @@ let function_body sub body =
       Tfunction_body (sub.expr sub body)
   | Tfunction_cases
       { fc_cases; fc_partial; fc_param; fc_loc; fc_exp_extra; fc_attributes;
-        fc_arg_mode; fc_arg_sort; }
+        fc_arg_mode; fc_arg_sort; fc_env; fc_ret_type; }
     ->
       let fc_loc = sub.location sub fc_loc in
       let fc_attributes = sub.attributes sub fc_attributes in
       let fc_cases = List.map (sub.case sub) fc_cases in
       let fc_exp_extra = Option.map (extra sub) fc_exp_extra in
+      let fc_env = sub.env sub fc_env in
       Tfunction_cases
         { fc_cases; fc_partial; fc_param; fc_loc; fc_exp_extra; fc_attributes;
-          fc_arg_mode; fc_arg_sort; }
+          fc_arg_mode; fc_arg_sort; fc_env; fc_ret_type; }
 
 let expr sub x =
   let extra x = extra sub x in

ocaml/typing/typecore.ml

@@ -7465,6 +7465,7 @@ and type_argument ?explanation ?recarg env (mode : expected_mode) sarg
               body =
                 Tfunction_cases
                   { fc_cases = cases; fc_partial = Total; fc_param = param;
+                    fc_env = env; fc_ret_type = ty_res;
                     fc_loc = cases_loc; fc_exp_extra = None;
                     fc_attributes = []; fc_arg_mode = Alloc.disallow_right marg;
                     fc_arg_sort = arg_sort;
@@ -8205,6 +8206,8 @@ and type_function_cases_expect
         fc_param = param;
         fc_loc = loc;
         fc_exp_extra = None;
+        fc_env = env;
+        fc_ret_type = ty_ret;
         fc_attributes = [];
         fc_arg_mode = Alloc.disallow_right arg_mode;
         fc_arg_sort = arg_sort;

ocaml/typing/typedtree.ml

@@ -243,8 +243,10 @@ and function_body =
 
 and function_cases =
   { fc_cases: value case list;
+    fc_env : Env.t;
     fc_arg_mode: Mode.Alloc.l;
     fc_arg_sort: Jkind.sort;
+    fc_ret_type : Types.type_expr;
     fc_partial: partial;
     fc_param: Ident.t;
     fc_loc: Location.t;

ocaml/typing/typedtree.mli

@@ -465,8 +465,13 @@ and function_body =
 
 and function_cases =
   { fc_cases: value case list;
+    fc_env : Env.t;
+    (** [fc_env] contains entries from all parameters except
+        for the last one being matched by the cases.
+    *)
     fc_arg_mode: Mode.Alloc.l;
     fc_arg_sort: Jkind.sort;
+    fc_ret_type : Types.type_expr;
     fc_partial: partial;
     fc_param: Ident.t;
     fc_loc: Location.t;