value_or_null atomics with correct result kind inference (#3807)

* Infer or_null result values for atomic ops

* stdlib changes

* tests

* `prove_is_immediate` for atomic comparsion arguments

* line limit

* rename to `prove_is_not_a_pointer`

---------

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

Author: dkalinichenko-js

lambda/lambda.ml

@@ -994,6 +994,7 @@ let nullable_value raw_kind =
 
 let layout_unit = non_null_value Pintval
 let layout_int = non_null_value Pintval
+let layout_int_or_null = nullable_value Pintval
 let layout_array kind = non_null_value (Parrayval kind)
 let layout_block = non_null_value Pgenval
 let layout_list =
@@ -2389,13 +2390,19 @@ let primitive_result_layout (p : primitive) =
   | (Parray_to_iarray | Parray_of_iarray) -> layout_any_value
   | Pget_header _ -> layout_boxed_int Boxed_nativeint
   | Prunstack | Presume | Pperform | Preperform -> layout_any_value
-  | Patomic_load { immediate_or_pointer = Immediate } -> layout_int
-  | Patomic_load { immediate_or_pointer = Pointer } -> layout_any_value
+  | Patomic_load { immediate_or_pointer = Immediate } ->
+    layout_int_or_null
+  | Patomic_load { immediate_or_pointer = Pointer } ->
+    layout_any_value
   | Patomic_set _ -> layout_unit
-  | Patomic_exchange { immediate_or_pointer = Immediate } -> layout_int
-  | Patomic_exchange { immediate_or_pointer = Pointer } -> layout_any_value
-  | Patomic_compare_exchange { immediate_or_pointer = Immediate } -> layout_int
-  | Patomic_compare_exchange { immediate_or_pointer = Pointer } -> layout_any_value
+  | Patomic_exchange { immediate_or_pointer = Immediate } ->
+    layout_int_or_null
+  | Patomic_exchange { immediate_or_pointer = Pointer } ->
+    layout_any_value
+  | Patomic_compare_exchange { immediate_or_pointer = Immediate } ->
+    layout_int_or_null
+  | Patomic_compare_exchange { immediate_or_pointer = Pointer } ->
+    layout_any_value
   | Patomic_compare_set _
   | Patomic_fetch_add -> layout_int
   | Pdls_get -> layout_any_value

middle_end/flambda2/simplify/simplify_ternary_primitive.ml

@@ -87,9 +87,9 @@ let simplify_atomic_compare_and_set_or_exchange_args
        also because the old value cannot be a pointer if the operation is to
        perform a write. *)
     let is_immediate ty =
-      match T.prove_is_a_tagged_immediate (DA.typing_env dacc) ty with
-      | Proved () -> true
-      | Unknown -> false
+      match T.prove_is_not_a_pointer (DA.typing_env dacc) ty with
+      | Proved true -> true
+      | Proved false | Unknown -> false
     in
     if is_immediate comparison_value_ty && is_immediate new_value_ty
     then Immediate
@@ -139,7 +139,7 @@ let simplify_atomic_compare_exchange
   in
   let result_var_ty =
     match atomic_kind (* N.B. not [args_kind] *) with
-    | Immediate -> T.any_tagged_immediate
+    | Immediate -> T.any_tagged_immediate_or_null
     | Any_value -> T.any_value
   in
   let dacc = DA.add_variable dacc result_var result_var_ty in

middle_end/flambda2/simplify/simplify_unary_primitive.ml

@@ -769,7 +769,7 @@ let simplify_atomic_load (block_access_field_kind : P.Block_access_field_kind.t)
     ~original_prim dacc ~original_term ~arg:_ ~arg_ty:_ ~result_var =
   match block_access_field_kind with
   | Immediate ->
-    let dacc = DA.add_variable dacc result_var T.any_tagged_immediate in
+    let dacc = DA.add_variable dacc result_var T.any_tagged_immediate_or_null in
     SPR.create original_term ~try_reify:false dacc
   | Any_value ->
     SPR.create_unknown dacc ~result_var

