Source file primitive.ml

(**************************************************************************)
(*                                                                        *)
(*                                 OCaml                                  *)
(*                                                                        *)
(*             Xavier Leroy, projet Cristal, INRIA Rocquencourt           *)
(*                                                                        *)
(*   Copyright 1996 Institut National de Recherche en Informatique et     *)
(*     en Automatique.                                                    *)
(*                                                                        *)
(*   All rights reserved.  This file is distributed under the terms of    *)
(*   the GNU Lesser General Public License version 2.1, with the          *)
(*   special exception on linking described in the file LICENSE.          *)
(*                                                                        *)
(**************************************************************************)

(* Description of primitive functions *)

open Misc
open Parsetree

module String = Misc.Stdlib.String

type unboxed_or_untagged_integer =
  | Unboxed_int64
  | Unboxed_nativeint
  | Unboxed_int32
  | Untagged_int16
  | Untagged_int8
  | Untagged_int

type unboxed_float = Unboxed_float64 | Unboxed_float32
type unboxed_vector = Unboxed_vec128 | Unboxed_vec256 | Unboxed_vec512

type boxed_integer = Boxed_int64 | Boxed_nativeint | Boxed_int32
type boxed_float = Boxed_float64 | Boxed_float32
type boxed_vector = Boxed_vec128 | Boxed_vec256 | Boxed_vec512

type native_repr =
  | Repr_poly
  | Same_as_ocaml_repr of Jkind_types.Sort.Const.t
  | Unboxed_float of boxed_float
  | Unboxed_vector of boxed_vector
  | Unboxed_or_untagged_integer of unboxed_or_untagged_integer
  | Unpacked_product of Jkind_types.Sort.Const.t

type effects = No_effects | Only_generative_effects | Arbitrary_effects
type coeffects = No_coeffects | Has_coeffects

type mode =
  | Prim_local
  | Prim_global
  | Prim_poly

