flambda-backend: Add [@layout_poly] attribute for layout representa…

…tion polymorphic externals (#2229)

* wip

fix printing

remove Repr_poly from lambda

add comments

proper jkind reason

test update

fix typo

produce correct error

check repr attr earlier and use get_unboxed_type_approximation

make more primitives work with non value layout

format

fix upstream

flambda2 only

rename tests

tests

code cleanup

more cleanup and an extra test

add additional checks around array primitives

* remove paren

* add test when type constructor constrains layout

* rename is_lambda_prim to is_builtin_prim

* test&comment about No_native_primitive_with_non_value and Old_style_float_with_non_value

* revert indentation change to reduce diff

* make ml/mli def of description_gen consistent

* fix long lines

* comment on make_prim_repr_args

* fix module tests

* rework jkind checks

* add comment on extern_repr

* move make_prim_repr_args comment to mli

* more comments

* prim_is_layout_representation_polymorphic rename

* add test about non-explicitly quantify tvars

* update type_sort_external comment

* test update

* rename to layout_poly and improve error messages

* fix format

* make the repr checks stronger

* rename test files

* register layout_poly in Builtin_attributes

* improve documentation

* code cleanup & rename

* more tests for instance_prim

* more documentation and tests around the jkind check

* Formatting and minor edits in comments

* make unused layout_poly an error

* update error message

* fix upstream build and bootstrap

---------

Co-authored-by: Chris Casinghino <ccasinghino@janestreet.com>

asmcomp/cmm_helpers.ml

@@ -1571,7 +1571,7 @@ let box_sized size mode dbg exp =
 (* Simplification of some primitives into C calls *)
 
 let default_prim name =
-  Primitive.simple_on_values ~name ~arity:0(*ignored*) ~alloc:true
+  Lambda.simple_prim_on_values ~name ~arity:0(*ignored*) ~alloc:true
 
 let simplif_primitive p : Clambda_primitives.primitive =
   match (p : Clambda_primitives.primitive) with

asmcomp/cmmgen.ml

@@ -590,7 +590,7 @@ let rec transl env e =
           transl_make_array dbg env kind alloc_heap args
       | (Pduparray _, [arg]) ->
           let prim_obj_dup =
-            Primitive.simple_on_values ~name:"caml_obj_dup" ~arity:1 ~alloc:true
+            Lambda.simple_prim_on_values ~name:"caml_obj_dup" ~arity:1 ~alloc:true
           in
           transl_ccall env prim_obj_dup [arg] dbg
       | (Pmakearray _, []) ->

bytecomp/bytegen.ml

@@ -894,7 +894,7 @@ let rec comp_expr stack_info env exp sz cont =
         (Lprim (Pmakearray (kind, mutability, m), args, loc)) sz cont
   | Lprim (Pduparray _, [arg], loc) ->
       let prim_obj_dup =
-        Primitive.simple_on_values ~name:"caml_obj_dup" ~arity:1 ~alloc:true
+        Lambda.simple_prim_on_values ~name:"caml_obj_dup" ~arity:1 ~alloc:true
       in
       comp_expr stack_info env (Lprim (Pccall prim_obj_dup, [arg], loc)) sz cont
   | Lprim (Pduparray _, _, _) ->

lambda/lambda.ml

@@ -163,7 +163,7 @@ type primitive =
   | Presume
   | Preperform
   (* External call *)
-  | Pccall of Primitive.description
+  | Pccall of external_call_description
   (* Exceptions *)
   | Praise of raise_kind
   (* Boolean operations *)
@@ -283,6 +283,15 @@ type primitive =
   (* Fetching domain-local state *)
   | Pdls_get
 
+and extern_repr =
+  | Same_as_ocaml_repr of Jkind.Sort.const
+  | Unboxed_float of boxed_float
+  | Unboxed_vector of Primitive.boxed_vector
+  | Unboxed_integer of Primitive.boxed_integer
+  | Untagged_int
+
+and external_call_description = extern_repr Primitive.description_gen
+
 and integer_comparison =
     Ceq | Cne | Clt | Cgt | Cle | Cge
 
@@ -1496,7 +1505,7 @@ let mod_field ?(read_semantics=Reads_agree) pos =
 let mod_setfield pos =
   Psetfield (pos, Pointer, Root_initialization)
 
-let alloc_mode_of_primitive_description (p : Primitive.description) =
+let alloc_mode_of_primitive_description (p : external_call_description) =
   if not Config.stack_allocation then
     if p.prim_alloc then Some alloc_heap else None
   else
@@ -1634,7 +1643,7 @@ let structured_constant_layout = function
   | Const_block _ | Const_immstring _ -> Pvalue Pgenval
   | Const_float_array _ | Const_float_block _ -> Pvalue (Parrayval Pfloatarray)
 
-let layout_of_native_repr : Primitive.native_repr -> _ = function
+let layout_of_extern_repr : extern_repr -> _ = function
   | Untagged_int ->  layout_int
   | Unboxed_vector v -> layout_boxed_vector v
   | Unboxed_float bf -> layout_boxed_float bf
@@ -1682,7 +1691,7 @@ let primitive_result_layout (p : primitive) =
   | Pbox_float (f, _) -> layout_boxed_float f
   | Pufloatfield _ -> Punboxed_float Pfloat64
   | Punbox_float float_kind -> Punboxed_float float_kind
-  | Pccall { prim_native_repr_res = _, repr_res } -> layout_of_native_repr repr_res
+  | Pccall { prim_native_repr_res = _, repr_res } -> layout_of_extern_repr repr_res
   | Praise _ -> layout_bottom
   | Psequor | Psequand | Pnot
   | Pnegint | Paddint | Psubint | Pmulint
@@ -1869,3 +1878,17 @@ let may_allocate_in_region lam =
     | () -> false
     | exception Exit -> true
   end
+
+let simple_prim_on_values ~name ~arity ~alloc =
+  Primitive.make
+    ~name
+    ~alloc
+    ~c_builtin:false
+    ~effects:Arbitrary_effects
+    ~coeffects:Has_coeffects
+    ~native_name:""
+    ~native_repr_args:
+      (Primitive.make_prim_repr_args arity
+        (Primitive.Prim_global,Same_as_ocaml_repr Jkind.Sort.Value))
+    ~native_repr_res:(Prim_global, Same_as_ocaml_repr Jkind.Sort.Value)
+    ~is_layout_poly:false

lambda/lambda.mli

@@ -121,7 +121,7 @@ type primitive =
   | Presume
   | Preperform
   (* External call *)
-  | Pccall of Primitive.description
+  | Pccall of external_call_description
   (* Exceptions *)
   | Praise of raise_kind
   (* Boolean operations *)
@@ -251,6 +251,17 @@ type primitive =
   (* Fetching domain-local state *)
   | Pdls_get
 
+(** This is the same as [Primitive.native_repr] but with [Repr_poly]
+    compiled away. *)
+and extern_repr =
+  | Same_as_ocaml_repr of Jkind.Sort.const
+  | Unboxed_float of boxed_float
+  | Unboxed_vector of Primitive.boxed_vector
+  | Unboxed_integer of Primitive.boxed_integer
+  | Untagged_int
+
+and external_call_description = extern_repr Primitive.description_gen
+
 and integer_comparison =
     Ceq | Cne | Clt | Cgt | Cle | Cge
 
@@ -372,12 +383,12 @@ val equal_boxed_vector_size : boxed_vector -> boxed_vector -> bool
 val must_be_value : layout -> value_kind
 
 (* This is the layout of ocaml values used as arguments to or returned from
-   primitives for this [native_repr].  So the legacy [Unboxed_float] - which is
+   primitives for this [extern_repr].  So the legacy [Unboxed_float] - which is
    a float that is unboxed before being passed to a C function - is mapped to
    [layout_any_value], while [Same_as_ocaml_repr Float64] is mapped to
    [layout_unboxed_float].
 *)
-val layout_of_native_repr : Primitive.native_repr -> layout
+val layout_of_extern_repr : extern_repr -> layout
 
 type structured_constant =
     Const_base of constant
@@ -788,7 +799,7 @@ val primitive_may_allocate : primitive -> alloc_mode option
   *)
 
 val alloc_mode_of_primitive_description :
