Propagates layouts through Flambda1

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

@@ -41,10 +41,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] ->
@@ -59,9 +59,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
@@ -74,14 +74,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);
@@ -105,7 +106,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 ~inline:Default_inline
+    ~body ~stub:true ~inline:Default_inline ~return_layout
     ~specialise:Default_specialise ~check:Default_check ~is_a_functor:false
     ~closure_origin:(Closure_origin.create (Closure_id.wrap closure_bound_var))
     ~poll:Default_poll (* don't propogate attribute to wrappers *)
@@ -216,7 +217,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 =
       let debug_info = Debuginfo.from_location loc in
@@ -228,7 +229,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
@@ -239,7 +240,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
@@ -251,6 +252,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;
@@ -263,23 +265,24 @@ 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 =
               let debug_info = Debuginfo.from_location loc in
               Variable.create_with_same_name_as_ident ~debug_info 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)
@@ -329,7 +332,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
@@ -340,7 +343,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)
@@ -529,7 +532,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
@@ -592,6 +595,7 @@ and close_functions t external_env function_declarations : Flambda.named =
     let body = Function_decl.body decl 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
@@ -629,7 +633,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
+        ~body ~stub ~return_layout
         ~inline:(Function_decl.inline decl)
         ~specialise:(Function_decl.specialise decl)
         ~check:(Function_decl.check decl)
@@ -644,7 +648,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)
@@ -684,7 +688,7 @@ 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 =
@@ -693,7 +697,7 @@ and close_let_bound_expression t ?let_rec_ident let_bound_var env
     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

middle_end/flambda/flambda.ml

@@ -28,6 +28,7 @@ type const =
 type apply = {
   func : Variable.t;
   args : Variable.t list;
+  result_layout : Lambda.layout;
   kind : call_kind;
   dbg : Debuginfo.t;
   reg_close : Lambda.region_close;
@@ -50,6 +51,7 @@ type send = {
   dbg : Debuginfo.t;
   reg_close : Lambda.region_close;
   mode : Lambda.alloc_mode;
+  result_layout : Lambda.layout;
 }
 
 type project_closure = Projection.project_closure
@@ -75,7 +77,7 @@ type t =
   | String_switch of Variable.t * (string * t) list * t option
                      * Lambda.layout
   | Static_raise of Static_exception.t * Variable.t list
-  | Static_catch of Static_exception.t * Variable.t list * t * t * Lambda.layout
+  | Static_catch of Static_exception.t * ( Variable.t * Lambda.layout ) list * t * t * Lambda.layout
   | Try_with of t * Variable.t * t * Lambda.layout
   | While of t * t
   | For of for_loop
@@ -129,6 +131,7 @@ and function_declarations = {
 and function_declaration = {
   closure_origin: Closure_origin.t;
   params : Parameter.t list;
+  return_layout : Lambda.layout;
   alloc_mode : Lambda.alloc_mode;
   region : bool;
   body : t;
@@ -210,7 +213,7 @@ let rec lam ppf (flam : t) =
   match flam with
   | Var (id) ->
       Variable.print ppf id
-  | Apply({func; args; kind; inlined; probe; dbg}) ->
+  | Apply({func; args; kind; inlined; probe; dbg; result_layout}) ->
     let direct ppf () =
       match kind with
       | Indirect -> ()
@@ -229,8 +232,9 @@ let rec lam ppf (flam : t) =
       | None -> ()
       | Some {name} -> fprintf ppf "<probe %s>" name
     in
-    fprintf ppf "@[<2>(apply%a%a%a<%s>@ %a%a)@]" direct () inlined () probe ()
+    fprintf ppf "@[<2>(apply%a%a%a<%s>%a@ %a%a)@]" direct () inlined () probe ()
       (Debuginfo.to_string dbg)
+      Printlambda.layout result_layout
       Variable.print func Variable.print_list args
   | Assign { being_assigned; new_value; } ->
     fprintf ppf "@[<2>(assign@ %a@ %a)@]"
@@ -336,7 +340,7 @@ let rec lam ppf (flam : t) =
            | [] -> ()
            | _ ->
                List.iter
-                 (fun x -> fprintf ppf " %a" Variable.print x)
+                 (fun (x, _layout) -> fprintf ppf " %a" Variable.print x)
                  vars)
         vars
         lam lhandler
@@ -606,7 +610,7 @@ let rec variables_usage ?ignore_uses_as_callee ?ignore_uses_as_argument
       | Static_raise (_, es) ->
         List.iter free_variable es
       | Static_catch (_, vars, e1, e2, _) ->
-        List.iter bound_variable vars;
+        List.iter (fun (var, _layout) -> bound_variable var) vars;
         aux e1;
         aux e2
       | Try_with (e1, var, e2, _kind) ->
@@ -1042,6 +1046,7 @@ let update_body_of_function_declaration (func_decl: function_declaration)
       ~body : function_declaration =
   { closure_origin = func_decl.closure_origin;
     params = func_decl.params;
+    return_layout = func_decl.return_layout;
     alloc_mode = func_decl.alloc_mode;
     region = func_decl.region;
     body;
@@ -1065,6 +1070,7 @@ let rec check_param_modes mode = function
      check_param_modes m params
 
 let create_function_declaration ~params ~alloc_mode ~region ~body ~stub
+      ~(return_layout : Lambda.layout)
       ~(inline : Lambda.inline_attribute)
       ~(specialise : Lambda.specialise_attribute)
       ~(check : Lambda.check_attribute)
@@ -1093,6 +1099,7 @@ let create_function_declaration ~params ~alloc_mode ~region ~body ~stub
   check_param_modes alloc_mode params;
   { closure_origin;
     params;
+    return_layout;
     alloc_mode;
     region;
     body;