Add result layout in Lapply and Lsend (#1102)

Author: Ekdohibs

middle_end/closure/closure.ml

@@ -1099,6 +1099,7 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
                  ap_loc=loc;
                  ap_func=(Lvar funct_var);
                  ap_args=internal_args;
+                 ap_result_layout=Lambda.layout_top;
                  ap_region_close=Rc_normal;
                  ap_mode=ret_mode;
                  ap_tailcall=Default_tailcall;
@@ -1165,7 +1166,7 @@ let rec close ({ backend; fenv; cenv ; mutable_vars; kinds; catch_env } as env)
           fail_if_probe ~probe "Unknown function";
           (Ugeneric_apply(ufunct, uargs, (pos, mode), dbg), Value_unknown)
       end
-  | Lsend(kind, met, obj, args, pos, mode, loc) ->
+  | Lsend(kind, met, obj, args, pos, mode, loc, _result_layout) ->
       let (umet, _) = close env met in
       let (uobj, _) = close env obj in
       let dbg = Debuginfo.from_location loc in

middle_end/flambda/closure_conversion.ml

@@ -329,7 +329,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
       in
       Let_rec (defs, close t env body)
     end
-  | Lsend (kind, meth, obj, args, reg_close, mode, loc) ->
+  | Lsend (kind, meth, obj, args, reg_close, mode, loc, _layout) ->
     let meth_var = Variable.create Names.meth in
     let obj_var = Variable.create Names.obj in
     let dbg = Debuginfo.from_location loc in

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -1053,6 +1053,7 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
   | Lapply
       { ap_func;
         ap_args;
+        ap_result_layout;
         ap_region_close;
         ap_mode;
         ap_loc;
@@ -1063,7 +1064,7 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
       } ->
     (* Note that we don't need kind information about [ap_args] since we already
        have it on the corresponding [Simple]s in the environment. *)
-    maybe_insert_let_cont "apply_result" Lambda.layout_top k acc env ccenv
+    maybe_insert_let_cont "apply_result" ap_result_layout k acc env ccenv
       (fun acc env ccenv k ->
         cps_tail_apply acc env ccenv ap_func ap_args ap_region_close ap_mode
           ap_loc ap_inlined ap_probe k k_exn)
@@ -1254,16 +1255,16 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
         let body acc ccenv = cps_tail acc body_env ccenv body k k_exn in
         CC.close_let_cont acc ccenv ~name:continuation ~is_exn_handler:false
           ~params ~recursive ~body ~handler)
