flambda-backend: Propagates layouts through Flambda1 (#1115)

Author: chambart

lambda/lambda.ml

@@ -291,6 +291,8 @@ let rec equal_value_kind x y =
 
 let equal_layout (Pvalue x) (Pvalue y) = equal_value_kind x y
 
+let compatible_layout (Pvalue _) (Pvalue _) = true
+
 let must_be_value layout =
   match layout with
   | Pvalue v -> v
@@ -555,6 +557,7 @@ let layout_block = Pvalue Pgenval
 let layout_list =
   Pvalue (Pvariant { consts = [0] ; non_consts = [0, [Pgenval; Pgenval]] })
 let layout_field = Pvalue Pgenval
+let layout_exception = Pvalue Pgenval
 let layout_function = Pvalue Pgenval
 let layout_object = Pvalue Pgenval
 let layout_class = Pvalue Pgenval
@@ -567,8 +570,12 @@ 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_letrec = layout_any_value
 
 let layout_top = Pvalue Pgenval
+let layout_bottom =
+  (* CR pchambart: this should be an actual bottom *)
+  Pvalue Pgenval
 
 let default_function_attribute = {
   inline = Default_inline;

lambda/lambda.mli

@@ -237,6 +237,8 @@ val equal_value_kind : value_kind -> value_kind -> bool
 
 val equal_layout : layout -> layout -> bool
 
+val compatible_layout : layout -> layout -> bool
+
 val equal_boxed_integer : boxed_integer -> boxed_integer -> bool
 
 val must_be_value : layout -> value_kind
@@ -470,6 +472,7 @@ val layout_int : layout
 val layout_array : array_kind -> layout
 val layout_block : layout
 val layout_list : layout
+val layout_exception : layout
 val layout_function : layout
 val layout_object : layout
 val layout_class : layout
@@ -485,8 +488,10 @@ 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_letrec : layout
 
 val layout_top : layout
+val layout_bottom : layout
 
 val name_lambda: let_kind -> lambda -> layout -> (Ident.t -> lambda) -> lambda
 val name_lambda_list: (lambda * layout) list -> (lambda list -> lambda) -> lambda

lambda/translclass.ml

@@ -214,7 +214,7 @@ let rec build_object_init ~scopes cl_table obj params inh_init obj_init cl =
       let (inh_init, obj_init) =
         build_object_init ~scopes cl_table obj params inh_init obj_init cl
       in
-      (inh_init, transl_apply ~scopes obj_init oexprs Loc_unknown)
+      (inh_init, transl_apply ~result_layout:Lambda.layout_top ~scopes obj_init oexprs Loc_unknown)
   | Tcl_let (rec_flag, defs, vals, cl) ->
       let (inh_init, obj_init) =
         build_object_init ~scopes cl_table obj (vals @ params)
@@ -485,7 +485,7 @@ let rec transl_class_rebind ~scopes obj_init cl vf =
   | Tcl_apply (cl, oexprs) ->
       let path, path_lam, obj_init =
         transl_class_rebind ~scopes obj_init cl vf in
-      (path, path_lam, transl_apply ~scopes obj_init oexprs Loc_unknown)
+      (path, path_lam, transl_apply ~result_layout:Lambda.layout_top ~scopes obj_init oexprs Loc_unknown)
   | Tcl_let (rec_flag, defs, _vals, cl) ->
       let path, path_lam, obj_init =
         transl_class_rebind ~scopes obj_init cl vf in

lambda/translcore.ml

@@ -402,19 +402,21 @@ and transl_exp0 ~in_new_scope ~scopes e =
         let e = { e with exp_desc = Texp_apply(funct, oargs, pos) } in
         let position = transl_apply_position pos in
         let mode = transl_exp_mode e in
+        let result_layout = Typeopt.layout e.exp_env e.exp_type in
         event_after ~scopes e
           (transl_apply ~scopes ~tailcall ~inlined ~specialised ~position ~mode
-             lam extra_args (of_location ~scopes e.exp_loc))
+             ~result_layout lam extra_args (of_location ~scopes e.exp_loc))
       end
   | Texp_apply(funct, oargs, position) ->
       let tailcall = Translattribute.get_tailcall_attribute funct in
       let inlined = Translattribute.get_inlined_attribute funct in
       let specialised = Translattribute.get_specialised_attribute funct in
+      let result_layout = Typeopt.layout e.exp_env e.exp_type in
       let e = { e with exp_desc = Texp_apply(funct, oargs, position) } in
       let position = transl_apply_position position in
       let mode = transl_exp_mode e in
       event_after ~scopes e
-        (transl_apply ~scopes ~tailcall ~inlined ~specialised
+        (transl_apply ~scopes ~tailcall ~inlined ~specialised ~result_layout
            ~position ~mode (transl_exp ~scopes funct)
            oargs (of_location ~scopes e.exp_loc))
   | Texp_match(arg, pat_expr_list, partial) ->
@@ -957,10 +959,10 @@ and transl_apply ~scopes
       ?(specialised = Default_specialise)
       ?(position=Rc_normal)
       ?(mode=alloc_heap)
+      ~result_layout
       lam sargs loc
   =
   let lapply funct args loc pos mode =
-    let result_layout = Lambda.layout_top in
     match funct, pos with
     | Lsend((Self | Public) as k, lmet, lobj, [], _, _, _, _), _ ->
         Lsend(k, lmet, lobj, args, pos, mode, loc, result_layout)
@@ -1584,9 +1586,9 @@ let transl_scoped_exp ~scopes exp =
   maybe_region_exp exp (transl_scoped_exp ~scopes exp)
 
 let transl_apply
-      ~scopes ?tailcall ?inlined ?specialised ?position ?mode fn args loc =
-  maybe_region_layout Lambda.layout_top (transl_apply
-      ~scopes ?tailcall ?inlined ?specialised ?position ?mode fn args loc)
+      ~scopes ?tailcall ?inlined ?specialised ?position ?mode ~result_layout fn args loc =
+  maybe_region_layout result_layout (transl_apply
+      ~scopes ?tailcall ?inlined ?specialised ?position ?mode ~result_layout fn args loc)
 
 (* Error report *)
 

lambda/translcore.mli

@@ -31,6 +31,7 @@ val transl_apply: scopes:scopes
                   -> ?specialised:specialise_attribute
                   -> ?position:region_close
                   -> ?mode:alloc_mode
+                  -> result_layout:Lambda.layout
                   -> lambda
                   -> (arg_label * apply_arg) list
                   -> scoped_location -> lambda

middle_end/clambda_primitives.ml

@@ -168,3 +168,5 @@ and raise_kind = Lambda.raise_kind =
   | Raise_notrace
 
 let equal (x: primitive) (y: primitive) = x = y
+
+let result_layout _p = Lambda.layout_any_value

middle_end/clambda_primitives.mli

@@ -171,3 +171,5 @@ and raise_kind = Lambda.raise_kind =
   | Raise_notrace
 
 val equal : primitive -> primitive -> bool
+
+val result_layout : primitive -> Lambda.layout

middle_end/flambda/augment_specialised_args.ml

@@ -470,6 +470,7 @@ module Make (T : S) = struct
           args =
             (Parameter.List.vars wrapper_params) @
             spec_args_bound_in_the_wrapper;
+          result_layout = function_decl.return_layout;
           kind = Direct (Closure_id.wrap new_fun_var);
           dbg = Debuginfo.none;
           reg_close = Rc_normal;
@@ -544,6 +545,7 @@ module Make (T : S) = struct
     let new_function_decl =
       Flambda.create_function_declaration
         ~params:wrapper_params
+        ~return_layout:function_decl.return_layout
         ~alloc_mode
         ~region:function_decl.region
         ~body:wrapper_body
@@ -651,6 +653,7 @@ module Make (T : S) = struct
       let rewritten_function_decl =
         Flambda.create_function_declaration
           ~params:all_params
+          ~return_layout:function_decl.return_layout
           ~alloc_mode:function_decl.alloc_mode
           ~region:function_decl.region
           ~body:function_decl.body

middle_end/flambda/closure_conversion.ml

@@ -40,10 +40,10 @@ let add_default_argument_wrappers lam =
     match lam with
     | Llet (( Strict | Alias | StrictOpt), _k, id,
         Lfunction {kind; params; body = fbody; attr; loc;
-                   mode; region}, body) ->
+                   mode; region; return }, body) ->
       begin match
         Simplif.split_default_wrapper ~id ~kind ~params
