From 7b721467fa377c85fb4fe15f7669319d11137b76 Mon Sep 17 00:00:00 2001
From: Aspen Smith <aspsmith@janestreet.com>
Date: Sun, 29 Jun 2025 13:54:04 -0400
Subject: [PATCH] Add support fo atomic locs in the stdlib

---
 stdlib/atomic.ml  | 52 +++++++++++++++++++++++++++++++++++++
 stdlib/atomic.mli | 65 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+)

diff --git a/stdlib/atomic.ml b/stdlib/atomic.ml
index 079549fbb80..f4230808533 100644
--- a/stdlib/atomic.ml
+++ b/stdlib/atomic.ml
@@ -58,3 +58,55 @@ module Contended = struct
     'a t @ contended local -> 'a -> 'a -> 'a @@ portable =
     "%atomic_compare_exchange"
 end
+
+module Loc = struct
+  type ('a : value_or_null) t : mutable_data with 'a = 'a atomic_loc
+  external get : ('a : value_or_null).
+    'a t @ local -> 'a @@ portable = "%atomic_load_loc"
+  external set : ('a : value_or_null).
+    'a t @ local -> 'a -> unit @@ portable = "%atomic_set_loc"
+  external exchange : ('a : value_or_null).
+    'a t @ local -> 'a -> 'a @@ portable = "%atomic_exchange_loc"
+  external compare_and_set : ('a : value_or_null).
+    'a t @ local -> 'a -> 'a -> bool @@ portable = "%atomic_cas_loc"
+  external compare_exchange : ('a : value_or_null).
+    'a t @ local -> 'a -> 'a -> 'a @@ portable = "%atomic_compare_exchange_loc"
+
+  external fetch_and_add
+    : int t @ contended local -> int -> int @@ portable = "%atomic_fetch_add_loc"
+
+  external add
+    : int t @ contended local -> int -> unit @@ portable = "%atomic_add_loc"
+
+  external sub
+    : int t @ contended local -> int -> unit @@ portable = "%atomic_sub_loc"
+
+  external logand
+    : int t @ contended local -> int -> unit @@ portable = "%atomic_land_loc"
+
+  external logor
+    : int t @ contended local -> int -> unit @@ portable = "%atomic_lor_loc"
+
+  external logxor
+    : int t @ contended local -> int -> unit @@ portable = "%atomic_lxor_loc"
+
+  let incr t = add t 1
+  let decr t = sub t 1
+
+  module Contended = struct
+    external get : ('a : value_or_null mod contended).
+      'a t @ contended local -> 'a @@ portable = "%atomic_load_loc"
+
+    external set : ('a : value_or_null mod portable).
+      'a t @ contended local -> 'a -> unit @@ portable = "%atomic_set_loc"
+
+    external exchange : ('a : value_or_null mod contended portable).
+      'a t @ contended local -> 'a -> 'a @@ portable = "%atomic_exchange_loc"
+
+    external compare_and_set : ('a : value_or_null mod portable).
+      'a t @ contended local -> 'a -> 'a -> bool @@ portable = "%atomic_cas_loc"
+
+    external compare_exchange : ('a : value_or_null mod contended portable).
+      'a t @ contended local -> 'a -> 'a -> 'a @@ portable = "%atomic_compare_exchange_loc"
+  end
+end
diff --git a/stdlib/atomic.mli b/stdlib/atomic.mli
index 128706ce7ae..ea8034c41e8 100644
--- a/stdlib/atomic.mli
+++ b/stdlib/atomic.mli
@@ -126,6 +126,71 @@ module Contended : sig
     'a t @ contended local -> 'a -> 'a -> 'a = "%atomic_compare_exchange"
 end
 
+(** Atomic "locations", such as record fields. *)
+module Loc : sig
+  (** This module exposes a dedicated type ['a Atomic.Loc.t] for
+      atomic locations (storing a value of type ['a]) inside objects
+      that may not be atomic references. It is used in particular for
+      atomic record fields: if a record [r] has an atomic field [f] of
+      type [foo], then [[%atomic.loc r.f]] has type [foo Atomic.Loc.t].
+
+      The API below mirrors the API to access {{!t}atomic references},
+      see the documentation above for more information. *)
+  type ('a : value_or_null) t : mutable_data with 'a = 'a atomic_loc
+
+  external get : ('a : value_or_null). 'a t @ local -> 'a = "%atomic_load_loc"
+
+  external set : ('a : value_or_null).
+    'a t @ local -> 'a -> unit = "%atomic_set_loc"
+
+  external exchange : ('a : value_or_null).
+    'a t @ local -> 'a -> 'a = "%atomic_exchange_loc"
+
+  external compare_and_set : ('a : value_or_null).
+    'a t @ local -> 'a -> 'a -> bool = "%atomic_cas_loc"
+
+  external compare_exchange : ('a : value_or_null).
+    'a t @ local -> 'a -> 'a -> 'a = "%atomic_compare_exchange_loc"
+
+  external fetch_and_add
+    : int t @ contended local -> int -> int = "%atomic_fetch_add_loc"
+
+  external add
+    : int t @ contended local -> int -> unit = "%atomic_add_loc"
+
+  external sub
+    : int t @ contended local -> int -> unit = "%atomic_sub_loc"
+
+  external logand
+    : int t @ contended local -> int -> unit = "%atomic_land_loc"
+
+  external logor
+    : int t @ contended local -> int -> unit = "%atomic_lor_loc"
+
+  external logxor
+    : int t @ contended local -> int -> unit = "%atomic_lxor_loc"
+
+  val incr : int t @ contended local -> unit
+  val decr : int t @ contended local -> unit
+
+  module Contended : sig
+    external get : ('a : value_or_null mod contended).
+      'a t @ contended local -> 'a = "%atomic_load_loc"
+
+    external set : ('a : value_or_null mod portable).
+      'a t @ contended local -> 'a -> unit = "%atomic_set_loc"
+
+    external exchange : ('a : value_or_null mod contended portable).
+      'a t @ contended local -> 'a -> 'a = "%atomic_exchange_loc"
+
+    external compare_and_set : ('a : value_or_null mod portable).
+      'a t @ contended local -> 'a -> 'a -> bool = "%atomic_cas_loc"
+
+    external compare_exchange : ('a : value_or_null mod contended portable).
+      'a t @ contended local -> 'a -> 'a -> 'a = "%atomic_compare_exchange_loc"
+  end
+end
+
 (** {1:examples Examples}
 
     {2 Basic Thread Coordination}