-  | Lsend (meth_kind, meth, obj, args, pos, mode, loc) ->
+  | Lsend (meth_kind, meth, obj, args, pos, mode, loc, layout) ->
     cps_non_tail_simple acc env ccenv obj
       (fun acc env ccenv obj ->
         cps_non_tail_var "meth" acc env ccenv meth
           Flambda_kind.With_subkind.any_value
           (fun acc env ccenv meth ->
             cps_non_tail_list acc env ccenv args
               (fun acc env ccenv args ->
-                maybe_insert_let_cont "send_result" Lambda.layout_top k acc env
-                  ccenv (fun acc env ccenv k ->
+                maybe_insert_let_cont "send_result" layout k acc env ccenv
+                  (fun acc env ccenv k ->
                     let exn_continuation : IR.exn_continuation =
                       { exn_handler = k_exn;
                         extra_args = extra_args_for_exn_continuation env k_exn

ocaml/bytecomp/bytegen.ml

@@ -591,7 +591,7 @@ let rec comp_expr env exp sz cont =
                       (Kapply nargs :: cont1))
         end
       end
-  | Lsend(kind, met, obj, args, rc, _, _) ->
+  | Lsend(kind, met, obj, args, rc, _, _, _) ->
       assert (kind <> Cached);
       let nargs = List.length args + 1 in
       let getmethod, args' =
@@ -994,7 +994,7 @@ let rec comp_expr env exp sz cont =
             match lam with
             | Lprim(prim, _, _) -> preserve_tailcall_for_prim prim
             | Lapply {ap_region_close=rc; _}
-            | Lsend(_, _, _, _, rc, _, _) ->
+            | Lsend(_, _, _, _, rc, _, _, _) ->
                not (is_nontail rc)
             | _ -> true
           in
@@ -1005,7 +1005,7 @@ let rec comp_expr env exp sz cont =
             let info =
               match lam with
                 Lapply{ap_args = args}  -> Event_return (List.length args)
-              | Lsend(_, _, _, args, _, _, _) ->
+              | Lsend(_, _, _, args, _, _, _, _) ->
                   Event_return (List.length args + 1)
               | Lprim(_,args,_)         -> Event_return (List.length args)
               | _                       -> Event_other

ocaml/lambda/lambda.ml

@@ -442,7 +442,7 @@ type lambda =
   | Lassign of Ident.t * lambda
   | Lsend of
       meth_kind * lambda * lambda * lambda list
-      * region_close * alloc_mode * scoped_location
+      * region_close * alloc_mode * scoped_location * layout
   | Levent of lambda * lambda_event
   | Lifused of Ident.t * lambda
   | Lregion of lambda
@@ -476,6 +476,7 @@ and lambda_for =
 and lambda_apply =
   { ap_func : lambda;
     ap_args : lambda list;
+    ap_result_layout : layout;
     ap_region_close : region_close;
     ap_mode : alloc_mode;
     ap_loc : scoped_location;
@@ -565,6 +566,7 @@ let layout_string = Pvalue Pgenval
 let layout_boxedint bi = Pvalue (Pboxedintval bi)
 let layout_lazy = Pvalue Pgenval
 let layout_lazy_contents = Pvalue Pgenval
+let layout_any_value = Pvalue Pgenval
 
 let layout_top = Pvalue Pgenval
 
@@ -650,8 +652,8 @@ let make_key e =
         Lsequence (tr_rec env e1,tr_rec env e2)
     | Lassign (x,e) ->
         Lassign (x,tr_rec env e)
-    | Lsend (m,e1,e2,es,pos,mo,_loc) ->
-        Lsend (m,tr_rec env e1,tr_rec env e2,tr_recs env es,pos,mo,Loc_unknown)
+    | Lsend (m,e1,e2,es,pos,mo,_loc,layout) ->
+        Lsend (m,tr_rec env e1,tr_rec env e2,tr_recs env es,pos,mo,Loc_unknown,layout)
     | Lifused (id,e) -> Lifused (id,tr_rec env e)
     | Lregion e -> Lregion (tr_rec env e)
     | Lletrec _|Lfunction _
@@ -746,7 +748,7 @@ let shallow_iter ~tail ~non_tail:f = function
       f for_from; f for_to; f for_body
   | Lassign(_, e) ->
       f e
-  | Lsend (_k, met, obj, args, _, _, _) ->
+  | Lsend (_k, met, obj, args, _, _, _, _) ->
       List.iter f (met::obj::args)
   | Levent (e, _evt) ->
       tail e
@@ -825,7 +827,7 @@ let rec free_variables = function
            (Ident.Set.remove for_id (free_variables for_body)))
   | Lassign(id, e) ->
       Ident.Set.add id (free_variables e)
-  | Lsend (_k, met, obj, args, _, _, _) ->
+  | Lsend (_k, met, obj, args, _, _, _, _) ->
       free_variables_list
         (Ident.Set.union (free_variables met) (free_variables obj))
         args
@@ -1007,9 +1009,9 @@ let subst update_env ?(freshen_bound_variables = false) s input_lam =
         assert (not (Ident.Map.mem id s));
         let id = try Ident.Map.find id l with Not_found -> id in
         Lassign(id, subst s l e)
-    | Lsend (k, met, obj, args, pos, mode, loc) ->
+    | Lsend (k, met, obj, args, pos, mode, loc, layout) ->
         Lsend (k, subst s l met, subst s l obj, subst_list s l args,
-               pos, mode, loc)
+               pos, mode, loc, layout)
     | Levent (lam, evt) ->
         let old_env = evt.lev_env in
         let env_updates =
@@ -1070,11 +1072,12 @@ let shallow_map ~tail ~non_tail:f = function
   | Lvar _
   | Lmutvar _
   | Lconst _ as lam -> lam
-  | Lapply { ap_func; ap_args; ap_region_close; ap_mode; ap_loc; ap_tailcall;
+  | Lapply { ap_func; ap_args; ap_result_layout; ap_region_close; ap_mode; ap_loc; ap_tailcall;
              ap_inlined; ap_specialised; ap_probe } ->
       Lapply {
         ap_func = f ap_func;
         ap_args = List.map f ap_args;
+        ap_result_layout;
         ap_region_close;
         ap_mode;
         ap_loc;
@@ -1131,8 +1134,8 @@ let shallow_map ~tail ~non_tail:f = function
                      for_body = f lf.for_body }
   | Lassign (v, e) ->
       Lassign (v, f e)
-  | Lsend (k, m, o, el, pos, mode, loc) ->
-      Lsend (k, f m, f o, List.map f el, pos, mode, loc)
+  | Lsend (k, m, o, el, pos, mode, loc, layout) ->
+      Lsend (k, f m, f o, List.map f el, pos, mode, loc, layout)
   | Levent (l, ev) ->
       Levent (tail l, ev)
   | Lifused (v, e) ->
@@ -1329,3 +1332,5 @@ let structured_constant_layout = function
   | Const_base const -> constant_layout const
   | Const_block _ | Const_immstring _ -> Pvalue Pgenval
   | Const_float_array _ | Const_float_block _ -> Pvalue (Parrayval Pfloatarray)
+
+let primitive_result_layout (_p : primitive) = layout_top

ocaml/lambda/lambda.mli

@@ -370,7 +370,7 @@ type lambda =
   | Lfor of lambda_for
   | Lassign of Ident.t * lambda
   | Lsend of meth_kind * lambda * lambda * lambda list
-             * region_close * alloc_mode * scoped_location
+             * region_close * alloc_mode * scoped_location * layout
   | Levent of lambda * lambda_event
   | Lifused of Ident.t * lambda
   | Lregion of lambda
@@ -409,6 +409,7 @@ and lambda_for =
 and lambda_apply =
   { ap_func : lambda;
     ap_args : lambda list;
+    ap_result_layout : layout;
     ap_region_close : region_close;
     ap_mode : alloc_mode;
     ap_loc : scoped_location;
@@ -482,6 +483,8 @@ val layout_boxedint : boxed_integer -> layout
 val layout_field : layout
 val layout_lazy : layout
 val layout_lazy_contents : layout
+(* A layout that is Pgenval because we are missing layout polymorphism *)
+val layout_any_value : layout
 
 val layout_top : layout
 
@@ -626,3 +629,5 @@ val mod_field: ?read_semantics: field_read_semantics -> int -> primitive
 val mod_setfield: int -> primitive
 
 val structured_constant_layout : structured_constant -> layout
+
+val primitive_result_layout : primitive -> layout

ocaml/lambda/matching.ml

@@ -1929,6 +1929,7 @@ let inline_lazy_force_cond arg pos loc =
                       ap_loc = loc;
                       ap_func = force_fun;
                       ap_args = [ varg ];
+                      ap_result_layout = Lambda.layout_lazy_contents;
                       ap_region_close = pos;
                       ap_mode = alloc_heap;
                       ap_inlined = Default_inlined;
@@ -1965,6 +1966,7 @@ let inline_lazy_force_switch arg pos loc =
                           ap_loc = loc;
                           ap_func = force_fun;
                           ap_args = [ varg ];
+                          ap_result_layout = Lambda.layout_lazy_contents;
                           ap_region_close = pos;
                           ap_mode = alloc_heap;
                           ap_inlined = Default_inlined;
@@ -1987,6 +1989,7 @@ let inline_lazy_force arg pos loc =
         ap_loc = loc;
         ap_func = Lazy.force code_force_lazy;
         ap_args = [ arg ];
+        ap_result_layout = Lambda.layout_lazy_contents;
         ap_region_close = pos;
         ap_mode = alloc_heap;
         ap_inlined = Default_inlined;

ocaml/lambda/printlambda.ml

@@ -771,7 +771,7 @@ let rec lam ppf = function
        lam for_to (alloc_mode mode) lam for_body
   | Lassign(id, expr) ->
       fprintf ppf "@[<2>(assign@ %a@ %a)@]" Ident.print id lam expr
-  | Lsend (k, met, obj, largs, pos, reg, _) ->
+  | Lsend (k, met, obj, largs, pos, reg, _, _) ->
       let args ppf largs =
         List.iter (fun l -> fprintf ppf "@ %a" lam l) largs in
       let kind =

ocaml/lambda/simplif.ml

@@ -88,9 +88,9 @@ let rec eliminate_ref id = function
                     for_body = eliminate_ref id lf.for_body }
   | Lassign(v, e) ->
       Lassign(v, eliminate_ref id e)
-  | Lsend(k, m, o, el, pos, mode, loc) ->
+  | Lsend(k, m, o, el, pos, mode, loc, layout) ->
       Lsend(k, eliminate_ref id m, eliminate_ref id o,
-            List.map (eliminate_ref id) el, pos, mode, loc)
+            List.map (eliminate_ref id) el, pos, mode, loc, layout)
   | Levent(l, ev) ->
       Levent(eliminate_ref id l, ev)
   | Lifused(v, e) ->
@@ -181,7 +181,7 @@ let simplify_exits lam =
       count ~try_depth lf.for_to;
       count ~try_depth lf.for_body
   | Lassign(_v, l) -> count ~try_depth l
-  | Lsend(_k, m, o, ll, _, _, _) -> List.iter (count ~try_depth) (m::o::ll)
+  | Lsend(_k, m, o, ll, _, _, _, _) -> List.iter (count ~try_depth) (m::o::ll)
   | Levent(l, _) -> count ~try_depth l
   | Lifused(_v, l) -> count ~try_depth l
   | Lregion l -> count ~try_depth l
@@ -324,9 +324,9 @@ let simplify_exits lam =
                     for_to = simplif ~try_depth lf.for_to;
                     for_body = simplif ~try_depth lf.for_body}
   | Lassign(v, l) -> Lassign(v, simplif ~try_depth l)
-  | Lsend(k, m, o, ll, pos, mode, loc) ->
+  | Lsend(k, m, o, ll, pos, mode, loc, layout) ->
       Lsend(k, simplif ~try_depth m, simplif ~try_depth o,
-      List.map (simplif ~try_depth) ll, pos, mode, loc)
+      List.map (simplif ~try_depth) ll, pos, mode, loc, layout)
   | Levent(l, ev) -> Levent(simplif ~try_depth l, ev)
   | Lifused(v, l) -> Lifused (v,simplif ~try_depth l)
   | Lregion l -> Lregion (simplif ~try_depth l)
@@ -458,7 +458,7 @@ let simplify_lets lam =
       (* Lalias-bound variables are never assigned, so don't increase
          v's refcount *)
       count bv l
-  | Lsend(_, m, o, ll, _, _, _) -> List.iter (count bv) (m::o::ll)
+  | Lsend(_, m, o, ll, _, _, _, _) -> List.iter (count bv) (m::o::ll)
   | Levent(l, _) -> count bv l
   | Lifused(v, l) ->
       if count_var v > 0 then count bv l
@@ -608,8 +608,8 @@ let simplify_lets lam =
                              for_to = simplif lf.for_to;
                              for_body = simplif lf.for_body}
   | Lassign(v, l) -> Lassign(v, simplif l)
-  | Lsend(k, m, o, ll, pos, mode, loc) ->
-      Lsend(k, simplif m, simplif o, List.map simplif ll, pos, mode, loc)
+  | Lsend(k, m, o, ll, pos, mode, loc, layout) ->
+      Lsend(k, simplif m, simplif o, List.map simplif ll, pos, mode, loc, layout)
   | Levent(l, ev) -> Levent(simplif l, ev)
   | Lifused(v, l) ->
       if count_var v > 0 then simplif l else lambda_unit
@@ -696,7 +696,7 @@ let rec emit_tail_infos is_tail lambda =
       emit_tail_infos false for_body
   | Lassign (_, lam) ->
       emit_tail_infos false lam
-  | Lsend (_, meth, obj, args, _, _, _loc) ->
+  | Lsend (_, meth, obj, args, _, _, _loc, _) ->
       emit_tail_infos false meth;
       emit_tail_infos false obj;
       list_emit_tail_infos false args
@@ -760,6 +760,7 @@ let split_default_wrapper ~id:fun_id ~kind ~params ~return ~body
           Lapply {
             ap_func = Lvar inner_id;
             ap_args = args;
+            ap_result_layout = return;
             ap_loc = Loc_unknown;
             ap_region_close = Rc_normal;
             ap_mode = alloc_heap;