type 'repr description_gen =
  { prim_name: string;         (* Name of primitive  or C function *)
    prim_arity: int;           (* Number of arguments *)
    prim_alloc: bool;          (* Does it allocates or raise? *)
    prim_c_builtin: bool;        (* Is the compiler allowed to replace it? *)
    prim_effects: effects;
    prim_coeffects: coeffects;
    prim_native_name: string;  (* Name of C function for the nat. code gen. *)
    prim_native_repr_args: (mode * 'repr) list;
    prim_native_repr_res: mode * 'repr;
    prim_is_layout_poly: bool }

type description = native_repr description_gen

(* Is [[@@immediate]], so poly-compare can be used to de-dup errors. If a [loc]
   is added, then errors will always be unique and de-duping will no longer be
   necessary. *)
type wrong_repr_error =
  | Product_arg
  | Expected_value_prim
  | Product_return
  | Unpacked_product_return
  | Repr_mismatch
[@@immediate]

let compare_wrong_repr_error (a : wrong_repr_error) (b : wrong_repr_error) =
  Stdlib.compare a b

type error =
  | Old_style_float_with_native_repr_attribute
  | Old_style_float_with_non_value
  | Old_style_noalloc_with_noalloc_attribute
  | No_native_primitive_with_repr_attribute
  | No_native_primitive_with_non_value
  | Inconsistent_attributes_for_effects
  | Inconsistent_noalloc_attributes_for_effects
  | Invalid_representation_polymorphic_attribute
  | Invalid_native_repr_for_primitive of
      { prim_name : string; errors : wrong_repr_error list }

exception Error of Location.t * error

type value_check = Bad_attribute | Bad_layout | Ok_value

let check_ocaml_value = function
  | _, Same_as_ocaml_repr (Base Scannable) -> Ok_value
  | _, Same_as_ocaml_repr _
  | _, Repr_poly -> Bad_layout
  | _, Unboxed_float _
  | _, Unboxed_vector _
  | _, Unboxed_or_untagged_integer _
  | _, Unpacked_product _ -> Bad_attribute

let is_builtin_prim_name name = String.length name > 0 && name.[0] = '%'

let rec make_prim_repr_args arity x =
  if arity = 0 then
    []
  else
    x :: make_prim_repr_args (arity - 1) x

let make ~name ~alloc ~c_builtin ~effects ~coeffects
      ~native_name ~native_repr_args ~native_repr_res
      ~is_layout_poly =
  {prim_name = name;
   prim_arity = List.length native_repr_args;
   prim_alloc = alloc;
   prim_c_builtin = c_builtin;
   prim_effects = effects;
   prim_coeffects = coeffects;
   prim_native_name = native_name;
   prim_native_repr_args = native_repr_args;
   prim_native_repr_res = native_repr_res;
   prim_is_layout_poly = is_layout_poly }

let parse_declaration valdecl ~native_repr_args ~native_repr_res ~is_layout_poly =
  let arity = List.length native_repr_args in
  let name, native_name, old_style_noalloc, old_style_float =
    match valdecl.pval_prim with
    | name :: "noalloc" :: name2 :: "float" :: _ -> (name, name2, true, true)
    | name :: "noalloc" :: name2 :: _ -> (name, name2, true, false)
    | name :: name2 :: "float" :: _ -> (name, name2, false, true)
    | name :: "noalloc" :: _ -> (name, "", true, false)
    | name :: name2 :: _ -> (name, name2, false, false)
    | name :: _ -> (name, "", false, false)
    | [] ->
        fatal_error "Primitive.parse_declaration"
  in
  let noalloc_attribute =
    Attr_helper.has_no_payload_attribute "noalloc" valdecl.pval_attributes
  in
  let builtin_attribute =
    Attr_helper.has_no_payload_attribute "builtin" valdecl.pval_attributes
  in
  let no_effects_attribute =
    Attr_helper.has_no_payload_attribute "no_effects" valdecl.pval_attributes
  in
  let only_generative_effects_attribute =
    Attr_helper.has_no_payload_attribute "only_generative_effects"
      valdecl.pval_attributes
  in
  let is_builtin_prim = is_builtin_prim_name name in
  let prim_is_layout_poly =
    match is_builtin_prim, is_layout_poly with
    | false, true ->  raise (Error (valdecl.pval_loc,
                        Invalid_representation_polymorphic_attribute))
    | _, b -> b
  in
  if no_effects_attribute && only_generative_effects_attribute then
    raise (Error (valdecl.pval_loc,
                  Inconsistent_attributes_for_effects));
  let effects =
    if no_effects_attribute then No_effects
    else if only_generative_effects_attribute then Only_generative_effects
    else Arbitrary_effects
  in
  let no_coeffects_attribute =
    Attr_helper.has_no_payload_attribute "no_coeffects" valdecl.pval_attributes
  in
  let coeffects =
    if no_coeffects_attribute then No_coeffects
    else Has_coeffects
  in
  if old_style_float then
    List.iter
      (fun repr -> match check_ocaml_value repr with
         | Ok_value -> ()
         | Bad_attribute ->
           raise (Error (valdecl.pval_loc,
                         Old_style_float_with_native_repr_attribute))
         | Bad_layout ->
           raise (Error (valdecl.pval_loc,
                         Old_style_float_with_non_value)))
      (native_repr_res :: native_repr_args);
  if old_style_noalloc && noalloc_attribute then
    raise (Error (valdecl.pval_loc,
                  Old_style_noalloc_with_noalloc_attribute));
  (* The compiler used to assume "noalloc" with "float", we just make this
     explicit now (GPR#167): *)
  let old_style_noalloc = old_style_noalloc || old_style_float in
  if old_style_float then
    Location.deprecated valdecl.pval_loc
      "[@@unboxed] + [@@noalloc] should be used\n\
       instead of \"float\""
  else if old_style_noalloc then
    Location.deprecated valdecl.pval_loc
      "[@@noalloc] should be used instead of \"noalloc\"";
  if native_name = "" then
    List.iter
      (fun repr -> match check_ocaml_value repr with
         | Ok_value -> ()
         | Bad_attribute ->
           raise (Error (valdecl.pval_loc,
                         No_native_primitive_with_repr_attribute))
         | Bad_layout ->
           (* Built-in primitives don't need a native version. *)
           if not is_builtin_prim then
             raise (Error (valdecl.pval_loc,
                           No_native_primitive_with_non_value)))
      (native_repr_res :: native_repr_args);
  let noalloc = old_style_noalloc || noalloc_attribute in
  if noalloc && only_generative_effects_attribute then
    raise (Error (valdecl.pval_loc,
                  Inconsistent_noalloc_attributes_for_effects));
  let native_repr_args, native_repr_res =
    if old_style_float then
      (make_prim_repr_args arity (Prim_global, Unboxed_float Boxed_float64),
       (Prim_global, Unboxed_float Boxed_float64))
    else
      (native_repr_args, native_repr_res)
  in
  {prim_name = name;
   prim_arity = arity;
   prim_alloc = not noalloc;
   prim_c_builtin = builtin_attribute;
   prim_effects = effects;
   prim_coeffects = coeffects;
   prim_native_name = native_name;
   prim_native_repr_args = native_repr_args;
   prim_native_repr_res = native_repr_res;
   prim_is_layout_poly }

open Outcometree

let add_attribute_list ty attrs =
  List.fold_left (fun ty attr -> Otyp_attribute(ty, attr)) ty attrs

let rec add_native_repr_attributes ty attrs =
  match ty, attrs with
    (* Otyp_poly and Otyp_newlayout case might have been added in e.g.
       tree_of_value_description *)
  | Otyp_poly (vars, ty), _ -> Otyp_poly (vars, add_native_repr_attributes ty attrs)
  | Otyp_newlayout (vars, ty), _ ->
    Otyp_newlayout (vars, add_native_repr_attributes ty attrs)
  | Otyp_arrow (label, am, a, Otyp_ret (rm, r)), attr_l :: rest ->
    let r = add_native_repr_attributes r rest in
    let a = add_attribute_list a attr_l in
    Otyp_arrow (label, am, a, Otyp_ret (rm, r))
  | _, [attr_l] -> add_attribute_list ty attr_l
  | _ ->
    assert (List.for_all (fun x -> x = []) attrs);
    ty

let oattr_unboxed = { oattr_name = "unboxed" }
let oattr_untagged = { oattr_name = "untagged" }
let oattr_unpacked = { oattr_name = "unpacked" }
let oattr_noalloc = { oattr_name = "noalloc" }
let oattr_builtin = { oattr_name = "builtin" }
let oattr_no_effects = { oattr_name = "no_effects" }
let oattr_only_generative_effects = { oattr_name = "only_generative_effects" }
let oattr_no_coeffects = { oattr_name = "no_coeffects" }
let oattr_local_opt = { oattr_name = "local_opt" }
let oattr_layout_poly = { oattr_name = "layout_poly" }

let print p osig_val_decl =
  let prims =
    if p.prim_native_name <> "" then
      [p.prim_name; p.prim_native_name]
    else
      [p.prim_name]
  in
  let for_all f =
    List.for_all f p.prim_native_repr_args && f p.prim_native_repr_res
  in
  let needs_unboxed_attribute = function
    | _, Same_as_ocaml_repr (Base Scannable)
    | _, Repr_poly
    | _, Unpacked_product _
    | _, Unboxed_or_untagged_integer (Untagged_int | Untagged_int8
                                    | Untagged_int16) -> false
    | _, Unboxed_float _
    | _, Unboxed_vector _
    | _, Unboxed_or_untagged_integer (Unboxed_int64 | Unboxed_int32
                                    | Unboxed_nativeint) ->
      true
    | _, Same_as_ocaml_repr (Base Void) -> false
      (* The [@unboxed] attribute is not supported on void in our compiler
         or the upstream compiler. *)
    | _, Same_as_ocaml_repr _ ->
      (* We require [@unboxed] for non-value types in upstream-compatible code,
         but treat it as optional otherwise. We thus print the [@unboxed]
         attribute only in the case it's required and leave it out when it's
         not. That's why we call [erasable_extensions_only] here. *)
      Language_extension.erasable_extensions_only ()
  in
  let is_untagged = function
    | _, Unboxed_or_untagged_integer (Untagged_int8 | Untagged_int16) ->
      Language_extension.erasable_extensions_only ()
    | _, Unboxed_or_untagged_integer (Untagged_int) -> true
    | _, Same_as_ocaml_repr _
    | _, Unboxed_float _
    | _, Unboxed_vector _
    | _, Unboxed_or_untagged_integer (Unboxed_int64 | Unboxed_int32
                                    | Unboxed_nativeint)
    | _, Unpacked_product _
    | _, Repr_poly -> false
  in
  let all_unboxed = for_all needs_unboxed_attribute in
  let all_untagged = for_all is_untagged in
  let attrs = if p.prim_alloc then [] else [oattr_noalloc] in
  let attrs = if p.prim_c_builtin then oattr_builtin::attrs else attrs in
  let attrs = match p.prim_effects with
    | No_effects -> oattr_no_effects::attrs
    | Only_generative_effects -> oattr_only_generative_effects::attrs
    | Arbitrary_effects -> attrs
  in
  let attrs = match p.prim_coeffects with
    | No_coeffects -> oattr_no_coeffects::attrs
    | Has_coeffects -> attrs
  in
  let attrs =
    if all_unboxed then
      oattr_unboxed :: attrs
    else if all_untagged then
      oattr_untagged :: attrs
    else
      attrs
  in
  let attrs =
    if p.prim_is_layout_poly then
      oattr_layout_poly :: attrs
    else
      attrs
  in
  let attrs_of_mode_and_repr (m, repr) =
    (match m with
     | Prim_local | Prim_global -> []
     | Prim_poly -> [oattr_local_opt])
    @
    (match repr with
     | Same_as_ocaml_repr (Base Scannable)
     | Repr_poly -> []
     | Unboxed_float _
     | Unboxed_vector _
     | Unboxed_or_untagged_integer (Unboxed_int32 | Unboxed_int64
                                   | Unboxed_nativeint) ->
       if all_unboxed then [] else [oattr_unboxed]
     | Unboxed_or_untagged_integer (Untagged_int | Untagged_int8
                                   | Untagged_int16) ->
       if all_untagged then [] else [oattr_untagged]
     | Unpacked_product _ -> [oattr_unpacked]
     | Same_as_ocaml_repr _->
       if all_unboxed || not (needs_unboxed_attribute (m, repr))
       then []
       else [oattr_unboxed])
  in
  let type_attrs =
    List.map attrs_of_mode_and_repr p.prim_native_repr_args @
    [attrs_of_mode_and_repr p.prim_native_repr_res]
  in
  { osig_val_decl with
    oval_prims = prims;
    oval_type = add_native_repr_attributes osig_val_decl.oval_type type_attrs;
    oval_attributes = attrs }

let native_name p =
  if p.prim_native_name <> ""
  then p.prim_native_name
  else p.prim_name

let byte_name p =
  p.prim_name

let unboxed_or_untagged_integer = function
  | Boxed_int32 -> Unboxed_int32
  | Boxed_nativeint -> Unboxed_nativeint
  | Boxed_int64 -> Unboxed_int64

let unboxed_float = function
  | Boxed_float32 -> Unboxed_float32
  | Boxed_float64 -> Unboxed_float64

let unboxed_vector = function
  | Boxed_vec128 -> Unboxed_vec128
  | Boxed_vec256 -> Unboxed_vec256
  | Boxed_vec512 -> Unboxed_vec512

(* Since these are just constant constructors, we can just use polymorphic equality and
   comparison at no performance loss. We still match on the variants to prove here that
   they are all constant constructors. *)
let equal_unboxed_or_untagged_integer
      ((Unboxed_int64 | Unboxed_nativeint | Unboxed_int32 | Untagged_int16
        | Untagged_int8 | Untagged_int) as i1) i2 =
  i1 = i2
let equal_unboxed_float
      ((Unboxed_float32 | Unboxed_float64) as f1) f2 = f1 = f2
let compare_unboxed_float
      ((Unboxed_float32 | Unboxed_float64) as f1) f2 = Stdlib.compare f1 f2
let equal_unboxed_vector
      ((Unboxed_vec128 | Unboxed_vec256 | Unboxed_vec512) as v1) v2 = v1 = v2
let compare_unboxed_vector
      ((Unboxed_vec128 | Unboxed_vec256 | Unboxed_vec512) as v1) v2 =
      Stdlib.compare v1 v2

let equal_boxed_integer bi1 bi2 =
  equal_unboxed_or_untagged_integer (unboxed_or_untagged_integer bi1)
    (unboxed_or_untagged_integer bi2)
let equal_boxed_float bf1 bf2 =
  equal_unboxed_float (unboxed_float bf1) (unboxed_float bf2)
let equal_boxed_vector bv1 bv2 =
  equal_unboxed_vector (unboxed_vector bv1) (unboxed_vector bv2)
let compare_boxed_float bf1 bf2 =
  compare_unboxed_float (unboxed_float bf1) (unboxed_float bf2)
let compare_boxed_vector bv1 bv2 =
  compare_unboxed_vector (unboxed_vector bv1) (unboxed_vector bv2)

let equal_unboxed_vector_size v1 v2 =
  (* For the purposes of layouts/native representations,
     vectors of the same width are equal. *)
  match v1, v2 with
  | Unboxed_vec128, Unboxed_vec128 -> true
  | Unboxed_vec256, Unboxed_vec256 -> true
  | Unboxed_vec512, Unboxed_vec512 -> true
  | (Unboxed_vec128 | Unboxed_vec256 | Unboxed_vec512), _ -> false

let equal_native_repr nr1 nr2 =
  match nr1, nr2 with
  | Repr_poly, Repr_poly -> true
  | Repr_poly, (Unboxed_float _ | Unboxed_or_untagged_integer _
               | Unboxed_vector _ | Same_as_ocaml_repr _
               | Unpacked_product _)
  | (Unboxed_float _ | Unboxed_or_untagged_integer _
    | Unboxed_vector _ | Same_as_ocaml_repr _
    | Unpacked_product _), Repr_poly
    -> false
  | Same_as_ocaml_repr s1, Same_as_ocaml_repr s2 ->
    Jkind_types.Sort.Const.equal s1 s2
  | Same_as_ocaml_repr _,
    (Unboxed_float _ | Unboxed_or_untagged_integer _ |
     Unboxed_vector _ | Unpacked_product _) -> false
  | Unboxed_float f1, Unboxed_float f2 -> equal_boxed_float f1 f2
  | Unboxed_float _,
    (Same_as_ocaml_repr _ | Unboxed_or_untagged_integer _ |
     Unboxed_vector _ | Unpacked_product _) -> false
  | Unboxed_vector vi1, Unboxed_vector vi2 ->
    equal_unboxed_vector_size (unboxed_vector vi1) (unboxed_vector vi2)
  | Unboxed_vector _,
    (Same_as_ocaml_repr _ | Unboxed_float _ |
     Unboxed_or_untagged_integer _ | Unpacked_product _) -> false
  | Unboxed_or_untagged_integer bi1, Unboxed_or_untagged_integer bi2 ->
    equal_unboxed_or_untagged_integer bi1 bi2
  | Unboxed_or_untagged_integer _,
    (Same_as_ocaml_repr _ | Unboxed_float _ |
     Unboxed_vector _ | Unpacked_product _) -> false
  | Unpacked_product s1, Unpacked_product s2 ->
    Jkind_types.Sort.Const.equal s1 s2
  | Unpacked_product _,
    (Same_as_ocaml_repr _ | Unboxed_float _ |
     Unboxed_vector _ | Unboxed_or_untagged_integer _) -> false

let equal_effects ef1 ef2 =
  match ef1, ef2 with
  | No_effects, No_effects -> true
  | No_effects, (Only_generative_effects | Arbitrary_effects) -> false
  | Only_generative_effects, Only_generative_effects -> true
  | Only_generative_effects, (No_effects | Arbitrary_effects) -> false
  | Arbitrary_effects, Arbitrary_effects -> true
  | Arbitrary_effects, (No_effects | Only_generative_effects) -> false

let equal_coeffects cf1 cf2 =
  match cf1, cf2 with
  | No_coeffects, No_coeffects -> true
  | No_coeffects, Has_coeffects -> false
  | Has_coeffects, Has_coeffects -> true
  | Has_coeffects, No_coeffects -> false

let native_name_is_external p =
  let nat_name = native_name p in
  nat_name <> "" && nat_name.[0] <> '%'

module Repr_check = struct

  type result =
    | Wrong_arity
    | Wrong_repr of wrong_repr_error list (* non-empty list *)
    | Success

  let args_res_reprs prim =
    (prim.prim_native_repr_args @ [prim.prim_native_repr_res])
    |> List.map snd

  let is expected_repr repr : wrong_repr_error list =
    if equal_native_repr expected_repr repr
    then []
    else [Repr_mismatch]

  let any = fun _ -> []

  let value_or_unboxed_or_untagged = function
    | Same_as_ocaml_repr (Base Scannable)
    | Unboxed_float _ | Unboxed_or_untagged_integer _ | Unboxed_vector _ -> true
    | Same_as_ocaml_repr _ | Repr_poly | Unpacked_product _ -> false

  let sort_is_product : Jkind_types.Sort.Const.t -> bool = function
    | Product _ -> true
    | Base _ -> false
    | Univar _ -> Misc.fatal_error "sort_is_product: univar"
    | Genvar _ -> Misc.fatal_error "sort_is_product: genvar"

  let c_stub_arg_errors = function
    | Same_as_ocaml_repr s ->
      if sort_is_product s then [Product_arg] else []
    | Unboxed_float _ | Unboxed_or_untagged_integer _ | Unboxed_vector _
    | Unpacked_product _ | Repr_poly -> []

  let c_stub_return_errors = function
    | Same_as_ocaml_repr (Base _)
    | Unboxed_float _ | Unboxed_or_untagged_integer _ | Unboxed_vector _
    | Repr_poly -> []
    | Unpacked_product _ -> [Unpacked_product_return]
    | Same_as_ocaml_repr (Product [s1; s2]) ->
      if (sort_is_product s1) ||
         (sort_is_product s2)
      then [Product_return]
      else []
    | Same_as_ocaml_repr (Product _) -> [Product_return]
    | Same_as_ocaml_repr (Univar _) ->
      Misc.fatal_error "c_stub_return_errors: univar"
    | Same_as_ocaml_repr (Genvar _) ->
      Misc.fatal_error "c_stub_return_errors: genvar"

  (* [checks = [check_arg1; check_arg2; ..; check_argn; check_ret]], where for
     each check, [check (repr : native_repr)] returns a [wrong_repr_error list].
     If the returned list is empty, then the check succeeded. For example,
     [check [any; is (Same_as_ocaml_repr C.scannable)] prim_desc] checks that
     [prim_desc] accepts a single argument of any layout and returns a
     scannable value. *)
  let check checks prim =
    let reprs = args_res_reprs prim in
    if List.length reprs <> List.length checks
    then Wrong_arity
    else begin
      let repr_errors =
        List.concat (List.map2 (fun f x -> f x) checks reprs)
      in
      if List.is_empty repr_errors
      then Success
      else Wrong_repr repr_errors
    end

  let exactly required =
    check (List.map is required)

  let same_arg_res_repr_with_arity arity prim =
    let repr = snd prim.prim_native_repr_res in
    exactly
      (List.init (arity+1) (fun _ -> repr))
      prim

  let no_non_value_repr prim =
    let arity = List.length prim.prim_native_repr_args in
    check
      (List.init (arity+1)
         (fun _ repr ->
            if value_or_unboxed_or_untagged repr
            then []
            else [Expected_value_prim]))
      prim

  let check_c_stub prim =
    (* C externals are allowed to return a tuple, but may not take products as
       arguments or return products with more than two elements. *)
    let arity = List.length prim.prim_native_repr_args in
    let checks =
      (List.init arity (fun _ -> c_stub_arg_errors))
      @ [c_stub_return_errors]
    in
    check checks prim
end

(* Note: [any] here is not the same as jkind [any]. It means we allow any
   [native_repr] for the corresponding argument or return.  It's [Typedecl]'s
   responsibility to check that types in externals are representable or marked
   with [@layout_poly] (see [make_native_repr] and the note above
   [error_if_containing_unexpected_jkind]).  Here we have more speicific checks
   for individual primitives. *)
let prim_has_valid_reprs ~loc prim =
  let open Repr_check in

  let module C = Jkind_types.Sort.Const in

  let check =
    (* Corresponds to [indexing_primitives] in [translprim.ml]. *)
    let stringlike_indexing_primitives =
      let widths : (_ * _ * Jkind_types.Sort.Const.t) list =
        [
          ("8", "", C.scannable);
          ("i8", "", C.scannable);
          ("16", "", C.scannable);
          ("i16", "", C.scannable);
          ("32", "", C.scannable);
          ("f32", "", C.scannable);
          ("64", "", C.scannable);
          ("a128", "", C.scannable);
          ("u128", "", C.scannable);
          ("a256", "", C.scannable);
          ("u256", "", C.scannable);
          ("a512", "", C.scannable);
          ("u512", "", C.scannable);
          ("8", "#", C.bits8);
          ("i8", "#", C.bits8);
          ("16", "#", C.bits16);
          ("i16", "#", C.bits16);
          ("32", "#", C.bits32);
          ("f32", "#", C.float32);
          ("64", "#", C.bits64);
          ("a128", "#", C.vec128);
          ("u128", "#", C.vec128);
          ("a256", "#", C.vec256);
          ("u256", "#", C.vec256);
          ("a512", "#", C.vec512);
          ("u512", "#", C.vec512);
        ]
      in
      let indices : (_ * Jkind_types.Sort.Const.t) list =
        [
          ("", C.scannable);
          ("_indexed_by_nativeint#", C.word);
          ("_indexed_by_int8#", C.bits8);
          ("_indexed_by_int16#", C.bits16);
          ("_indexed_by_int32#", C.bits32);
          ("_indexed_by_int64#", C.bits64);
        ]
      in
      let combiners =
        [
          ( Printf.sprintf "%%caml_%s_get%s%s%s%s",
            fun index_kind width_kind ->
              [
                Same_as_ocaml_repr C.scannable;
                Same_as_ocaml_repr index_kind;
                Same_as_ocaml_repr width_kind;
              ] );
          ( Printf.sprintf "%%caml_%s_set%s%s%s%s",
            fun index_kind width_kind ->
              [
                Same_as_ocaml_repr C.scannable;
                Same_as_ocaml_repr index_kind;
                Same_as_ocaml_repr width_kind;
                Same_as_ocaml_repr C.scannable;
              ] );
        ]
      in
      (let ( let* ) x f = List.concat_map f x in
       let* container = [ "bigstring"; "bytes"; "string" ] in
       let* safe_sigil = [ ""; "u" ] in
       let* index_sigil, index_kind = indices in
       let* width_sigil, unboxed_sigil, width_kind = widths in
       let* combine_string, combine_repr = combiners in
       let string =
         combine_string container width_sigil safe_sigil unboxed_sigil
           index_sigil
       in
       let reprs = combine_repr index_kind width_kind in
       [ (string, reprs) ])
      |> List.to_seq
      |> fun seq -> String.Map.add_seq seq String.Map.empty
    in
    (* Corresponds to [array_vec_primitives] in [translprim.ml]. *)
    let vector_array_indexing_primitives =
      let vector_sizes = [
        ("128", "", C.scannable);
        ("128", "#", C.vec128);
        ("256", "", C.scannable);
        ("256", "#", C.vec256);
        ("512", "", C.scannable);
        ("512", "#", C.vec512);
      ] in
      let array_types = [
        "float_array";
        "floatarray";
        "unboxed_float_array";
        "unboxed_float32_array";
        "int_array";
        "unboxed_int64_array";
        "unboxed_int32_array";
        "untagged_int16_array";
        "untagged_int8_array";
        "unboxed_nativeint_array";
      ] in
      let safe_sigils = [""; "u"] in
      let indices = [
        ("", C.scannable);
        ("_indexed_by_nativeint#", C.word);
        ("_indexed_by_int8#", C.bits8);
        ("_indexed_by_int16#", C.bits16);
        ("_indexed_by_int32#", C.bits32);
        ("_indexed_by_int64#", C.bits64);
      ] in
      let combiners =
        [
          ( Printf.sprintf "%%caml_%s_get%s%s%s%s",
            fun index_kind vector_kind ->
              [
                Same_as_ocaml_repr C.scannable;
                Same_as_ocaml_repr index_kind;
                Same_as_ocaml_repr vector_kind;
              ] );
          ( Printf.sprintf "%%caml_%s_set%s%s%s%s",
            fun index_kind vector_kind ->
              [
                Same_as_ocaml_repr C.scannable;
                Same_as_ocaml_repr index_kind;
                Same_as_ocaml_repr vector_kind;
                Same_as_ocaml_repr C.scannable;
              ] );
        ]
      in
      (let ( let* ) x f = List.concat_map f x in
       let* array_type = array_types in
       let* safe_sigil = safe_sigils in
       let* size_str, unboxed_sigil, vector_kind = vector_sizes in
       let* index_suffix, index_kind = indices in
       let* combine_string, combine_repr = combiners in
       let string =
         combine_string array_type size_str safe_sigil unboxed_sigil
           index_suffix
       in
       let reprs = combine_repr index_kind vector_kind in
       [ (string, reprs) ])
      |> List.to_seq
      |> fun seq -> String.Map.add_seq seq String.Map.empty
    in
    match prim.prim_name with
    | "%identity"
    | "%opaque"
    | "%obj_magic" ->
      same_arg_res_repr_with_arity 1

    | "%ignore" ->
      check [any; is (Same_as_ocaml_repr C.scannable)]
    | "%revapply" ->
      check [any; is (Same_as_ocaml_repr C.scannable); any]
    | "%apply" ->
      check [is (Same_as_ocaml_repr C.scannable); any; any]

    (* This doesn't prevent

       {|
          external get : float# array -> int -> int32# =
            "%array_safe_get"
       |}

       but the same is true for types:

       {|
          external get : float array -> int -> int32 =
            "%array_safe_get"
       |}
    *)
    | "%array_safe_get" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.scannable);
        any]
    | "%array_safe_set" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.scannable);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_unsafe_get" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.scannable);
        any]
    | "%array_unsafe_set" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.scannable);
        any;
        is (Same_as_ocaml_repr C.scannable)]

    | "%array_safe_get_indexed_by_int64#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits64);
        any]
    | "%array_safe_set_indexed_by_int64#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits64);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_unsafe_get_indexed_by_int64#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits64);
        any]
    | "%array_unsafe_set_indexed_by_int64#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits64);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_safe_get_indexed_by_int32#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits32);
        any]
    | "%array_safe_set_indexed_by_int32#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits32);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_unsafe_get_indexed_by_int32#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits32);
        any]
    | "%array_unsafe_set_indexed_by_int32#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits32);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_safe_get_indexed_by_int16#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits16);
        any]
    | "%array_safe_set_indexed_by_int16#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits16);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_unsafe_get_indexed_by_int16#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits16);
        any]
    | "%array_unsafe_set_indexed_by_int16#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits16);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_safe_get_indexed_by_int8#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits8);
        any]
    | "%array_safe_set_indexed_by_int8#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits8);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_unsafe_get_indexed_by_int8#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits8);
        any]
    | "%array_unsafe_set_indexed_by_int8#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits8);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_safe_get_indexed_by_nativeint#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.word);
        any]
    | "%array_safe_set_indexed_by_nativeint#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.word);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%array_unsafe_get_indexed_by_nativeint#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.word);
        any]
    | "%array_unsafe_set_indexed_by_nativeint#" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.word);
        any;
        is (Same_as_ocaml_repr C.scannable)]
    | "%makearray_dynamic" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        any;
        is (Same_as_ocaml_repr C.scannable);
      ]
    | "%makearray_dynamic_uninit" ->
      (* Restrictions on this primitive are checked in [Translprim] *)
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.scannable);
      ]
    | "%array_element_size_in_bytes" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.scannable);
      ]
    | "%peek" | "%poke" ->
      (* Arities and layouts of these primitives are checked in [Translprim] *)
      fun _ -> Success
    | "%get_idx" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits64);
        any
      ]
    | "%get_idx_imm" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits64);
        any
      ]
    | "%set_idx" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits64);
        any;
        is (Same_as_ocaml_repr C.scannable);
      ]
    | "%unsafe_array_idx" ->
      check [
        is (Same_as_ocaml_repr C.scannable);
        is (Same_as_ocaml_repr C.bits64);
      ]
    | "%unsafe_array_idx_indexed_by_int8#" ->
      check [
        is (Same_as_ocaml_repr C.bits8);
        is (Same_as_ocaml_repr C.bits64);
      ]
    | "%unsafe_array_idx_indexed_by_int16#" ->
      check [
        is (Same_as_ocaml_repr C.bits16);
        is (Same_as_ocaml_repr C.bits64);
      ]
    | "%unsafe_array_idx_indexed_by_int32#" ->
      check [
        is (Same_as_ocaml_repr C.bits32);
        is (Same_as_ocaml_repr C.bits64);
      ]
    | "%unsafe_array_idx_indexed_by_int64#" ->
      check [
        is (Same_as_ocaml_repr C.bits64);
        is (Same_as_ocaml_repr C.bits64);
      ]
    | "%unsafe_array_idx_indexed_by_nativeint#" ->
      check [
        is (Same_as_ocaml_repr C.word);
        is (Same_as_ocaml_repr C.bits64);
      ]
    | "%unsafe_get_ptr" ->
      check [
        is (Same_as_ocaml_repr (C.Product [C.scannable; C.bits64]));
        any
      ]
    | "%unsafe_get_ptr_imm" ->
      check [
        is (Same_as_ocaml_repr (C.Product [C.scannable; C.bits64]));
        any
      ]
    | "%unsafe_set_ptr" ->
      check [
        is (Same_as_ocaml_repr (C.Product [C.scannable; C.bits64]));
        any;
        is (Same_as_ocaml_repr C.scannable);
      ]
    | "%box_float" ->
      exactly [Same_as_ocaml_repr C.float64; Same_as_ocaml_repr C.scannable]
    | "%unbox_float" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.float64]
    | "%box_float32" ->
      exactly [Same_as_ocaml_repr C.float32; Same_as_ocaml_repr C.scannable]
    | "%unbox_float32" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.float32]
    | "%box_nativeint" ->
      exactly [Same_as_ocaml_repr C.word; Same_as_ocaml_repr C.scannable]
    | "%unbox_nativeint" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.word]
    | "%domain_index" ->
      exactly
        [ Same_as_ocaml_repr C.scannable;
          Same_as_ocaml_repr C.untagged_immediate ]
    | "%tag_int" ->
      exactly
        [ Same_as_ocaml_repr C.untagged_immediate;
          Same_as_ocaml_repr C.scannable ]
    | "%untag_int" ->
      exactly
        [ Same_as_ocaml_repr C.scannable;
          Same_as_ocaml_repr C.untagged_immediate ]
    | "%tag_int8" ->
      exactly [Same_as_ocaml_repr C.bits8; Same_as_ocaml_repr C.scannable]
    | "%untag_int8" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.bits8]
    | "%tag_int16" ->
      exactly [Same_as_ocaml_repr C.bits16; Same_as_ocaml_repr C.scannable]
    | "%untag_int16" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.bits16]
    | "%box_int32" ->
      exactly [Same_as_ocaml_repr C.bits32; Same_as_ocaml_repr C.scannable]
    | "%unbox_int32" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.bits32]
    | "%box_int64" ->
      exactly [Same_as_ocaml_repr C.bits64; Same_as_ocaml_repr C.scannable]
    | "%unbox_int64" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.bits64]
    | "%unbox_unit" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.void]
    | "%box_vec128" ->
      exactly [Same_as_ocaml_repr C.vec128; Same_as_ocaml_repr C.scannable]
    | "%unbox_vec128" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.vec128]
    | "%box_vec256" ->
      exactly [Same_as_ocaml_repr C.vec256; Same_as_ocaml_repr C.scannable]
    | "%unbox_vec256" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.vec256]
    | "%join_vec256" ->
      exactly [Same_as_ocaml_repr C.vec128; Same_as_ocaml_repr C.vec128;
               Same_as_ocaml_repr C.vec256]
    | "%split_vec256" ->
      exactly [Same_as_ocaml_repr C.vec256;
               Same_as_ocaml_repr (Product [C.vec128; C.vec128])]
    | "%box_vec512" ->
      exactly [Same_as_ocaml_repr C.vec512; Same_as_ocaml_repr C.scannable]
    | "%unbox_vec512" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.vec512]

    | "%reinterpret_tagged_int63_as_unboxed_int64" ->
      exactly [Same_as_ocaml_repr C.scannable; Same_as_ocaml_repr C.bits64]
    | "%reinterpret_unboxed_int64_as_tagged_int63" ->
      exactly [Same_as_ocaml_repr C.bits64; Same_as_ocaml_repr C.scannable]

    | name -> (
        match String.Map.find_opt name stringlike_indexing_primitives with
        | Some reprs -> exactly reprs
        | None ->
            match String.Map.find_opt name vector_array_indexing_primitives with
            | Some reprs -> exactly reprs
            | None ->
              let module I = Scalar.Operation in
              match I.With_percent_prefix.of_string name with
              | intrinsic ->
                exactly
                  (List.map (fun sort -> Same_as_ocaml_repr sort)
                     (I.sort intrinsic))
              | exception Not_found ->
                if is_builtin_prim_name name then no_non_value_repr
                  (* These can probably support non-value reprs if the
                     need arises:
                     {|
                       | "%send"
                       | "%sendself"
                       | "%sendcache"
                     |}
                  *)
                else check_c_stub)
  in
  match check prim with
  | Success -> ()
  | Wrong_arity ->
    (* There's already an arity check in translprim that catches this. We will
       defer to that logic.  We are only checking the arity of some built-in
       primitives here but not all, and it would be weird to raise different
       errors dependent on the [prim_name]. *)
    ()
  | Wrong_repr errors ->
    raise (Error (loc,
            Invalid_native_repr_for_primitive
              { prim_name = prim.prim_name; errors }))