-  Primitive.description -> alloc_mode option
+  external_call_description -> alloc_mode option
   (** Like [primitive_may_allocate], for [external] calls. *)
 
 (***********************)
@@ -837,3 +848,11 @@ val is_check_enabled : opt:bool -> property -> bool
 
 (* Returns true if the given lambda can allocate on the local stack *)
 val may_allocate_in_region : lambda -> bool
+
+(* Returns [external_call_description] for [Pccall] assuming arguments
+   and result all have layout [value] *)
+val simple_prim_on_values
+:  name:string
+-> arity:int
+-> alloc:bool
+-> external_call_description

lambda/matching.ml

@@ -1915,7 +1915,7 @@ let get_pat_args_lazy p rem =
 *)
 
 let prim_obj_tag =
-  Primitive.simple_on_values ~name:"caml_obj_tag" ~arity:1 ~alloc:false
+  Lambda.simple_prim_on_values ~name:"caml_obj_tag" ~arity:1 ~alloc:false
 
 let get_mod_field modname field =
   lazy
@@ -2255,11 +2255,11 @@ let divide_array ~scopes kind ctx pm =
 let strings_test_threshold = 8
 
 let prim_string_notequal =
-  Pccall (Primitive.simple_on_values ~name:"caml_string_notequal" ~arity:2
+  Pccall (Lambda.simple_prim_on_values ~name:"caml_string_notequal" ~arity:2
             ~alloc:false)
 
 let prim_string_compare =
-  Pccall (Primitive.simple_on_values ~name:"caml_string_compare" ~arity:2
+  Pccall (Lambda.simple_prim_on_values ~name:"caml_string_compare" ~arity:2
             ~alloc:false)
 
 let bind_sw arg layout k =

lambda/transl_comprehension_utils.ml

@@ -124,7 +124,7 @@ module Lambda_utils = struct
     (** The Lambda primitive for calling a simple C primitive *)
     (* CR layouts v4: To change when non-values are allowed in arrays. *)
     let c_prim name arity =
-      Pccall (Primitive.simple_on_values ~name ~arity ~alloc:true)
+      Pccall (Lambda.simple_prim_on_values ~name ~arity ~alloc:true)
 
     (** Create a function that produces the Lambda representation for a
         one-argument C primitive when provided with a Lambda argument *)

lambda/translcore.ml

@@ -102,7 +102,7 @@ let declare_probe_handlers lam =
 
 let prim_fresh_oo_id =
   Pccall
-    (Primitive.simple_on_values ~name:"caml_fresh_oo_id" ~arity:1 ~alloc:false)
+    (Lambda.simple_prim_on_values ~name:"caml_fresh_oo_id" ~arity:1 ~alloc:false)
 
 let transl_extension_constructor ~scopes env path ext =
   let path =
@@ -345,8 +345,8 @@ let rec iter_exn_names f pat =
 
 let transl_ident loc env ty path desc kind =
   match desc.val_kind, kind with
-  | Val_prim p, Id_prim poly_mode ->
-      Translprim.transl_primitive loc p env ty ~poly_mode (Some path)
+  | Val_prim p, Id_prim (poly_mode, poly_sort) ->
+      Translprim.transl_primitive loc p env ty ~poly_mode ~poly_sort (Some path)
   | Val_anc _, Id_value ->
       raise(Error(to_location loc, Free_super_var))
   | (Val_reg | Val_self _), Id_value ->
@@ -404,7 +404,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
       transl_function ~in_new_scope ~scopes e params body
         ~alloc_mode ~ret_mode ~ret_sort ~region
   | Texp_apply({ exp_desc = Texp_ident(path, _, {val_kind = Val_prim p},
-                                       Id_prim pmode, _);
+                                       Id_prim (pmode, psort), _);
                 exp_type = prim_type; } as funct, oargs, pos, ap_mode)
     when can_apply_primitive p pmode pos oargs ->
       let rec cut_args prim_repr oargs =
@@ -413,7 +413,10 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
         | _, [] -> failwith "Translcore cut_args"
         | ((_, arg_repr) :: prim_repr), ((_, Arg (x, _)) :: oargs) ->
           let arg_exps, extra_args = cut_args prim_repr oargs in
-          let arg_sort = Jkind.Sort.of_const (sort_of_native_repr arg_repr) in
+          let arg_sort =
+            Jkind.Sort.of_const
+              (Translprim.sort_of_native_repr arg_repr ~poly_sort:psort)
+          in
           (x, arg_sort) :: arg_exps, extra_args
         | _, ((_, Omitted _) :: _) -> assert false
       in
@@ -426,7 +429,8 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
       in
       let lam =
         Translprim.transl_primitive_application
-          (of_location ~scopes e.exp_loc) p e.exp_env prim_type pmode
+          (of_location ~scopes e.exp_loc) p e.exp_env prim_type
+          ~poly_mode:pmode ~poly_sort:psort
           path prim_exp args (List.map fst arg_exps) position
       in
       if extra_args = [] then lam

lambda/translmod.ml

@@ -129,8 +129,11 @@ let rec apply_coercion loc strict restr arg =
         [{name = param; layout = Lambda.layout_module;
           attributes = Lambda.default_param_attribute; mode = alloc_heap}]
         [carg] cc_res
-  | Tcoerce_primitive { pc_desc; pc_env; pc_type; pc_poly_mode } ->
-      Translprim.transl_primitive loc pc_desc pc_env pc_type ~poly_mode:pc_poly_mode None
+  | Tcoerce_primitive { pc_desc; pc_env; pc_type; pc_poly_mode; pc_poly_sort } ->
+      Translprim.transl_primitive loc pc_desc pc_env pc_type
+        ~poly_mode:pc_poly_mode
+        ~poly_sort:pc_poly_sort
+        None
   | Tcoerce_alias (env, path, cc) ->
       let lam = transl_module_path loc env path in
       name_lambda strict arg Lambda.layout_module
@@ -268,7 +271,7 @@ let record_primitive = function
 let preallocate_letrec ~bindings ~body =
   assert (Clflags.is_flambda2 ());
   let caml_update_dummy_prim =
-    Primitive.simple_on_values ~name:"caml_update_dummy" ~arity:2 ~alloc:true
+    Lambda.simple_prim_on_values ~name:"caml_update_dummy" ~arity:2 ~alloc:true
   in
   let update_dummy var expr =
     Lprim (Pccall caml_update_dummy_prim, [Lvar var; expr], Loc_unknown)
@@ -282,7 +285,7 @@ let preallocate_letrec ~bindings ~body =
   List.fold_left
     (fun body (id, _def, size) ->
        let desc =
-         Primitive.simple_on_values ~name:"caml_alloc_dummy" ~arity:1
+         Lambda.simple_prim_on_values ~name:"caml_alloc_dummy" ~arity:1
            ~alloc:true
        in
        let size : lambda = Lconst (Const_base (Const_int size)) in
@@ -687,7 +690,10 @@ and transl_structure ~scopes loc fields cc rootpath final_env = function
                       | Tcoerce_primitive p ->
                           let loc = of_location ~scopes p.pc_loc in
                           Translprim.transl_primitive
-                            loc p.pc_desc p.pc_env p.pc_type ~poly_mode:p.pc_poly_mode None
+                            loc p.pc_desc p.pc_env p.pc_type
+                            ~poly_mode:p.pc_poly_mode
+                            ~poly_sort:p.pc_poly_sort
+                            None
                       | _ -> apply_coercion loc Strict cc (get_field pos))
                     pos_cc_list, loc)
             and id_pos_list =