middle_end/flambda2/types/flambda2_types.mli

@@ -327,6 +327,8 @@ val any_value : t
 
 val any_tagged_immediate : t
 
+val any_tagged_immediate_or_null : t
+
 val any_tagged_bool : t
 
 val any_boxed_float32 : t
@@ -646,6 +648,9 @@ val prove_unique_fully_constructed_immutable_heap_block :
 
 val prove_is_int : Typing_env.t -> t -> bool proof_of_property
 
+(* Either a tagged integer or a null poitner. *)
+val prove_is_not_a_pointer : Typing_env.t -> t -> bool proof_of_property
+
 (* Returns the result of [Is_flat_float_array] *)
 val meet_is_flat_float_array : Typing_env.t -> t -> bool meet_shortcut
 

middle_end/flambda2/types/grammar/more_type_creators.ml

@@ -74,6 +74,10 @@ let any_tagged_immediate_non_null =
     ~immediates:Unknown ~blocks:(Known TG.Row_like_for_blocks.bottom)
     ~extensions:No_extensions
 
+let any_tagged_immediate_or_null =
+  TG.create_from_head_value
+    { non_null = Ok any_tagged_immediate_non_null; is_null = Maybe_null }
+
 let these_tagged_immediates0 imms =
   match Targetint_31_63.Set.get_singleton imms with
   | Some imm -> TG.this_tagged_immediate imm

middle_end/flambda2/types/grammar/more_type_creators.mli

@@ -43,6 +43,8 @@ val these_naked_vec128s :
 
 val any_tagged_immediate : Type_grammar.t
 
+val any_tagged_immediate_or_null : Type_grammar.t
+
 val these_tagged_immediates0 : Targetint_31_63.Set.t -> Type_grammar.t
 
 val these_tagged_immediates : Targetint_31_63.Set.t -> Type_grammar.t

middle_end/flambda2/types/provers.ml

@@ -153,6 +153,25 @@ let prove_is_int env t =
 let meet_is_int_variant_only env t =
   gen_value_to_meet (prove_is_int_generic_value ~variant_only:true) env t
 