let prim_can_contain_layout_any prim =
  match prim.prim_name with
  | "%array_length"
  | "%array_blit"
  | "%array_safe_get"
  | "%array_safe_set"
  | "%array_unsafe_get"
  | "%array_unsafe_set"
  | "%array_safe_get_indexed_by_int64#"
  | "%array_safe_set_indexed_by_int64#"
  | "%array_unsafe_get_indexed_by_int64#"
  | "%array_unsafe_set_indexed_by_int64#"
  | "%array_safe_get_indexed_by_int32#"
  | "%array_safe_set_indexed_by_int32#"
  | "%array_unsafe_get_indexed_by_int32#"
  | "%array_unsafe_set_indexed_by_int32#"
  | "%array_safe_get_indexed_by_nativeint#"
  | "%array_safe_set_indexed_by_nativeint#"
  | "%array_unsafe_get_indexed_by_nativeint#"
  | "%array_unsafe_set_indexed_by_nativeint#" -> false
  | _ -> true

module Style = Misc.Style

let report_error ppf err =
  match err with
  | Old_style_float_with_native_repr_attribute ->
    Format_doc.fprintf ppf "Cannot use %a in conjunction with %a/%a."
      Style.inline_code "float"
      Style.inline_code "[@unboxed]"
      Style.inline_code  "[@untagged]"
  | Old_style_float_with_non_value ->
    Format_doc.fprintf ppf "Cannot use %a in conjunction with \
                        types of non-value layouts."
      Style.inline_code "float"
  | Old_style_noalloc_with_noalloc_attribute ->
    Format_doc.fprintf ppf "Cannot use %a in conjunction with %a."
      Style.inline_code "noalloc"
      Style.inline_code "[@@noalloc]"
  | No_native_primitive_with_repr_attribute ->
    Format_doc.fprintf ppf
      "@[The native code version of the primitive is mandatory@ \
      when attributes %a or %a are present.@]"
      Style.inline_code "[@untagged]"
      Style.inline_code "[@unboxed]"
  | No_native_primitive_with_non_value ->
    Format_doc.fprintf ppf
      "@[The native code version of the primitive is mandatory@ \
       for types with non-value layouts.@]"
  | Inconsistent_attributes_for_effects ->
    Format_doc.fprintf ppf "At most one of %a and %a can be specified."
      Style.inline_code "[@no_effects]"
      Style.inline_code "[@only_generative_effects]"
  | Inconsistent_noalloc_attributes_for_effects ->
    Format_doc.fprintf ppf "Cannot use %a in conjunction with %a."
      Style.inline_code "[@@no_generative_effects]"
      Style.inline_code "[@@noalloc]"
  | Invalid_representation_polymorphic_attribute ->
    Format_doc.fprintf ppf "Attribute %a can only be used \
                        on built-in primitives."
      Style.inline_code "[@layout_poly]"
  | Invalid_native_repr_for_primitive { prim_name; errors } ->
    Format_doc.fprintf ppf
      "The primitive [%s] is used in an invalid declaration.@ \
       The declaration contains argument/return types with the@ \
       wrong layout."
      prim_name;
    let errors = List.sort_uniq compare_wrong_repr_error errors in
    List.iter
      (function
      | Product_arg ->
        Format_doc.fprintf ppf
          "@.@{<hint>Hint@}: @[<v>\
           Types with product layouts in C stub arguments require the@ \
           %a attribute.@]"
          Style.inline_code "[@unpacked]"
      | Expected_value_prim ->
        Format_doc.fprintf ppf
          "@.@{<hint>Hint@}: @[<v>\
           This was expected to be a value-only primitive. You might've@ \
           misspelled the primitive name.@]"
      | Product_return ->
        Format_doc.fprintf ppf
          "@.@{<hint>Hint@}: @[<v>\
           Unboxed products in C stub returns must be a pair of non-products.@]"
      | Unpacked_product_return ->
        Format_doc.fprintf ppf
          "@.@{<hint>Hint@}: @[<v>\
           The %a attribute is not allowed on C stub returns.@]"
          Style.inline_code "[@unpacked]"
      | Repr_mismatch -> ()
        (* The error message already says "wrong layout", so a hint here would
           be redundant. *))
      errors

let () =
  Location.register_error_of_exn
    (function
      | Error (loc, err) ->
        Some (Location.error_of_printer ~loc report_error err)
      | _ ->
        None
    )