Peek and poke (ocaml-flambda#3309)

Author: mshinwell

bytecomp/bytegen.ml

@@ -199,7 +199,7 @@ let preserve_tailcall_for_prim = function
   | Patomic_exchange | Patomic_compare_exchange
   | Patomic_cas | Patomic_fetch_add | Patomic_load _
   | Pdls_get | Preinterpret_tagged_int63_as_unboxed_int64
-  | Preinterpret_unboxed_int64_as_tagged_int63 | Ppoll ->
+  | Preinterpret_unboxed_int64_as_tagged_int63 | Ppoll | Ppeek _ | Ppoke _ ->
       false
 
 (* Add a Kpop N instruction in front of a continuation *)
@@ -737,6 +737,8 @@ let comp_primitive stack_info p sz args =
   | Punbox_float _ | Pbox_float (_, _) | Punbox_int _ | Pbox_int _
     ->
       fatal_error "Bytegen.comp_primitive"
+  | Ppeek _ | Ppoke _ ->
+      fatal_error "Bytegen.comp_primitive: Ppeek/Ppoke not supported in bytecode"
 
 let is_immed n = immed_min <= n && n <= immed_max
 

lambda/lambda.ml

@@ -334,6 +334,8 @@ type primitive =
   | Parray_to_iarray
   | Parray_of_iarray
   | Pget_header of locality_mode
+  | Ppeek of peek_or_poke
+  | Ppoke of peek_or_poke
   (* Fetching domain-local state *)
   | Pdls_get
   (* Poll for runtime actions *)
@@ -490,6 +492,14 @@ and boxed_integer = Primitive.boxed_integer =
 and boxed_vector = Primitive.boxed_vector =
   | Boxed_vec128
 
+and peek_or_poke =
+  | Ppp_tagged_immediate
+  | Ppp_unboxed_float32
+  | Ppp_unboxed_float
+  | Ppp_unboxed_int32
+  | Ppp_unboxed_int64
+  | Ppp_unboxed_nativeint
+
 and bigarray_kind =
     Pbigarray_unknown
   | Pbigarray_float16
@@ -1941,7 +1951,8 @@ let primitive_may_allocate : primitive -> locality_mode option = function
   | Patomic_fetch_add
   | Pdls_get
   | Preinterpret_unboxed_int64_as_tagged_int63
-  | Parray_element_size_in_bytes _ -> None
+  | Parray_element_size_in_bytes _
+  | Ppeek _ | Ppoke _ -> None
   | Preinterpret_tagged_int63_as_unboxed_int64 ->
     if !Clflags.native_code then None
     else
@@ -2107,7 +2118,7 @@ let primitive_can_raise prim =
   | Prunstack | Pperform | Presume | Preperform -> true (* XXX! *)
   | Pdls_get | Ppoll | Preinterpret_tagged_int63_as_unboxed_int64
   | Preinterpret_unboxed_int64_as_tagged_int63
-  | Parray_element_size_in_bytes _ ->
+  | Parray_element_size_in_bytes _ | Ppeek _ | Ppoke _ ->
     false
 
 let constant_layout: constant -> layout = function
@@ -2342,6 +2353,16 @@ let primitive_result_layout (p : primitive) =
   | Ppoll -> layout_unit
   | Preinterpret_tagged_int63_as_unboxed_int64 -> layout_unboxed_int64
   | Preinterpret_unboxed_int64_as_tagged_int63 -> layout_int
+  | Ppeek layout -> (
+      match layout with
+      | Ppp_tagged_immediate -> layout_int
+      | Ppp_unboxed_float32 -> layout_unboxed_float Unboxed_float32
+      | Ppp_unboxed_float -> layout_unboxed_float Unboxed_float64
+      | Ppp_unboxed_int32 -> layout_unboxed_int32
+      | Ppp_unboxed_int64 -> layout_unboxed_int64
+      | Ppp_unboxed_nativeint -> layout_unboxed_nativeint
+    )
+  | Ppoke _ -> layout_unit
 
 let compute_expr_layout free_vars_kind lam =
   let rec compute_expr_layout kinds = function

lambda/lambda.mli

@@ -343,6 +343,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
   (* 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
@@ -524,6 +526,14 @@ and boxed_integer = Primitive.boxed_integer =
 and boxed_vector = Primitive.boxed_vector =
   | Boxed_vec128
 
+and peek_or_poke =
+  | Ppp_tagged_immediate
+  | Ppp_unboxed_float32
+  | Ppp_unboxed_float
+  | Ppp_unboxed_int32
+  | Ppp_unboxed_int64
+  | Ppp_unboxed_nativeint
+
 and bigarray_kind =
     Pbigarray_unknown
   | Pbigarray_float16

lambda/printlambda.ml

@@ -478,6 +478,15 @@ let field_read_semantics ppf sem =
   | Reads_agree -> ()
   | Reads_vary -> fprintf ppf "_mut"
 
+let peek_or_poke ppf (pp : peek_or_poke) =
+  match pp with
+  | Ppp_tagged_immediate -> fprintf ppf "tagged_immediate"
+  | Ppp_unboxed_float32 -> fprintf ppf "unboxed_float32"
+  | Ppp_unboxed_float -> fprintf ppf "unboxed_float"
+  | Ppp_unboxed_int32 -> fprintf ppf "unboxed_int32"
+  | Ppp_unboxed_int64 -> fprintf ppf "unboxed_int64"
+  | Ppp_unboxed_nativeint -> fprintf ppf "unboxed_nativeint"
+
 let primitive ppf = function
   | Pbytes_to_string -> fprintf ppf "bytes_to_string"
   | Pbytes_of_string -> fprintf ppf "bytes_of_string"
@@ -930,6 +939,12 @@ let primitive ppf = function
       fprintf ppf "reinterpret_tagged_int63_as_unboxed_int64"
   | Preinterpret_unboxed_int64_as_tagged_int63 ->
       fprintf ppf "reinterpret_unboxed_int64_as_tagged_int63"
+  | Ppeek layout ->
+      fprintf ppf "(peek@ %a)"
+        peek_or_poke layout
+  | Ppoke layout ->
+      fprintf ppf "(poke@ %a)"
+        peek_or_poke layout
 
 let name_of_primitive = function
   | Pbytes_of_string -> "Pbytes_of_string"
@@ -1107,6 +1122,8 @@ let name_of_primitive = function
       "Preinterpret_tagged_int63_as_unboxed_int64"
   | Preinterpret_unboxed_int64_as_tagged_int63 ->
       "Preinterpret_unboxed_int64_as_tagged_int63"
+  | Ppeek _ -> "Ppeek"
+  | Ppoke _ -> "Ppoke"
 
 let zero_alloc_attribute ppf check =
   match check with

lambda/tmc.ml

@@ -974,6 +974,7 @@ let rec choice ctx t =
     | Pint_as_pointer _
     | Psequand | Psequor
     | Ppoll
+    | Ppeek _ | Ppoke _
       ->
         let primargs = traverse_list ctx primargs in
         Choice.lambda (Lprim (prim, primargs, loc))

lambda/translprim.ml

@@ -28,6 +28,7 @@ module String = Misc.Stdlib.String
 type error =
   | Unknown_builtin_primitive of string
   | Wrong_arity_builtin_primitive of string
+  | Wrong_layout_for_peek_or_poke of string
   | Invalid_floatarray_glb
   | Product_iarrays_unsupported
   | Invalid_array_kind_for_uninitialized_makearray_dynamic
@@ -123,6 +124,10 @@ 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
+    (* For [Peek] and [Poke] the [option] is [None] until the primitive
+       specialization code (below) has been run. *)
   | Unsupported of Lambda.primitive
 
 let units_with_used_primitives = Hashtbl.create 7
@@ -916,6 +921,8 @@ let lookup_primitive loc ~poly_mode ~poly_sort pos p =
       Primitive(Preinterpret_tagged_int63_as_unboxed_int64, 1)
     | "%reinterpret_unboxed_int64_as_tagged_int63" ->
       Primitive(Preinterpret_unboxed_int64_as_tagged_int63, 1)
+    | "%peek" -> Peek None
+    | "%poke" -> Poke None
     | s when String.length s > 0 && s.[0] = '%' ->
       (match String.Map.find_opt s indexing_primitives with
        | Some prim -> prim ~mode
@@ -1185,6 +1192,27 @@ let glb_array_set_type loc t1 t2 =
   (* Pfloatarray is a minimum *)
   | Pfloatarray_set, Pfloatarray -> Pfloatarray_set
 
+let peek_or_poke_layout_from_type ~prim_name error_loc env ty
+      : Lambda.peek_or_poke option =
+  match Ctype.type_sort ~why:Peek_or_poke ~fixed:true env ty with
+  | Error _ -> None
+  | Ok sort ->
+    let sort = Jkind.Sort.default_to_value_and_get sort in
+    let layout = Typeopt.layout env error_loc sort ty in
+    match layout with
+    | Punboxed_float Unboxed_float32 -> Some Ppp_unboxed_float32
+    | Punboxed_float Unboxed_float64 -> Some Ppp_unboxed_float
+    | Punboxed_int Unboxed_int32 -> Some Ppp_unboxed_int32
+    | Punboxed_int Unboxed_int64 -> Some Ppp_unboxed_int64
+    | Punboxed_int Unboxed_nativeint -> Some Ppp_unboxed_nativeint
+    | Pvalue { raw_kind = Pintval ; _ } -> Some Ppp_tagged_immediate
+    | Ptop
+    | Pvalue _
+    | Punboxed_vector _
+    | Punboxed_product _
+    | Pbottom ->
+      raise (Error (error_loc, Wrong_layout_for_peek_or_poke prim_name))
+
 (* 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. *)
@@ -1367,6 +1395,25 @@ let specialize_primitive env loc ty ~has_constant_constructor prim =
     end else begin
       None
     end
+  | Peek _, _ -> (
+    match is_function_type env ty with
+    | None -> None
+    | Some (_p1, result_ty) ->
+      match
+        peek_or_poke_layout_from_type ~prim_name:"peek"
+          (to_location loc) env result_ty
+      with
+      | None -> None
+      | Some contents_layout -> Some (Peek (Some contents_layout))
+  )
+  | Poke _, _ptr_ty :: new_value_ty :: _ -> (
+    match
+      peek_or_poke_layout_from_type ~prim_name:"poke"
+        (to_location loc) env new_value_ty
+    with
+    | None -> None
+    | Some contents_layout -> Some (Poke (Some contents_layout))
+  )
   | _ -> None
 
 let caml_equal =
@@ -1613,6 +1660,12 @@ let lambda_of_prim prim_name prim loc args arg_exps =
         ap_region_close = pos;
         ap_mode = alloc_heap;
       }
+  | Peek None, _ | Poke None, _ ->
+      raise(Error(to_location loc, Wrong_layout_for_peek_or_poke prim_name))
+  | Peek (Some layout), [ptr] ->
+      Lprim (Ppeek layout, [ptr], loc)
+  | Poke (Some layout), [ptr; new_value] ->
+      Lprim (Ppoke layout, [ptr; new_value], loc)
   | Unsupported prim, _ ->
       let exn =
         transl_extension_path loc (Lazy.force Env.initial)
@@ -1631,7 +1684,7 @@ let lambda_of_prim prim_name prim loc args arg_exps =
   | (Raise _ | Raise_with_backtrace
     | Lazy_force _ | Loc _ | Primitive _ | Sys_argv | Comparison _
     | Send _ | Send_self _ | Send_cache _ | Frame_pointers | Identity
-    | Apply _ | Revapply _), _ ->
+    | Apply _ | Revapply _ | Peek _ | Poke _), _ ->
       raise(Error(to_location loc, Wrong_arity_builtin_primitive prim_name))
 
 let check_primitive_arity loc p =
@@ -1663,8 +1716,8 @@ let check_primitive_arity loc p =
     | Send _ | Send_self _ -> p.prim_arity = 2
     | Send_cache _ -> p.prim_arity = 4
     | Frame_pointers -> p.prim_arity = 0
-    | Identity -> p.prim_arity = 1
-    | Apply _ | Revapply _ -> p.prim_arity = 2
+    | Identity | Peek _ -> p.prim_arity = 1
+    | Apply _ | Revapply _ | Poke _ -> p.prim_arity = 2
     | Unsupported _ -> true
   in
   if not ok then raise(Error(loc, Wrong_arity_builtin_primitive p.prim_name))
@@ -1850,7 +1903,7 @@ let lambda_primitive_needs_event_after = function
   | Pintofbint _ | Pctconst _ | Pbswap16 | Pint_as_pointer _ | Popaque _
   | Pdls_get
   | Pobj_magic _ | Punbox_float _ | Punbox_int _ | Punbox_vector _
-  | Preinterpret_unboxed_int64_as_tagged_int63
+  | Preinterpret_unboxed_int64_as_tagged_int63 | Ppeek _ | Ppoke _
   (* These don't allocate in bytecode; they're just identity functions: *)
   | Pbox_float (_, _) | Pbox_int _ | Pbox_vector (_, _)
     -> false
@@ -1864,7 +1917,7 @@ let primitive_needs_event_after = function
   | Lazy_force _ | Send _ | Send_self _ | Send_cache _
   | Apply _ | Revapply _ -> true
   | Raise _ | Raise_with_backtrace | Loc _ | Frame_pointers | Identity
-  | Unsupported _ -> false
+  | Peek _ | Poke _ | Unsupported _ -> false
 
 let transl_primitive_application loc p env ty ~poly_mode ~poly_sort
     path exp args arg_exps pos =
@@ -1908,6 +1961,8 @@ let report_error ppf = function
   | Wrong_arity_builtin_primitive prim_name ->
       fprintf ppf "Wrong arity for builtin primitive %a"
         Style.inline_code prim_name
+  | Wrong_layout_for_peek_or_poke prim_name ->
+      fprintf ppf "Unsupported layout for the %s primitive" prim_name
   | Invalid_floatarray_glb ->
       fprintf ppf
         "@[Floatarray primitives can't be used on arrays containing@ \

lambda/translprim.mli

@@ -62,6 +62,7 @@ val sort_of_native_repr :
 type error =
   | Unknown_builtin_primitive of string
   | Wrong_arity_builtin_primitive of string
+  | Wrong_layout_for_peek_or_poke of string
   | Invalid_floatarray_glb
   | Product_iarrays_unsupported
   | Invalid_array_kind_for_uninitialized_makearray_dynamic

lambda/value_rec_compiler.ml

@@ -350,7 +350,9 @@ let compute_static_size lam =
     | Patomic_cas
     | Patomic_fetch_add
     | Popaque _
-    | Pdls_get ->
+    | Pdls_get
+    | Ppeek _
+    | Ppoke _ ->
         dynamic_size lam
 
     (* Primitives specific to flambda-backend *)

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1049,7 +1049,7 @@ let close_primitive acc env ~let_bound_ids_with_kinds named
       | Patomic_exchange | Patomic_compare_exchange | Patomic_cas
       | Patomic_fetch_add | Pdls_get | Ppoll | Patomic_load _
       | Preinterpret_tagged_int63_as_unboxed_int64
-      | Preinterpret_unboxed_int64_as_tagged_int63 ->
+      | Preinterpret_unboxed_int64_as_tagged_int63 | Ppeek _ | Ppoke _ ->
         (* Inconsistent with outer match *)
         assert false
     in

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -100,6 +100,16 @@ let standard_int_or_float_of_unboxed_integer (ubint : L.unboxed_integer) :
 let standard_int_or_float_of_boxed_integer bint =
   standard_int_or_float_of_unboxed_integer (Primitive.unboxed_integer bint)
 
+let standard_int_or_float_of_peek_or_poke (layout : L.peek_or_poke) :
+    K.Standard_int_or_float.t =
+  match layout with
+  | Ppp_tagged_immediate -> Tagged_immediate
+  | Ppp_unboxed_float32 -> Naked_float32
+  | Ppp_unboxed_float -> Naked_float
+  | Ppp_unboxed_int32 -> Naked_int32
+  | Ppp_unboxed_int64 -> Naked_int64
+  | Ppp_unboxed_nativeint -> Naked_nativeint
+
 let convert_block_access_field_kind i_or_p : P.Block_access_field_kind.t =
   match i_or_p with L.Immediate -> Immediate | L.Pointer -> Any_value
 
@@ -2396,6 +2406,12 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
         "Preinterpret_tagged_int63_as_unboxed_int64 can only be used on 64-bit \
          targets";
     [Unary (Reinterpret_64_bit_word Tagged_int63_as_unboxed_int64, i)]
+  | Ppeek layout, [[ptr]] ->
+    let kind = standard_int_or_float_of_peek_or_poke layout in
+    [Unary (Peek kind, ptr)]
+  | Ppoke layout, [[ptr]; [new_value]] ->
+    let kind = standard_int_or_float_of_peek_or_poke layout in
+    [Binary (Poke kind, ptr, new_value)]
   | ( ( Pdivbint { is_safe = Unsafe; size = _; mode = _ }
       | Pmodbint { is_safe = Unsafe; size = _; mode = _ }
       | Psetglobal _ | Praise _ | Pccall _ ),
@@ -2428,7 +2444,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       | Pufloatfield _ | Patomic_load _ | Pmixedfield _
       | Preinterpret_unboxed_int64_as_tagged_int63
       | Preinterpret_tagged_int63_as_unboxed_int64
-      | Parray_element_size_in_bytes _ ),
+      | Parray_element_size_in_bytes _ | Ppeek _ ),
       ([] | _ :: _ :: _ | [([] | _ :: _ :: _)]) ) ->
     Misc.fatal_errorf
       "Closure_conversion.convert_primitive: Wrong arity for unary primitive \
@@ -2471,7 +2487,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
             _,
             _ )
       | Pcompare_ints | Pcompare_floats _ | Pcompare_bints _ | Patomic_exchange
-      | Patomic_fetch_add ),
+      | Patomic_fetch_add | Ppoke _ ),
       ( []
       | [_]
       | _ :: _ :: _ :: _