Fix up kind checking for peek/poke

lambda/lambda.ml

@@ -325,8 +325,8 @@ type primitive =
   | Parray_to_iarray
   | Parray_of_iarray
   | Pget_header of locality_mode
-  | Ppeek of peek_or_poke
-  | Ppoke of peek_or_poke
+  | Ppeek of peek_or_poke_kind
+  | Ppoke of peek_or_poke_kind
   (* Fetching domain-local state *)
   | Pdls_get
   (* Poll for runtime actions *)
@@ -477,7 +477,7 @@ and unboxed_integer = boxed_integer
 
 and unboxed_vector = boxed_vector
 
-and peek_or_poke =
+and peek_or_poke_kind =
   | Ppp_tagged_immediate
   | Ppp_unboxed_float32
   | Ppp_unboxed_float

lambda/lambda.mli

@@ -331,8 +331,8 @@ type primitive =
   | Parray_of_iarray (* Unsafely reinterpret an immutable array as a mutable
                         one; O(1) *)
   | Pget_header of locality_mode
-  | Ppeek of peek_or_poke
-  | Ppoke of peek_or_poke
+  | Ppeek of peek_or_poke_kind
+  | Ppoke of peek_or_poke_kind
   (* Get the header of a block. This primitive is invalid if provided with an
      immediate value.
      Note: The GC color bits in the header are not reliable except for checking
@@ -508,7 +508,7 @@ and unboxed_integer = boxed_integer
 
 and unboxed_vector = boxed_vector
 
-and peek_or_poke =
+and peek_or_poke_kind =
   | Ppp_tagged_immediate
   | Ppp_unboxed_float32
   | Ppp_unboxed_float

lambda/printlambda.ml

@@ -475,7 +475,7 @@ let field_read_semantics ppf sem =
   | Reads_agree -> ()
   | Reads_vary -> fprintf ppf "_mut"
 
-let peek_or_poke ppf (pp : peek_or_poke) =
+let peek_or_poke ppf (pp : peek_or_poke_kind) =
   match pp with
   | Ppp_tagged_immediate -> fprintf ppf "tagged_immediate"
   | Ppp_unboxed_float32 -> fprintf ppf "unboxed_float32"

lambda/translprim.ml

@@ -110,8 +110,8 @@ type prim =
   | Identity
   | Apply of Lambda.region_close * Lambda.layout
   | Revapply of Lambda.region_close * Lambda.layout
-  | Peek of Lambda.peek_or_poke option
-  | Poke of Lambda.peek_or_poke option
+  | Peek of Lambda.peek_or_poke_kind option
+  | Poke of Lambda.peek_or_poke_kind option
     (* For [Peek] and [Poke] the [option] is [None] until the primitive
        specialization code (below) has been run. *)
   | Unsupported of Lambda.primitive
@@ -1162,31 +1162,6 @@ let glb_array_set_type loc t1 t2 =
   (* Pfloatarray is a minimum *)
   | Pfloatarray_set, Pfloatarray -> Pfloatarray_set
 
-let peek_or_poke_layout_from_type error_loc env ty
-      : Lambda.peek_or_poke option =
-  match
-    (* XXX mshinwell: fix [why] *)
-    Ctype.type_sort ~why:Layout_poly_in_external ~fixed:true env ty
-  with
-  | Error _ -> None
-  | Ok sort ->
-    let layout = Typeopt.layout env error_loc sort ty in
-    match layout with
-    | Punboxed_float Pfloat32 -> Some Ppp_unboxed_float32
-    | Punboxed_float Pfloat64 -> Some Ppp_unboxed_float
-    | Punboxed_int Pint32 -> Some Ppp_unboxed_int32
-    | Punboxed_int Pint64 -> Some Ppp_unboxed_int64
-    | Punboxed_int Pnativeint -> Some Ppp_unboxed_nativeint
-    | Pvalue { raw_kind = Pintval ; _ } -> Some Ppp_tagged_immediate
-    | Ptop
-    | Pvalue _
-    | Punboxed_vector _
-    | Punboxed_product _
-    | Pbottom ->
-      (* XXX should be a proper error *)
-      Misc.fatal_errorf "Bad layout for %%peek:@ %a"
-        Printlambda.layout layout
-
 (* Specialize a primitive from available type information. *)
 (* CR layouts v7: This function had a loc argument added just to support the void
    check error message.  Take it out when we remove that. *)
@@ -1348,14 +1323,12 @@ let specialize_primitive env loc ty ~has_constant_constructor prim =
          (if not, it seems like a layout of value is inferred)
          let test_read32_s (x : int32# t) : int32# = read_s x
       *)
-      match peek_or_poke_layout_from_type (to_location loc) env result_ty with
-      | None -> None
-      | Some contents_layout -> Some (Peek (Some contents_layout))
+      let kind = Typeopt.peek_or_poke_kind env (to_location loc) result_ty in
+      Some (Peek (Some kind))
   )
   | Poke _, _ptr_ty :: new_value_ty :: _ -> (
-    match peek_or_poke_layout_from_type (to_location loc) env new_value_ty with
-    | None -> None
-    | Some contents_layout -> Some (Poke (Some contents_layout))
+    let kind = Typeopt.peek_or_poke_kind env (to_location loc) new_value_ty in
+    Some (Poke (Some kind))
   )
   | _ -> None
 

typing/jkind.ml

@@ -1463,6 +1463,9 @@ module Format_history = struct
         "it's the layout polymorphic type in an external declaration@ \
          ([@@layout_poly] forces all variables of layout 'any' to be@ \
          representable at call sites)"
+    | Peeked_or_poked ->
+      fprintf ppf
+        "it's the type of something used with the peek/poke primitives"
 
   let format_concrete_legacy_creation_reason ppf :
       History.concrete_legacy_creation_reason -> unit = function