-          ~body:fbody ~return:Lambda.layout_top ~attr ~loc ~mode ~region
+          ~body:fbody ~return ~attr ~loc ~mode ~region
       with
       | [fun_id, def] -> Llet (Alias, Lambda.layout_function, fun_id, def, body)
       | [fun_id, def; inner_fun_id, def_inner] ->
@@ -58,9 +58,9 @@ let add_default_argument_wrappers lam =
             (List.map
                (function
                  | (id, Lambda.Lfunction {kind; params; body; attr; loc;
-                                          mode; region}) ->
+                                          mode; region; return }) ->
                    Simplif.split_default_wrapper ~id ~kind ~params ~body
-                     ~return:Lambda.layout_top ~attr ~loc ~mode ~region
+                     ~return ~attr ~loc ~mode ~region
                  | _ -> assert false)
                defs)
         in
@@ -73,14 +73,15 @@ let add_default_argument_wrappers lam =
 (** Generate a wrapper ("stub") function that accepts a tuple argument and
     calls another function with arguments extracted in the obvious
     manner from the tuple. *)
-let tupled_function_call_stub original_params unboxed_version ~closure_bound_var ~region
+let tupled_function_call_stub original_params unboxed_version ~closure_bound_var ~region ~return_layout
       : Flambda.function_declaration =
   let tuple_param_var = Variable.rename unboxed_version in
   let params = List.map (fun p -> Variable.rename p) original_params in
   let call : Flambda.t =
     Apply ({
         func = unboxed_version;
         args = params;
+        result_layout = return_layout;
         (* CR-someday mshinwell for mshinwell: investigate if there is some
            redundancy here (func is also unboxed_version) *)
         kind = Direct (Closure_id.wrap unboxed_version);
@@ -104,7 +105,7 @@ let tupled_function_call_stub original_params unboxed_version ~closure_bound_var
   let alloc_mode = Lambda.alloc_heap in
   let tuple_param = Parameter.wrap tuple_param_var alloc_mode Lambda.layout_block in
   Flambda.create_function_declaration ~params:[tuple_param] ~alloc_mode ~region
-    ~body ~stub:true ~dbg:Debuginfo.none ~inline:Default_inline
+    ~body ~stub:true ~dbg:Debuginfo.none ~inline:Default_inline ~return_layout
     ~specialise:Default_specialise ~is_a_functor:false
     ~closure_origin:(Closure_origin.create (Closure_id.wrap closure_bound_var))
     ~poll:Default_poll (* don't propogate attribute to wrappers *)
@@ -215,7 +216,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
            initial_value = var;
            body;
            contents_kind = block_kind })
-  | Lfunction { kind; params; body; attr; loc; mode; region } ->
+  | Lfunction { kind; params; body; attr; loc; mode; region; return } ->
     let name = Names.anon_fn_with_loc loc in
     let closure_bound_var = Variable.create name in
     (* CR-soon mshinwell: some of this is now very similar to the let rec case
@@ -224,7 +225,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
     let set_of_closures =
       let decl =
         Function_decl.create ~let_rec_ident:None ~closure_bound_var ~kind ~mode
-          ~region ~params ~body ~attr ~loc
+          ~region ~params ~body ~attr ~loc ~return_layout:return
       in
       close_functions t env (Function_decls.create [decl])
     in
@@ -235,7 +236,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
     in
     Flambda.create_let set_of_closures_var set_of_closures
       (name_expr (Project_closure (project_closure)) ~name)
-  | Lapply { ap_func; ap_args; ap_loc; ap_region_close; ap_mode;
+  | Lapply { ap_func; ap_args; ap_loc; ap_region_close; ap_mode; ap_result_layout;
              ap_tailcall = _; ap_inlined; ap_specialised; ap_probe; } ->
     Lift_code.lifting_helper (close_list t env ap_args)
       ~evaluation_order:`Right_to_left
@@ -247,6 +248,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
           (Apply ({
               func = func_var;
               args;
+              result_layout = ap_result_layout;
               kind = Indirect;
               dbg = Debuginfo.from_location ap_loc;
               reg_close = ap_region_close;
@@ -259,22 +261,23 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
   | Lletrec (defs, body) ->
     let env =
       List.fold_right (fun (id,  _) env ->
-          Env.add_var env id (Variable.create_with_same_name_as_ident id) Lambda.layout_top)
+          Env.add_var env id (Variable.create_with_same_name_as_ident id)
+            Lambda.layout_letrec)
         defs env
     in
     let function_declarations =
       (* Identify any bindings in the [let rec] that are functions.  These
          will be named after the corresponding identifier in the [let rec]. *)
       List.map (function
           | (let_rec_ident,
-             Lambda.Lfunction { kind; params; body; attr; loc; mode; region }) ->
+             Lambda.Lfunction { kind; params; return; body; attr; loc; mode; region }) ->
             let closure_bound_var =
               Variable.create_with_same_name_as_ident let_rec_ident
             in
             let function_declaration =
               Function_decl.create ~let_rec_ident:(Some let_rec_ident)
                 ~closure_bound_var ~kind ~mode ~region
-                ~params ~body ~attr ~loc
+                ~params ~body ~attr ~loc ~return_layout:return
             in
             Some function_declaration
           | _ -> None)
@@ -324,7 +327,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, _layout) ->
+  | Lsend (kind, meth, obj, args, reg_close, mode, loc, result_layout) ->
     let meth_var = Variable.create Names.meth in
     let obj_var = Variable.create Names.obj in
     let dbg = Debuginfo.from_location loc in
@@ -335,7 +338,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
           ~name:Names.send_arg
           ~create_body:(fun args ->
               Send { kind; meth = meth_var; obj = obj_var; args;
-                     dbg; reg_close; mode })))
+                     dbg; reg_close; mode; result_layout })))
   | Lprim ((Pdivint Safe | Pmodint Safe
            | Pdivbint { is_safe = Safe } | Pmodbint { is_safe = Safe }) as prim,
            [arg1; arg2], loc)
@@ -524,7 +527,7 @@ let rec close t env (lam : Lambda.lambda) : Flambda.t =
       List.map (fun (ident, kind) ->
           (Variable.create_with_same_name_as_ident ident, kind)) ids
     in
-    Static_catch (st_exn, List.map fst vars, close t env body,
+    Static_catch (st_exn, vars, close t env body,
       close t (Env.add_vars env (List.map fst ids) vars) handler, kind)
   | Ltrywith (body, id, handler, kind) ->
     let var = Variable.create_with_same_name_as_ident id in
@@ -589,6 +592,7 @@ and close_functions t external_env function_declarations : Flambda.named =
     let dbg = Debuginfo.from_location loc in
     let region = Function_decl.region decl in
     let params = Function_decl.params decl in
+    let return_layout = Function_decl.return_layout decl in
     (* Create fresh variables for the elements of the closure (cf.
        the comment on [Function_decl.closure_env_without_parameters], above).
        This induces a renaming on [Function_decl.free_idents]; the results of
@@ -625,7 +629,7 @@ and close_functions t external_env function_declarations : Flambda.named =
     let fun_decl =
       Flambda.create_function_declaration
         ~params ~alloc_mode:(Function_decl.mode decl) ~region
-        ~body ~stub ~dbg
+        ~body ~stub ~dbg ~return_layout
         ~inline:(Function_decl.inline decl)
         ~specialise:(Function_decl.specialise decl)
         ~is_a_functor:(Function_decl.is_a_functor decl)
@@ -639,7 +643,7 @@ and close_functions t external_env function_declarations : Flambda.named =
       let unboxed_version = Variable.rename closure_bound_var in
       let generic_function_stub =
         tupled_function_call_stub (List.map fst param_vars) unboxed_version
-          ~closure_bound_var ~region
+          ~closure_bound_var ~region ~return_layout
       in
       Variable.Map.add unboxed_version fun_decl
         (Variable.Map.add closure_bound_var generic_function_stub map)
@@ -679,13 +683,13 @@ and close_list t sb l = List.map (close t sb) l
 and close_let_bound_expression t ?let_rec_ident let_bound_var env
       (lam : Lambda.lambda) : Flambda.named =
   match lam with
-  | Lfunction { kind; params; body; attr; loc; mode; region } ->
+  | Lfunction { kind; params; return; body; attr; loc; mode; region } ->
     (* Ensure that [let] and [let rec]-bound functions have appropriate
        names. *)
     let closure_bound_var = Variable.rename let_bound_var in
     let decl =
       Function_decl.create ~let_rec_ident ~closure_bound_var ~kind ~mode ~region
-        ~params ~body ~attr ~loc
+        ~params ~body ~attr ~loc ~return_layout:return
     in
     let set_of_closures_var = Variable.rename let_bound_var in
     let set_of_closures =

middle_end/flambda/closure_conversion_aux.ml

@@ -93,14 +93,15 @@ module Function_decls = struct
       mode : Lambda.alloc_mode;
       region : bool;
       params : (Ident.t * Lambda.layout) list;
+      return_layout : Lambda.layout;
       body : Lambda.lambda;
       free_idents_of_body : Ident.Set.t;
       attr : Lambda.function_attribute;
       loc : Lambda.scoped_location
     }
 
     let create ~let_rec_ident ~closure_bound_var ~kind ~mode ~region
-          ~params ~body ~attr ~loc =
+          ~params ~return_layout ~body ~attr ~loc =
       let let_rec_ident =
         match let_rec_ident with
         | None -> Ident.create_local "unnamed_function"
@@ -112,6 +113,7 @@ module Function_decls = struct
         mode;
         region;
         params;
+        return_layout;
         body;
         free_idents_of_body = Lambda.free_variables body;
         attr;
@@ -124,6 +126,7 @@ module Function_decls = struct
     let mode t = t.mode
     let region t = t.region
     let params t = t.params
+    let return_layout t = t.return_layout
     let body t = t.body
     let free_idents t = t.free_idents_of_body
     let inline t = t.attr.inline

middle_end/flambda/closure_conversion_aux.mli

@@ -61,6 +61,7 @@ module Function_decls : sig
       -> mode:Lambda.alloc_mode
       -> region:bool
       -> params:(Ident.t * Lambda.layout) list
+      -> return_layout:Lambda.layout
       -> body:Lambda.lambda
       -> attr:Lambda.function_attribute
       -> loc:Lambda.scoped_location
@@ -72,6 +73,7 @@ module Function_decls : sig
     val mode : t -> Lambda.alloc_mode
     val region : t -> bool
     val params : t -> (Ident.t * Lambda.layout) list
+    val return_layout : t -> Lambda.layout
     val body : t -> Lambda.lambda
     val inline : t -> Lambda.inline_attribute
     val specialise : t -> Lambda.specialise_attribute