@@ -1134,8 +1140,11 @@ let field_of_str loc str =
   let ids = Array.of_list (defined_idents str.str_items) in
   fun (pos, cc) ->
     match cc with
-    | Tcoerce_primitive { pc_desc; pc_env; pc_type; pc_poly_mode } ->
-        Translprim.transl_primitive loc pc_desc pc_env pc_type ~poly_mode:pc_poly_mode None
+    | Tcoerce_primitive { pc_desc; pc_env; pc_type; pc_poly_mode; pc_poly_sort } ->
+        Translprim.transl_primitive loc pc_desc pc_env pc_type
+          ~poly_mode:pc_poly_mode
+          ~poly_sort:pc_poly_sort
+          None
     | Tcoerce_alias (env, path, cc) ->
         let lam = transl_module_path loc env path in
         apply_coercion loc Alias cc lam
@@ -1471,7 +1480,10 @@ let transl_store_structure ~scopes glob map prims aliases str =
     Lsequence(Lprim(mod_setfield pos,
                     [Lprim(Pgetglobal glob, [], Loc_unknown);
                      Translprim.transl_primitive Loc_unknown
-                       prim.pc_desc prim.pc_env prim.pc_type ~poly_mode:prim.pc_poly_mode None],
+                       prim.pc_desc prim.pc_env prim.pc_type
+                       ~poly_mode:prim.pc_poly_mode
+                       ~poly_sort:prim.pc_poly_sort
+                       None],
                     Loc_unknown),
               cont)
 

lambda/translobj.ml

@@ -84,7 +84,7 @@ let int n = Lconst (Const_base (Const_int n))
 
 (* CR layouts v5: To change when we have arrays of other sorts *)
 let prim_makearray =
-  Primitive.simple_on_values ~name:"caml_make_vect" ~arity:2 ~alloc:true
+  Lambda.simple_prim_on_values ~name:"caml_make_vect" ~arity:2 ~alloc:true
 
 (* Also use it for required globals *)
 let transl_label_init_general f =