@@ -1957,6 +1960,7 @@ module Debug_printers = struct
     | Optional_arg_default -> fprintf ppf "Optional_arg_default"
     | Layout_poly_in_external -> fprintf ppf "Layout_poly_in_external"
     | Unboxed_tuple_element -> fprintf ppf "Unboxed_tuple_element"
+    | Peeked_or_poked -> fprintf ppf "Peeked_or_poked"
 
   let concrete_legacy_creation_reason ppf :
       History.concrete_legacy_creation_reason -> unit = function

typing/jkind_intf.ml

@@ -187,6 +187,7 @@ module History = struct
     | Optional_arg_default
     | Layout_poly_in_external
     | Unboxed_tuple_element
+    | Peeked_or_poked
 
   (* For sort variables that are in the "legacy" position
      on the jkind lattice, defaulting exactly to [value]. *)

typing/typeopt.ml

@@ -20,6 +20,21 @@ open Types
 open Typedtree
 open Lambda
 
+(* [classification]s are used for two things: things in arrays, and things in
+   lazys. In the former case, we need detailed information about unboxed
+   products and in the latter it would be wasteful to compute that information,
+   so this type is polymorphic in what it remembers about products. *)
+type 'a classification =
+  | Int   (* any immediate type *)
+  | Float
+  | Unboxed_float of unboxed_float
+  | Unboxed_int of unboxed_integer
+  | Unboxed_vector of unboxed_vector
+  | Lazy
+  | Addr  (* any value except a float or a lazy *)
+  | Any
+  | Product of 'a
+
 type error =
     Non_value_layout of type_expr * Jkind.Violation.t option
   | Non_value_sort of Jkind.Sort.t * type_expr
@@ -34,6 +49,7 @@ type error =
   | Unsupported_vector_in_product_array
   | Mixed_product_array of Jkind.Sort.Const.t
   | Product_iarrays_unsupported
+  | Unsupported_peek_or_poke_type of unit classification * type_expr
 
 exception Error of Location.t * error
 
@@ -116,21 +132,6 @@ let type_legacy_sort ~why env loc ty =
   | Ok sort -> sort
   | Error err -> raise (Error (loc, Not_a_sort (ty, err)))
 
-(* [classification]s are used for two things: things in arrays, and things in
-   lazys. In the former case, we need detailed information about unboxed
-   products and in the latter it would be wasteful to compute that information,
-   so this type is polymorphic in what it remembers about products. *)
-type 'a classification =
-  | Int   (* any immediate type *)
-  | Float
-  | Unboxed_float of unboxed_float
-  | Unboxed_int of unboxed_integer
-  | Unboxed_vector of unboxed_vector
-  | Lazy
-  | Addr  (* any value except a float or a lazy *)
-  | Any
-  | Product of 'a
-
 (* Classify a ty into a [classification]. Looks through synonyms, using
    [scrape_ty].  Returning [Any] is safe, though may skip some optimizations.
    See comment on [classification] above to understand [classify_product]. *)
@@ -279,6 +280,23 @@ let array_pattern_kind pat elt_sort =
     ~elt_sort:(Some elt_sort)
     pat.pat_env pat.pat_loc pat.pat_type
 
+let peek_or_poke_kind env loc ty =
+  let elt_sort =
+    match Ctype.type_sort ~why:Peeked_or_poked ~fixed:true env ty with
+    | Ok s -> s
+    | Error err -> raise (Error (loc, Not_a_sort (ty, err)))
+  in
+  let classify_product _ty _sorts = () in
+  match classify ~classify_product env loc ty elt_sort with
+  | Any | Float | Addr | Lazy | Unboxed_vector _ | Product _ as c ->
+    raise (Error (loc, Unsupported_peek_or_poke_type (c, ty)))
+  | Int -> Ppp_tagged_immediate
+  | Unboxed_int Pnativeint -> Ppp_unboxed_nativeint
+  | Unboxed_int Pint32 -> Ppp_unboxed_int32
+  | Unboxed_int Pint64 -> Ppp_unboxed_int64
+  | Unboxed_float Pfloat32 -> Ppp_unboxed_float32
+  | Unboxed_float Pfloat64 -> Ppp_unboxed_float
+
 let bigarray_decode_type env ty tbl dfl =
   match scrape env ty with
   | Tconstr(Pdot(Pident mod_id, type_name), [], _)
@@ -1047,6 +1065,18 @@ let report_error ppf = function
   | Product_iarrays_unsupported ->
       fprintf ppf
         "Immutable arrays of unboxed products are not yet supported."
+  | Unsupported_peek_or_poke_type (c, ty) ->
+      let explanation =
+        match c with
+        | Int | Unboxed_float _ | Unboxed_int _ -> assert false
+        | Unboxed_vector _ -> "Vectors are not supported."
+        | Float | Lazy | Addr | Any ->
+          "They only work with types that can be ignored by the GC."
+        | Product _ -> "Unboxed products are not supported."
+      in
+      fprintf ppf
+        "The peek and poke primitives do not work with type %a.@ %s"
+        Printtyp.type_expr ty explanation
 
 let () =
   Location.register_error_of_exn

typing/typeopt.mli

@@ -36,6 +36,8 @@ val array_kind :
   Typedtree.expression -> Jkind.Sort.t -> Lambda.array_kind
 val array_pattern_kind :
   Typedtree.pattern -> Jkind.Sort.t -> Lambda.array_kind
+val peek_or_poke_kind :
+  Env.t -> Location.t -> Types.type_expr -> Lambda.peek_or_poke_kind
 
 (* If [kind] or [layout] is unknown, attempt to specialize it by examining the
    type parameters of the bigarray. If [kind] or [length] is not unknown, returns