flambda-backend: Don't get the layout of arguments from patterns (#1179)

Author: Ekdohibs

lambda/translcore.ml

@@ -361,11 +361,18 @@ and transl_exp0 ~in_new_scope ~scopes e =
         body_layout (event_before ~scopes body (transl_exp ~scopes body))
   | Texp_function { arg_label = _; param; cases; partial;
                     region; curry; warnings; arg_mode; alloc_mode } ->
+      (* CR ncourant: it would be better if we had [arg_layout] here *)
+      let arg_layout =
+        match is_function_type e.exp_env e.exp_type with
+        | None -> Misc.fatal_error "Translcore.transl_exp0: Type of a function is not a function type"
+        | Some (arg_type, _) ->
+            Typeopt.layout e.exp_env arg_type
+      in
       let scopes =
         if in_new_scope then scopes
         else enter_anonymous_function ~scopes
       in
-      transl_function ~scopes e alloc_mode param arg_mode cases partial warnings region curry
+      transl_function ~scopes e alloc_mode param arg_mode arg_layout cases partial warnings region curry
   | Texp_apply({ exp_desc = Texp_ident(path, _, {val_kind = Val_prim p},
                                        Id_prim pmode);
                 exp_type = prim_type; } as funct, oargs, pos, alloc_mode)
@@ -1072,10 +1079,10 @@ and transl_apply ~scopes
 
 and transl_curried_function
       ~scopes loc return
-      repr ~region ~curry partial warnings (param:Ident.t) cases =
+      repr ~region ~curry partial warnings (param:Ident.t) arg_layout cases =
   let max_arity = Lambda.max_arity () in
   let rec loop ~scopes loc return ~arity ~region ~curry
