flambda-backend: Refactor check_representable in typedecl.ml (#2656)

* Move `check_representable`

* Call `check_representable` earlier

* Fix

* Temp promote test

* Push through `unboxed`

* Remove redundant checks

* Comments and formatting

* Small fixes

---------

Co-authored-by: Diana Kalinichenko <[email protected]>

Author: dkalinichenko-js

testsuite/tests/typing-layouts/basics.ml

@@ -605,14 +605,7 @@ Error: Layout void is more experimental than allowed by the enabled layouts exte
 
 type ('a : any) any4 = Any4 of 'a
 [%%expect{|
-Line 1, characters 23-33:
-1 | type ('a : any) any4 = Any4 of 'a
-                           ^^^^^^^^^^
-Error: Constructor argument types must have a representable layout.
-       The layout of 'a is any, because
-         of the annotation on 'a in the declaration of the type any4.
-       But the layout of 'a must be representable, because
-         it's the type of a constructor field.
+type 'a any4 = Any4 of 'a
 |}];;
 
 (************************************************************)

typing/typedecl.ml

@@ -381,7 +381,34 @@ let set_private_row env loc p decl =
   in
   set_type_desc rv (Tconstr (p, decl.type_params, ref Mnil))
 
-let transl_labels ~new_var_jkind env univars closed lbls =
+(* Makes sure a type is representable. When called with a type variable, will
+   lower [any] to a sort variable if [allow_unboxed = true], and to [value]
+   if [allow_unboxed = false]. *)
+(* CR layouts: Many places where [check_representable] is called in this file
+   should be replaced with checks at the places where values of those types are
+   constructed.  We've been conservative here in the first version. This is the
+   same issue as with arrows. *)
+let check_representable ~why ~allow_unboxed env loc kloc typ =
+  match Ctype.type_sort ~why env typ with
+  (* CR layouts v5: This is a convenient place to rule out non-value types in
+      structures that don't support them yet. (A callsite passes
+      [~allow_unboxed:true] to indicate that non-value types are allowed.)
+      When we support mixed blocks everywhere, this [check_representable]
+      will have outlived its usefulness and we can delete it.
+  *)
+  (* CR layouts v2.5: This rules out non-value types in [@@unboxed] types. No
+      real need to rule that out - I just haven't had time to write tests for it
+      yet. *)
+  | Ok s -> begin
+    if not allow_unboxed then
+      match Jkind.Sort.get_default_value s with
+      | Void | Value -> ()
+      | Float64 | Float32 | Word | Bits32 | Bits64 as const ->
+        raise (Error (loc, Invalid_jkind_in_block (typ, const, kloc)))
+    end
+  | Error err -> raise (Error (loc,Jkind_sort {kloc; typ; err}))
+
+let transl_labels ~new_var_jkind ~allow_unboxed env univars closed lbls kloc =
   assert (lbls <> []);
   let all_labels = ref String.Set.empty in
   List.iter
@@ -421,6 +448,8 @@ let transl_labels ~new_var_jkind env univars closed lbls =
       (fun ld ->
          let ty = ld.ld_type.ctyp_type in
          let ty = match get_desc ty with Tpoly(t,[]) -> t | _ -> ty in
+         check_representable ~why:(Label_declaration ld.ld_id)
+          ~allow_unboxed env ld.ld_loc kloc ty;
          {Types.ld_id = ld.ld_id;
           ld_mutable = ld.ld_mutable;
           ld_global = ld.ld_global;
@@ -435,24 +464,42 @@ let transl_labels ~new_var_jkind env univars closed lbls =
       lbls in
   lbls, lbls'
 
-let transl_types_gf ~new_var_jkind env univars closed tyl =
+let transl_types_gf ~new_var_jkind ~allow_unboxed
+  env loc univars closed tyl kloc =
   let mk arg =
     let cty = transl_simple_type ~new_var_jkind env ?univars ~closed Mode.Alloc.Const.legacy arg in
     let gf = Typemode.transl_global_flags
       (Jane_syntax.Mode_expr.of_attrs arg.ptyp_attributes |> fst) in
     (cty, gf)
   in
   let tyl_gfl = List.map mk tyl in
-  let tyl_gfl' = List.map (fun (cty, gf) -> cty.ctyp_type, gf) tyl_gfl in
+  let tyl_gfl' = List.mapi (fun idx (cty, gf) ->
+    check_representable ~why:(Constructor_declaration idx) ~allow_unboxed
+      env loc kloc cty.ctyp_type;
+    cty.ctyp_type, gf) tyl_gfl
+  in
   tyl_gfl, tyl_gfl'
 
-let transl_constructor_arguments ~new_var_jkind env univars closed = function
+let transl_constructor_arguments ~new_var_jkind ~unboxed
+  env loc univars closed = function
   | Pcstr_tuple l ->
-      let flds, flds' = transl_types_gf ~new_var_jkind env univars closed l in
-      Types.Cstr_tuple flds',
-      Cstr_tuple flds
+      let flds, flds' =
+        (* CR layouts: we forbid [@@unboxed] variants from being
+           non-value, see comment in [check_representable]. *)
+        transl_types_gf ~new_var_jkind ~allow_unboxed:(not unboxed)
+          env loc univars closed l (Cstr_tuple { unboxed })
+      in
+      Types.Cstr_tuple flds', Cstr_tuple flds
   | Pcstr_record l ->
-      let lbls, lbls' = transl_labels ~new_var_jkind env univars closed l in
+      let lbls, lbls' =
+        (* CR layouts: we forbid fields of inlined records from being
+           non-value, see comment in [check_representable].
+           When we allow mixed inline records, we still want to
+           disallow non-value types in unboxed records, so this
+           should become `not unboxed`, as in the `Pcstr_tuple` case. *)
+        transl_labels ~new_var_jkind ~allow_unboxed:false
+          env univars closed l (Inlined_record { unboxed })
+      in
       Types.Cstr_record lbls',
       Cstr_record lbls
 
@@ -462,7 +509,7 @@ let transl_constructor_arguments ~new_var_jkind env univars closed = function
    defined types. It is updated later by [update_constructor_arguments_jkinds]
 *)
 let make_constructor
-      env loc ~cstr_path ~type_path type_params (svars : _ Either.t)
+      env loc ~cstr_path ~type_path ~unboxed type_params (svars : _ Either.t)
       sargs sret_type =
   let tvars = match svars with
     | Left vars_only -> List.map (fun v -> v.txt, None) vars_only
@@ -484,7 +531,8 @@ let make_constructor
   match sret_type with
   | None ->
       let args, targs =
-        transl_constructor_arguments ~new_var_jkind:Any env None true sargs
+        transl_constructor_arguments ~new_var_jkind:Any ~unboxed
+          env loc None true sargs
       in
         tvars, targs, None, args, None
   | Some sret_type ->
@@ -510,7 +558,8 @@ let make_constructor
           in
           let univars = if closed then Some univar_list else None in
           let args, targs =
-            transl_constructor_arguments ~new_var_jkind:Sort env univars closed sargs
+            transl_constructor_arguments ~new_var_jkind:Sort ~unboxed
+              env loc univars closed sargs
           in
           let tret_type =
             transl_simple_type ~new_var_jkind:Sort env ?univars ~closed Mode.Alloc.Const.legacy
@@ -759,7 +808,7 @@ let transl_declaration env sdecl (id, uid) =
               attributes
           in
           let tvars, targs, tret_type, args, ret_type =
-            make_constructor env scstr.pcd_loc
+            make_constructor ~unboxed:unbox env scstr.pcd_loc
               ~cstr_path:(Path.Pident name) ~type_path:path params
               svars scstr.pcd_args scstr.pcd_res
           in
@@ -813,7 +862,12 @@ let transl_declaration env sdecl (id, uid) =
         in
           Ttype_variant tcstrs, Type_variant (cstrs, rep), jkind
       | Ptype_record lbls ->
-          let lbls, lbls' = transl_labels ~new_var_jkind:Any env None true lbls in
+          let lbls, lbls' =
+            (* CR layouts: we forbid [@@unboxed] records from being
+               non-value, see comment in [check_representable]. *)
+            transl_labels ~new_var_jkind:Any ~allow_unboxed:(not unbox)
+            env None true lbls (Record { unboxed = unbox })
+          in
           let rep, jkind =
             (* Note this is inaccurate, using `Record_boxed` in cases where the
                correct representation is [Record_float], [Record_ufloat], or
@@ -1089,45 +1143,13 @@ let check_coherence env loc dpath decl =
 let check_abbrev env sdecl (id, decl) =
   (id, check_coherence env sdecl.ptype_loc (Path.Pident id) decl)
 
-(* Makes sure a type is representable.  Will lower "any" to "value". *)
-(* CR layouts: In the places where this is used, we first call this to
-   ensure a type is representable, and then call [Ctype.type_jkind] to get the
-   most precise jkind.  These could be combined into some new function
-   [Ctype.type_jkind_representable] that avoids duplicated work *)
-(* CR layouts: Many places where [check_representable] is called in this file
-   should be replaced with checks at the places where values of those types are
-   constructed.  We've been conservative here in the first version. This is the
-   same issue as with arrows. *)
-let check_representable ~why ~allow_unboxed env loc kloc typ =
-  match Ctype.type_sort ~why env typ with
-  (* CR layouts v3: This is a convenient place to rule out non-value types in
-     structures that don't support them yet. (A callsite passes
-     [~allow_unboxed:true] to indicate that non-value types are allowed.)
-     When we support mixed blocks everywhere, this [check_representable]
-     will have outlived its usefulness and we can delete it.
-  *)
-  (* CR layouts v2.5: This rules out non-value types in [@@unboxed] types. No
-     real need to rule that out - I just haven't had time to write tests for it
-     yet. *)
-  | Ok s -> begin
-      match Jkind.Sort.get_default_value s with
-      (* All calls to this are part of [update_decl_jkind], which happens after
-         all the defaulting, so we don't expect this actually defaults the
-         sort - we just want the [const]. *)
-      | Void | Value -> ()
-      | Float64 | Float32 | Word | Bits32 | Bits64 as const ->
-          if not allow_unboxed then
-            raise (Error (loc, Invalid_jkind_in_block (typ, const, kloc)))
-    end
-  | Error err -> raise (Error (loc,Jkind_sort {kloc; typ; err}))
-
 (* The [update_x_jkinds] functions infer more precise jkinds in the type kind,
    including which fields of a record are void.  This would be hard to do during
    [transl_declaration] due to mutually recursive types.
 *)
 (* [update_label_jkinds] additionally returns whether all the jkinds
    were void *)
-let update_label_jkinds env loc lbls named ~is_inlined =
+let update_label_jkinds env loc lbls named =
   (* [named] is [Some jkinds] for top-level records (we will update the
      jkinds) and [None] for inlined records. *)
   (* CR layouts v5: it wouldn't be too hard to support records that are all
@@ -1137,15 +1159,8 @@ let update_label_jkinds env loc lbls named ~is_inlined =
     | None -> fun _ _ -> ()
     | Some jkinds -> fun idx jkind -> jkinds.(idx) <- jkind
   in
-  let kloc =
-    if is_inlined
-    then Inlined_record { unboxed = false }
-    else Record { unboxed = false }
-  in
   let lbls =
-    List.mapi (fun idx (Types.{ld_type; ld_id; ld_loc} as lbl) ->
-      check_representable ~why:(Label_declaration ld_id)
-        ~allow_unboxed:(Option.is_some named) env ld_loc kloc ld_type;
+    List.mapi (fun idx (Types.{ld_type} as lbl) ->
       let ld_jkind = Ctype.type_jkind env ld_type in
       update idx ld_jkind;
       {lbl with ld_jkind}
@@ -1163,13 +1178,11 @@ let update_constructor_arguments_jkinds env loc cd_args jkinds =
   match cd_args with
   | Types.Cstr_tuple tys ->
     List.iteri (fun idx (ty,_) ->
-      check_representable ~why:(Constructor_declaration idx) ~allow_unboxed:true
-        env loc (Cstr_tuple { unboxed = false }) ty;
       jkinds.(idx) <- Ctype.type_jkind env ty) tys;
     cd_args, Array.for_all Jkind.is_void_defaulting jkinds
   | Types.Cstr_record lbls ->
     let lbls, all_void =
-      update_label_jkinds env loc lbls None ~is_inlined:true
+      update_label_jkinds env loc lbls None
     in
     jkinds.(0) <- Jkind.value ~why:Boxed_record;
     Types.Cstr_record lbls, all_void
@@ -1401,14 +1414,12 @@ let update_decl_jkind env dpath decl =
   (* returns updated labels, updated rep, and updated jkind *)
   let update_record_kind loc lbls rep =
     match lbls, rep with
-    | [Types.{ld_type; ld_id; ld_loc} as lbl], Record_unboxed ->
-      check_representable ~why:(Label_declaration ld_id) ~allow_unboxed:false
-        env ld_loc (Record { unboxed = true }) ld_type;
+    | [Types.{ld_type} as lbl], Record_unboxed ->
       let ld_jkind = Ctype.type_jkind env ld_type in
       [{lbl with ld_jkind}], Record_unboxed, ld_jkind
     | _, Record_boxed jkinds ->
       let lbls, all_void =
-        update_label_jkinds env loc lbls (Some jkinds) ~is_inlined:false
+        update_label_jkinds env loc lbls (Some jkinds)
       in
       let jkind = Jkind.for_boxed_record ~all_void in
       let reprs =
@@ -1515,18 +1526,13 @@ let update_decl_jkind env dpath decl =
   let update_variant_kind cstrs rep =
     (* CR layouts: factor out duplication *)
     match cstrs, rep with
-    | [{Types.cd_args;cd_loc} as cstr], Variant_unboxed -> begin
+    | [{Types.cd_args} as cstr], Variant_unboxed -> begin
         match cd_args with
         | Cstr_tuple [ty,_] -> begin
-            check_representable ~why:(Constructor_declaration 0)
-              ~allow_unboxed:false env cd_loc (Cstr_tuple { unboxed = true }) ty;
             let jkind = Ctype.type_jkind env ty in
             cstrs, Variant_unboxed, jkind
           end
-        | Cstr_record [{ld_type; ld_id; ld_loc} as lbl] -> begin
-            check_representable ~why:(Label_declaration ld_id)
-              ~allow_unboxed:false env ld_loc (Inlined_record { unboxed = true })
-              ld_type;
+        | Cstr_record [{ld_type} as lbl] -> begin
             let ld_jkind = Ctype.type_jkind env ld_type in
             [{ cstr with Types.cd_args =
                            Cstr_record [{ lbl with ld_jkind }] }],
@@ -2217,7 +2223,7 @@ let transl_extension_constructor_decl
       env type_path typext_params loc id svars sargs sret_type =
   let tvars, targs, tret_type, args, ret_type =
     make_constructor env loc
-      ~cstr_path:(Pident id) ~type_path typext_params
+      ~cstr_path:(Pident id) ~type_path ~unboxed:false typext_params
       svars sargs sret_type
   in
   let num_args =