+let prove_is_not_a_pointer_generic_value env t =
+  match expand_head env t with
+  | Value Unknown -> Unknown
+  | Value Bottom -> Invalid
+  | Value (Ok { is_null = Maybe_null; non_null = Bottom }) -> Proved true
+  | Value (Ok { is_null = Not_null; non_null = Bottom }) -> Invalid
+  | Value (Ok { is_null = Maybe_null | Not_null; non_null = Unknown }) ->
+    Unknown
+  | Value (Ok { is_null; non_null = Ok head }) -> (
+    match prove_is_int_generic_value ~variant_only:false env head, is_null with
+    | Proved true, (Maybe_null | Not_null) -> Proved true
+    | (Proved false | Unknown), Maybe_null | Unknown, Not_null -> Unknown
+    | Proved false, Not_null -> Proved false
+    | Invalid, _ -> Invalid (* Ought to be impossible. *))
+  | _ -> wrong_kind "Value" t Invalid
+
+let prove_is_not_a_pointer env t =
+  as_property (prove_is_not_a_pointer_generic_value env t)
+
 (* Note: this function returns a generic proof because we want to propagate the
    Invalid cases to prove_naked_immediates_generic, but it's not suitable for
    implementing [meet_get_tag] because it doesn't ignore the immediates part of

middle_end/flambda2/types/provers.mli

@@ -123,6 +123,9 @@ val prove_unique_tag_and_size :
 
 val prove_is_int : Typing_env.t -> Type_grammar.t -> bool proof_of_property
 
+val prove_is_not_a_pointer :
+  Typing_env.t -> Type_grammar.t -> bool proof_of_property
+
 val meet_is_int_variant_only :
   Typing_env.t -> Type_grammar.t -> bool meet_shortcut
 

stdlib/atomic.ml

@@ -12,15 +12,25 @@
 (*                                                                        *)
 (**************************************************************************)
 
-type !'a t : mutable_data with 'a
+type (!'a : value_or_null) t : mutable_data with 'a
 
-external make : 'a -> ('a t[@local_opt]) @@ portable = "%makemutable"
-external make_contended : 'a -> ('a t[@local_opt]) @@ portable = "caml_atomic_make_contended"
-external get : 'a t @ local -> 'a @@ portable = "%atomic_load"
-external set : 'a t @ local -> 'a -> unit @@ portable = "%atomic_set"
-external exchange : 'a t @ local -> 'a -> 'a @@ portable = "%atomic_exchange"
-external compare_and_set : 'a t @ local -> 'a -> 'a -> bool @@ portable = "%atomic_cas"
-external compare_exchange : 'a t @ local -> 'a -> 'a -> 'a @@ portable = "%atomic_compare_exchange"
+external make : ('a : value_or_null).
+  'a -> ('a t[@local_opt]) @@ portable = "%makemutable"
+external make_contended : ('a : value_or_null).
+  'a -> ('a t[@local_opt]) @@ portable = "caml_atomic_make_contended"
+external get : ('a : value_or_null).
+  'a t @ local -> 'a @@ portable = "%atomic_load"
+external set : ('a : value_or_null).
+  'a t @ local -> 'a -> unit @@ portable = "%atomic_set"
+external exchange : ('a : value_or_null).
+  'a t @ local -> 'a -> 'a @@ portable =
+  "%atomic_exchange"
+external compare_and_set : ('a : value_or_null).
+  'a t @ local -> 'a -> 'a -> bool @@ portable =
+  "%atomic_cas"
+external compare_exchange : ('a : value_or_null).
+  'a t @ local -> 'a -> 'a -> 'a @@ portable =
+  "%atomic_compare_exchange"
 external fetch_and_add : int t @ contended local -> int -> int @@ portable = "%atomic_fetch_add"
 external add : int t @ contended local -> int -> unit @@ portable = "%atomic_add"
 external sub : int t @ contended local -> int -> unit @@ portable = "%atomic_sub"
@@ -32,9 +42,19 @@ let incr r = add r 1
 let decr r = sub r 1
 
 module Contended = struct
-  external get : ('a : value mod contended). 'a t @ contended local -> 'a @@ portable = "%atomic_load"
-  external set : ('a : value mod portable). 'a t @ contended local -> 'a -> unit @@ portable = "%atomic_set"
-  external exchange : ('a : value mod contended portable). 'a t @ contended local -> 'a -> 'a @@ portable = "%atomic_exchange"
-  external compare_and_set : ('a : value mod portable). 'a t @ contended local -> 'a -> 'a -> bool @@ portable = "%atomic_cas"
-  external compare_exchange : ('a : value mod contended portable). 'a t @ contended local -> 'a -> 'a -> 'a @@ portable = "%atomic_compare_exchange"
+  external get : ('a : value_or_null mod contended).
+    'a t @ contended local -> 'a @@ portable =
+    "%atomic_load"
+  external set : ('a : value_or_null mod portable).
+    'a t @ contended local -> 'a -> unit @@ portable =
+    "%atomic_set"
+  external exchange : ('a : value_or_null mod contended portable).
+    'a t @ contended local -> 'a -> 'a @@ portable =
+    "%atomic_exchange"
+  external compare_and_set : ('a : value_or_null mod portable).
+    'a t @ contended local -> 'a -> 'a -> bool @@ portable =
+    "%atomic_cas"
+  external compare_exchange : ('a : value_or_null mod contended portable).
+    'a t @ contended local -> 'a -> 'a -> 'a @@ portable =
+    "%atomic_compare_exchange"
 end

stdlib/atomic.mli

@@ -25,10 +25,11 @@
 *)
 
 (** An atomic (mutable) reference to a value of type ['a]. *)
-type !'a t : mutable_data with 'a
+type (!'a : value_or_null) t : mutable_data with 'a
 
 (** Create an atomic reference. *)
-external make : 'a -> ('a t[@local_opt]) = "%makemutable"
+external make : ('a : value_or_null) .
+  'a -> ('a t[@local_opt]) = "%makemutable"
 
 (** Create an atomic reference that is alone on a cache line. It occupies 4-16x
     the memory of one allocated with [make v].
@@ -44,28 +45,33 @@ external make : 'a -> ('a t[@local_opt]) = "%makemutable"
 
     CR ocaml 5 all-runtime5: does not support runtime4 *)
 
-external make_contended : 'a -> ('a t[@local_opt]) = "caml_atomic_make_contended"
+external make_contended : ('a : value_or_null).
+  'a -> ('a t[@local_opt]) = "caml_atomic_make_contended"
 
 (** Get the current value of the atomic reference. *)
-val get : 'a t @ local -> 'a
+val get : ('a : value_or_null). 'a t @ local -> 'a
 
 (** Set a new value for the atomic reference. *)
-external set : 'a t @ local -> 'a -> unit = "%atomic_set"
+external set : ('a : value_or_null).
+  'a t @ local -> 'a -> unit = "%atomic_set"
 
 (** Set a new value for the atomic reference, and return the current value. *)
-external exchange : 'a t @ local -> 'a -> 'a = "%atomic_exchange"
+external exchange : ('a : value_or_null).
+  'a t @ local -> 'a -> 'a = "%atomic_exchange"
 
 (** [compare_and_set r seen v] sets the new value of [r] to [v] only
     if its current value is physically equal to [seen] -- the
     comparison and the set occur atomically. Returns [true] if the
     comparison succeeded (so the set happened) and [false]
     otherwise. *)
-external compare_and_set : 'a t @ local -> 'a -> 'a -> bool = "%atomic_cas"
+external compare_and_set : ('a : value_or_null).
+  'a t @ local -> 'a -> 'a -> bool = "%atomic_cas"
 
 (** [compare_exchange r seen v] sets the new value of [r] to [v] only
     if its current value is physically equal to [seen] -- the comparison
     and the set occur atomically. Returns the previous value. *)
-external compare_exchange : 'a t @ local -> 'a -> 'a -> 'a = "%atomic_compare_exchange"
+external compare_exchange : ('a : value_or_null).
+  'a t @ local -> 'a -> 'a -> 'a = "%atomic_compare_exchange"
 
 (** [fetch_and_add r n] atomically increments the value of [r] by [n],
     and returns the current value (before the increment). *)
@@ -96,24 +102,28 @@ val decr : int t @ contended local -> unit
     via modes. *)
 module Contended : sig
   (** Like {!get}, but can be called on an atomic that came from another domain. *)
-  val get : ('a : value mod contended). 'a t @ contended local -> 'a
+  val get : ('a : value_or_null mod contended).
+    'a t @ contended local -> 'a
 
   (** Like {!set}, but can be called on an atomic that came from another domain. *)
   external set
-    : ('a : value mod portable). 'a t @ contended local -> 'a -> unit = "%atomic_set"
+    : ('a : value_or_null mod portable).
+    'a t @ contended local -> 'a -> unit ="%atomic_set"
 
   (** Like {!exchange}, but can be called on an atomic that came from another domain. *)
-  external exchange : ('a : value mod contended portable). 'a t @ contended local -> 'a -> 'a
-    = "%atomic_exchange"
+  external exchange :
+    ('a : value_or_null mod contended portable).
+    'a t @ contended local -> 'a -> 'a = "%atomic_exchange"
 
   (** Like {!compare_and_set}, but can be called on an atomic that came from another domain. *)
   external compare_and_set
-    : ('a : value mod portable). 'a t @ contended local -> 'a -> 'a -> bool = "%atomic_cas"
+    : ('a : value_or_null mod portable).
+    'a t @ contended local -> 'a -> 'a -> bool = "%atomic_cas"
 
   (** Like {!compare_exchange}, but can be called on an atomic that came from another domain. *)
   external compare_exchange
-    : ('a : value mod contended portable). 'a t @ contended local -> 'a -> 'a -> 'a
-    = "%atomic_compare_exchange"
+    : ('a : value_or_null mod contended portable).
+    'a t @ contended local -> 'a -> 'a -> 'a = "%atomic_compare_exchange"
 end
 
 (** {1:examples Examples}

testsuite/tests/typing-layouts-or-null/atomics.ml

@@ -0,0 +1,98 @@
+(* TEST
+ flambda2;
+ {
+   native;
+ } {
+   flags = "-O3";
+   native;
+ } {
+   flags = "-Oclassic";
+   native;
+ } {
+   bytecode;
+ }
+*)
+
+external get_int : int or_null Atomic.t -> int or_null = "%atomic_load"
+
+external exchange_int : int or_null Atomic.t -> int or_null -> int or_null
+  = "%atomic_exchange"
+
+external compare_and_exchange_int
+  : int or_null Atomic.t -> int or_null -> int or_null -> int or_null
+  = "%atomic_compare_exchange"
+
+let () =
+  let x = Sys.opaque_identity (Atomic.make Null) in
+  match get_int x with
+  | Null -> ()
+  | This _ -> assert false
+;;
+
+let () =
+  let x = Sys.opaque_identity (Atomic.make (This 47)) in
+  match get_int x with
+  | This 47 -> ()
+  | This _ -> assert false
+  | Null -> assert false
+;;
+
+let () =
+  let x = Sys.opaque_identity (Atomic.make Null) in
+  Atomic.set x (This 5);
+  Atomic.set x (This 7);
+  match get_int x with
+  | This 7 -> ()
+  | Null -> assert false
+  | This _ -> assert false
+;;
+
+let () =
+  let x = Sys.opaque_identity (Atomic.make Null) in
+  match exchange_int x (This 11) with
+  | Null -> ()
+  | This _ -> assert false
+;;
+
+let () =
+  let x = Sys.opaque_identity (Atomic.make (This 11)) in
+  (match exchange_int x Null with
+  | This 11 -> ()
+  | This _ -> assert false
+  | Null -> assert false);
+  match exchange_int x (This 0) with
+  | Null -> ()
+  | This _ -> assert false
+;;
+
+let () =
+  let x = Sys.opaque_identity (Atomic.make Null) in
+  match compare_and_exchange_int x (This 5) (This 8) with
+  | Null -> ()
+  | This _ -> assert false
+;;
+
+let () =
+  let x = Sys.opaque_identity (Atomic.make Null) in
+  (match compare_and_exchange_int x Null (This 42) with
+  | Null -> ()
+  | This _ -> assert false);
+  match compare_and_exchange_int x Null (This 45) with
+  | This 42 -> ()
+  | This _ -> assert false
+  | Null -> assert false
+;;
+
+let () =
+  let aaa = This (ref "aaa") in
+  let bb = Sys.opaque_identity (This (ref "bb")) in
+  let zzzzz = This (Sys.opaque_identity (ref "zzzzz")) in
+  let arr = This (ref "arr") in
+  let y = Atomic.make aaa in
+  assert (Atomic.exchange y Null == aaa);
+  assert (Atomic.exchange y bb == Null);
+  assert (Atomic.compare_exchange y aaa zzzzz == bb);
+  assert (Atomic.compare_exchange y aaa Null == bb);
+  assert (Atomic.compare_exchange y bb arr == bb);
+  assert (Atomic.compare_and_set y arr Null)
+;;