-            partial warnings (param:Ident.t) cases =
+            partial warnings (param:Ident.t) arg_layout cases =
     match curry, cases with
       More_args {partial_mode},
       [{c_lhs=pat; c_guard=None;
@@ -1091,12 +1098,18 @@ and transl_curried_function
       if Parmatch.inactive ~partial pat
       then
         let partial_mode = transl_alloc_mode partial_mode in
-        let layout = layout pat.pat_env pat.pat_type in
         let return_layout = function_return_layout exp_env exp_type in
+        let arg_layout' =
+          match is_function_type exp_env exp_type with
+          | None ->
+              Misc.fatal_error "Translcore.transl_curried_function: \
+                                Type of Texp_function is not function"
+          | Some (lhs, _) -> layout exp_env lhs
+        in
         let ((fnkind, params, return, region), body) =
           loop ~scopes exp_loc return_layout
             ~arity:(arity + 1) ~region:region' ~curry:curry'
-            partial' warnings' param' cases'
+            partial' warnings' param' arg_layout' cases'
         in
         let fnkind =
           match partial_mode, fnkind with
@@ -1109,8 +1122,8 @@ and transl_curried_function
              assert (nlocal = List.length params);
              Curried {nlocal = nlocal + 1}
         in
-        ((fnkind, (param, layout) :: params, return, region),
-         Matching.for_function ~scopes return_layout loc None (Lvar param, layout)
+        ((fnkind, (param, arg_layout) :: params, return, region),
+         Matching.for_function ~scopes return_layout loc None (Lvar param, arg_layout)
            [pat, body] partial)
       else begin
         begin match partial with
@@ -1123,18 +1136,18 @@ and transl_curried_function
         | Partial -> ()
         end;
         transl_tupled_function ~scopes ~arity ~region ~curry
-          loc return repr partial param cases
+          loc return repr partial param arg_layout cases
       end
     | curry, cases ->
       transl_tupled_function ~scopes ~arity ~region ~curry
-        loc return repr partial param cases
+        loc return repr partial param arg_layout cases
   in
   loop ~scopes loc return ~arity:1 ~region ~curry
-    partial warnings param cases
+    partial warnings param arg_layout cases
 
 and transl_tupled_function
       ~scopes ~arity ~region ~curry loc return
-      repr partial (param:Ident.t) cases =
+      repr partial (param:Ident.t) arg_layout cases =
   let partial_mode =
     match curry with
     | More_args {partial_mode} | Final_arg {partial_mode} ->
@@ -1153,21 +1166,13 @@ and transl_tupled_function
               (Matching.flatten_pattern size c_lhs, c_guard, c_rhs))
             cases in
         let kinds =
-          (* All the patterns might not share the same types. We must take the
-             union of the patterns types *)
-          match pats_expr_list with
-          | [] -> assert false
-          | (pats, _, _) :: cases ->
-              let first_case_layouts =
-                List.map (fun pat -> layout pat.pat_env pat.pat_type) pats
-              in
-              List.fold_left
-                (fun kinds (pats, _, _) ->
-                   List.map2 (fun kind pat ->
-                       layout_union kind
-                         (layout pat.pat_env pat.pat_type))
-                     kinds pats)
-                first_case_layouts cases
+          match arg_layout with
+          | Pvalue (Pvariant { consts = []; non_consts = [0, kinds] }) ->
+              List.map (fun vk -> Pvalue vk) kinds
+          | _ ->
+              Misc.fatal_error
+                "Translcore.transl_tupled_function: \
+                 Argument should be a tuple, but couldn't get the kinds"
         in
         let tparams =
           List.map (fun kind -> Ident.create_local "param", kind) kinds
@@ -1181,29 +1186,16 @@ and transl_tupled_function
         ((Tupled, tparams, return, region), body)
     with Matching.Cannot_flatten ->
       transl_function0 ~scopes loc ~region ~partial_mode
-        return repr partial param cases
+        return repr partial param arg_layout cases
       end
   | _ -> transl_function0 ~scopes loc ~region ~partial_mode
-           return repr partial param cases
+           return repr partial param arg_layout cases
 
 and transl_function0
       ~scopes loc ~region ~partial_mode return
-      repr partial (param:Ident.t) cases =
-    let layout =
-      match cases with
-      | [] ->
-        (* With Camlp4, a pattern matching might be empty *)
-        Lambda.layout_bottom
-      | {c_lhs=pat} :: other_cases ->
-        (* All the patterns might not share the same types. We must take the
-           union of the patterns types *)
-        List.fold_left (fun ly {c_lhs=pat} ->
-          Typeopt.layout_union ly
-            (layout pat.pat_env pat.pat_type))
-          (layout pat.pat_env pat.pat_type) other_cases
-    in
+      repr partial (param:Ident.t) arg_layout cases =
     let body =
-      Matching.for_function ~scopes return loc repr (Lvar param, layout)
+      Matching.for_function ~scopes return loc repr (Lvar param, arg_layout)
         (transl_cases ~scopes cases) partial
     in
     let region = region || not (may_allocate_in_region body) in
@@ -1213,17 +1205,17 @@ and transl_function0
         | Alloc_local -> 1
         | Alloc_heap -> 0
     in
-    ((Curried {nlocal}, [param, layout], return, region), body)
+    ((Curried {nlocal}, [param, arg_layout], return, region), body)
 
-and transl_function ~scopes e alloc_mode param arg_mode cases partial warnings region curry =
+and transl_function ~scopes e alloc_mode param arg_mode arg_layout cases partial warnings region curry =
   let mode = transl_alloc_mode alloc_mode in
   let ((kind, params, return, region), body) =
     event_function ~scopes e
       (function repr ->
          let pl = push_defaults e.exp_loc arg_mode cases partial warnings in
          let return_layout = function_return_layout e.exp_env e.exp_type in
          transl_curried_function ~scopes e.exp_loc return_layout
-           repr ~region ~curry partial warnings param pl)
+           repr ~region ~curry partial warnings param arg_layout pl)
   in
   let attr = default_function_attribute in
   let loc = of_location ~scopes e.exp_loc in
@@ -1570,13 +1562,31 @@ and transl_letop ~scopes loc env let_ ands param case partial warnings =
   in
   let exp = loop (layout let_.bop_exp.exp_env let_.bop_exp.exp_type) (transl_exp ~scopes let_.bop_exp) ands in
   let func =
+    let arg_layout =
+      match Typeopt.is_function_type env let_.bop_op_type with
+      | None ->
+          Misc.fatal_error
+            "Translcore.transl_letop: letop should be a function"
+      | Some (_, rhs) ->
+          match Typeopt.is_function_type env rhs with
+          | None ->
+              Misc.fatal_error
+                "Translcore.transl_letop: letop should have at least two arguments"
+          | Some (lhs, _) ->
+              match Typeopt.is_function_type env lhs with
+              | None ->
+                  Misc.fatal_error
+                    "Translcore.transl_letop: letop second argument should be a function"
+              | Some (arg_type, _) ->
+                  Typeopt.layout env arg_type
+    in
     let return_layout = layout case.c_rhs.exp_env case.c_rhs.exp_type in
     let curry = More_args { partial_mode = Alloc_mode.global } in
     let (kind, params, return, _region), body =
       event_function ~scopes case.c_rhs
         (function repr ->
            transl_curried_function ~scopes case.c_rhs.exp_loc return_layout
-             repr ~region:true ~curry partial warnings param [case])
+             repr ~region:true ~curry partial warnings param arg_layout [case])
     in
     let attr = default_function_attribute in
     let loc = of_location ~scopes case.c_rhs.exp_loc in