Add support for the [@atomic] attribute on mutable fields

glittershark · glittershark · commit c3e354653575 · 2025-06-13T17:09:11.000-04:00
Add support to only the frontend for specifying the `[@atomic]` attribute on (only) mutable fields. This is a no-op currently, but is forwarded down through the ld_mutable on the label description in the typedtree. This is the first of many commits pulling in upstream support[0] for atomic record fields. I've slightly diverged from the design there, by adding atomic_flag to mutability only under Mutable rather than as a new ld_atomic field, since that feels like a better design (it avoids immutable atomic fields by construction, rather than having to assert that they never happen) [0] ocaml/ocaml#13404
diff --git a/chamelon/compat.jst.ml b/chamelon/compat.jst.ml
@@ -308,7 +308,12 @@ type tpat_array_identifier = mutability * Jkind.sort
 
 let mkTpat_array
     ?id:(mut, arg_sort =
-        (Mutable Alloc.Comonadic.Const.legacy, Jkind.Sort.value)) l =
+        ( Mutable
+            {
+              modal_upper_bound = Alloc.Comonadic.Const.legacy;
+              atomic = Nonatomic;
+            },
+          Jkind.Sort.value )) l =
   Tpat_array (mut, arg_sort, l)
 
 type tpat_tuple_identifier = string option list
diff --git a/parsing/asttypes.mli b/parsing/asttypes.mli
@@ -40,6 +40,8 @@ type private_flag = Private | Public
 
 type mutable_flag = Immutable | Mutable
 
+type atomic_flag = Nonatomic | Atomic
+
 type virtual_flag = Virtual | Concrete
 
 type override_flag = Override | Fresh
diff --git a/parsing/builtin_attributes.ml b/parsing/builtin_attributes.ml
@@ -69,6 +69,7 @@ let warn_unused () =
    misplaced attribute warnings. *)
 let builtin_attrs =
   [ "inline"
+  ; "atomic"
   ; "inlined"
   ; "specialise"
   ; "specialised"
@@ -1107,3 +1108,5 @@ let get_tracing_probe_payload (payload : Parsetree.payload) =
     | _ -> Error ()
   in
   Ok { name; name_loc; enabled_at_init; arg }
+
+let has_atomic attrs = has_attribute "atomic" attrs
diff --git a/parsing/builtin_attributes.mli b/parsing/builtin_attributes.mli
@@ -345,3 +345,5 @@ type tracing_probe =
 *)
 val get_tracing_probe_payload :
   Parsetree.payload -> (tracing_probe, unit) result
+
+val has_atomic: Parsetree.attributes -> bool
diff --git a/testsuite/tests/compiler-libs/test_untypeast.ml b/testsuite/tests/compiler-libs/test_untypeast.ml
@@ -8,17 +8,28 @@ let run s =
   let pe = Parse.expression (Lexing.from_string s) in
   let te = Typecore.type_expression (Lazy.force Env.initial) pe in
   let ute = Untypeast.untype_expression te in
-  Format.asprintf "%a" Pprintast.expression ute
+  Format.printf "%a@." Pprintast.expression ute
 ;;
 
 [%%expect{|
-val run : string -> string = <fun>
+val run : string -> unit = <fun>
 |}];;
 
 run {| match None with Some (Some _) -> () | _ -> () |};;
 
 [%%expect{|
-- : string = "match None with | Some (Some _) -> () | _ -> ()"
+match None with | Some (Some _) -> () | _ -> ()
+- : unit = ()
+|}];;
+
+run {| let open struct type t = { mutable x : int [@atomic] } end in
+       let _ = fun (v : t) -> v.x in () |};;
+
+[%%expect{|
+let open struct type t = {
+                  mutable x: int [@atomic ]} end in
+  let _ = fun (v : t) -> v.x in ()
+- : unit = ()
 |}];;
 
 (***********************************)
@@ -28,20 +39,23 @@ run {| match None with Some (Some _) -> () | _ -> () |};;
 run {| fun x y z -> function w -> x y z w |};;
 
 [%%expect{|
-- : string = "fun x y z -> function | w -> x y z w"
+fun x y z -> function | w -> x y z w
+- : unit = ()
 |}];;
 
 (* 3-ary function returning a 1-ary function *)
 run {| fun x y z -> (function w -> x y z w) |};;
 
 [%%expect{|
-- : string = "fun x y z -> (function | w -> x y z w)"
+fun x y z -> (function | w -> x y z w)
+- : unit = ()
 |}];;
 
 run {| match None with Some (Some _) -> () | _ -> () |};;
 
 [%%expect{|
-- : string = "match None with | Some (Some _) -> () | _ -> ()"
+match None with | Some (Some _) -> () | _ -> ()
+- : unit = ()
 |}];;
 
 (***********************************)
@@ -51,14 +65,16 @@ run {| match None with Some (Some _) -> () | _ -> () |};;
 run {| fun x y z -> function w -> x y z w |};;
 
 [%%expect{|
-- : string = "fun x y z -> function | w -> x y z w"
+fun x y z -> function | w -> x y z w
+- : unit = ()
 |}];;
 
 (* 3-ary function returning a 1-ary function *)
 run {| fun x y z -> (function w -> x y z w) |};;
 
 [%%expect{|
-- : string = "fun x y z -> (function | w -> x y z w)"
+fun x y z -> (function | w -> x y z w)
+- : unit = ()
 |}];;
 
 (***********************************)
@@ -67,27 +83,30 @@ run {| fun x y z -> (function w -> x y z w) |};;
 run {| let foo : 'a. 'a -> 'a = fun x -> x in foo |}
 
 [%%expect{|
-- : string = "let foo : ('a : value) . 'a -> 'a = fun x -> x in foo"
+let foo : ('a : value) . 'a -> 'a = fun x -> x in foo
+- : unit = ()
 |}];;
 
 run {| let foo : type a . a -> a = fun x -> x in foo |}
 
 [%%expect{|
-- : string =
-"let foo : ('a : value) . 'a -> 'a = fun (type a) -> ( (fun x -> x : a -> a)) in\nfoo"
+let foo : ('a : value) . 'a -> 'a = fun (type a) -> ( (fun x -> x : a -> a)) in
+foo
+- : unit = ()
 |}];;
 
 (* CR: untypeast/pprintast are totally busted on programs with modes in value
    bindings. Fix this. *)
 run {| let foo : ('a -> 'a) @ portable = fun x -> x in foo |}
 
 [%%expect{|
-- : string =
-"let (foo : 'a -> 'a) = ((fun x -> x : 'a -> 'a) : _ @ portable) in foo"
+let (foo : 'a -> 'a) = ((fun x -> x : 'a -> 'a) : _ @ portable) in foo
+- : unit = ()
 |}];;
 
 run {| let foo : 'a . ('a -> 'a) @ portable = fun x -> x in foo |}
 
 [%%expect{|
-- : string = "let foo : ('a : value) . 'a -> 'a = fun x -> x in foo"
+let foo : ('a : value) . 'a -> 'a = fun x -> x in foo
+- : unit = ()
 |}];;
diff --git a/typing/includecore.ml b/typing/includecore.ml
@@ -154,9 +154,9 @@ let value_descriptions ~loc env name
              let ty2, mode_l2, mode_y2, _ = Ctype.instance_prim p2 vd2.val_type in
              Option.iter (Mode.Locality.equate_exn loc) mode_l2;
              Option.iter (Mode.Yielding.equate_exn yield) mode_y2;
-             try 
+             try
                Ctype.moregeneral env true ty1 ty2
-             with Ctype.Moregen err -> 
+             with Ctype.Moregen err ->
                raise (Dont_match (Type err))
            ) yielding
          ) locality;
@@ -239,6 +239,7 @@ type kind_mismatch = type_kind * type_kind
 type label_mismatch =
   | Type of Errortrace.equality_error
   | Mutability of position
+  | Atomicity of position
   | Modality of Modality.Value.equate_error
 
 type record_change =
@@ -402,6 +403,10 @@ let report_label_mismatch first second env ppf err =
       Format.fprintf ppf "%s is mutable and %s is not."
         (String.capitalize_ascii (choose ord first second))
         (choose_other ord first second)
+  | Atomicity ord ->
+      Format.fprintf ppf "%s is atomic and %s is not."
+        (String.capitalize_ascii (choose ord first second))
+        (choose_other ord first second)
   | Modality err_ -> report_modality_equate_error first second ppf err_
 
 let pp_record_diff first second prefix decl env ppf (x : record_change) =
@@ -685,21 +690,27 @@ module Record_diffing = struct
   let compare_labels env params1 params2
         (ld1 : Types.label_declaration)
         (ld2 : Types.label_declaration) =
-        let mut =
+        let err =
           match ld1.ld_mutable, ld2.ld_mutable with
           | Immutable, Immutable -> None
-          | Mutable _, Immutable -> Some First
-          | Immutable, Mutable _ -> Some Second
-          | Mutable m1, Mutable m2 ->
-            let open Mode.Alloc.Comonadic.Const in
-            (if not (Misc.Le_result.equal ~le m1 legacy) then
-              Misc.fatal_errorf "Unexpected mutable(%a)" print m1);
-            (if not (Misc.Le_result.equal ~le m2 legacy) then
-              Misc.fatal_errorf "Unexpected mutable(%a)" print m2);
-            None
+          | Mutable _, Immutable -> Some (Mutability First)
+          | Immutable, Mutable _ -> Some (Mutability Second)
+          | Mutable { modal_upper_bound = m1; atomic = atomic1 },
+            Mutable { modal_upper_bound = m2; atomic = atomic2 } ->
+            begin match atomic1, atomic2 with
+            | Atomic, Nonatomic -> Some (Atomicity First)
+            | Nonatomic, Atomic -> Some (Atomicity Second)
+            | Atomic, Atomic | Nonatomic, Nonatomic ->
+                let open Mode.Alloc.Comonadic.Const in
+                (if not (Misc.Le_result.equal ~le m1 legacy) then
+                   Misc.fatal_errorf "Unexpected mutable(%a)" print m1);
+                (if not (Misc.Le_result.equal ~le m2 legacy) then
+                   Misc.fatal_errorf "Unexpected mutable(%a)" print m2);
+                None
+            end
         in
-        begin match mut with
-        | Some mut -> Some (Mutability mut)
+        begin match err with
+        | Some err -> Some err
         | None ->
           match
             Modality.Value.Const.equate ld1.ld_modalities ld2.ld_modalities
diff --git a/typing/includecore.mli b/typing/includecore.mli
@@ -63,6 +63,7 @@ type kind_mismatch = type_kind * type_kind
 type label_mismatch =
   | Type of Errortrace.equality_error
   | Mutability of position
+  | Atomicity of position
   | Modality of Mode.Modality.Value.equate_error
 
 type record_change =
diff --git a/typing/printtyp.ml b/typing/printtyp.ml
@@ -1838,10 +1838,10 @@ let param_jkind ty =
 let tree_of_label l =
   let mut =
     match l.ld_mutable with
-    | Mutable m ->
+    | Mutable { modal_upper_bound; _ } ->
         let mut =
           let open Alloc.Comonadic.Const in
-          if Misc.Le_result.equal ~le m legacy then
+          if Misc.Le_result.equal ~le modal_upper_bound legacy then
             Om_mutable None
           else
             Om_mutable (Some "<non-legacy>")
diff --git a/typing/printtyped.ml b/typing/printtyped.ml
@@ -91,8 +91,10 @@ let fmt_mutable_flag f x =
 let fmt_mutable_mode_flag f (x : Types.mutability) =
   match x with
   | Immutable -> fprintf f "Immutable"
-  | Mutable m ->
-    fprintf f "Mutable(%a)" Mode.Alloc.Comonadic.Const.print m
+  | Mutable { modal_upper_bound; atomic = Nonatomic } ->
+    fprintf f "Mutable(%a)" Mode.Alloc.Comonadic.Const.print modal_upper_bound
+  | Mutable { modal_upper_bound; atomic = Atomic } ->
+    fprintf f "Atomic(%a)" Mode.Alloc.Comonadic.Const.print modal_upper_bound
 
 let fmt_virtual_flag f x =
   match x with
diff --git a/typing/typecore.ml b/typing/typecore.ml
@@ -1018,7 +1018,7 @@ let mutable_mode m0 =
 let check_construct_mutability ~loc ~env mutability ty ?modalities block_mode =
   match mutability with
   | Immutable -> ()
-  | Mutable m0 ->
+  | Mutable { modal_upper_bound = m0; _ } ->
       let m0 = mutable_mode m0 in
       let m0 = cross_left env ty ?modalities m0 in
       submode ~loc ~env m0 block_mode
@@ -3034,7 +3034,11 @@ and type_pat_aux
   | Ppat_array (mut, spl) ->
       let mut =
         match mut with
-        | Mutable -> Mutable Alloc.Comonadic.Const.legacy
+        | Mutable  -> Mutable {
+          modal_upper_bound = Alloc.Comonadic.Const.legacy;
+          (* CR aspsmith: Revisit once we support atomic arrays *)
+          atomic = Nonatomic
+        }
         | Immutable ->
             Language_extension.assert_enabled ~loc Immutable_arrays ();
             Immutable
@@ -6193,7 +6197,8 @@ and type_expect_
       in
       let (label_loc, label, newval) =
         match label.lbl_mut with
-        | Mutable m0 ->
+        | Mutable { modal_upper_bound = m0; atomic } ->
+          ignore atomic;  (* CR aspsmith: TODO *)
           submode ~loc:record.exp_loc ~env rmode mode_mutate_mutable;
           let mode = mutable_mode m0 |> mode_default in
           let mode = mode_modality label.lbl_modalities mode in
@@ -6215,7 +6220,11 @@ and type_expect_
   | Pexp_array(mut, sargl) ->
       let mutability =
         match mut with
-        | Mutable -> Mutable Alloc.Comonadic.Const.legacy
+        | Mutable -> Mutable {
+          modal_upper_bound = Alloc.Comonadic.Const.legacy;
+          (* CR aspsmith: Revisit once we support atomic record fields *)
+          atomic = Nonatomic;
+        }
         | Immutable ->
             Language_extension.assert_enabled ~loc Immutable_arrays ();
             Immutable
@@ -9802,10 +9811,15 @@ and type_comprehension_expr ~loc ~env ~ty_expected ~attributes cexpr =
         Predef.list_argument_jkind
     | Pcomp_array_comprehension (amut, comp) ->
         let container_type, mut = match amut with
-          | Mutable   -> Predef.type_array, Mutable Alloc.Comonadic.Const.legacy
-          | Immutable ->
-              Language_extension.assert_enabled ~loc Immutable_arrays ();
-              Predef.type_iarray, Immutable
+        | Mutable   ->
+          Predef.type_array, Mutable {
+            modal_upper_bound = Alloc.Comonadic.Const.legacy;
+            (* CR aspsmith: Revisit once we support atomic arrays *)
+            atomic = Nonatomic
+          }
+        | Immutable ->
+          Language_extension.assert_enabled ~loc Immutable_arrays ();
+          Predef.type_iarray, Immutable
         in
         (Array_comprehension mut : comprehension_type),
         container_type,
diff --git a/typing/typedecl.ml b/typing/typedecl.ml
@@ -141,6 +141,7 @@ type error =
   | Unsafe_mode_crossing_on_invalid_type_kind
   | Illegal_baggage of jkind_l
   | No_unboxed_version of Path.t
+  | Atomic_field_must_be_mutable of string
 
 open Typedtree
 
@@ -481,19 +482,32 @@ let transl_labels (type rep) ~(record_form : rep record_form) ~new_var_jkind
           pld_type=arg;pld_loc=loc;pld_attributes=attrs} =
     Builtin_attributes.warning_scope attrs
       (fun () ->
+         let is_atomic = Builtin_attributes.has_atomic attrs in
          let mut : mutability =
-          match mut with
-          | Immutable -> Immutable
-          | Mutable ->
+          match mut, is_atomic with
+          | Immutable, false -> Immutable
+          | Immutable, true ->
+            raise (Error (loc, Atomic_field_must_be_mutable name.txt))
+          | Mutable, is_atomic ->
               match record_form with
-              | Legacy -> Mutable Mode.Alloc.Comonadic.Const.legacy
+              | Legacy -> Mutable {
+                modal_upper_bound = Mode.Alloc.Comonadic.Const.legacy;
+                atomic = if is_atomic then Atomic else Nonatomic
+              }
               | Unboxed_product -> raise(Error(loc, Unboxed_mutable_label))
          in
          let modalities =
           Typemode.transl_modalities ~maturity:Stable mut modalities
          in
          let arg = Ast_helper.Typ.force_poly arg in
          let cty = transl_simple_type ~new_var_jkind env ?univars ~closed Mode.Alloc.Const.legacy arg in
+         let is_atomic = Builtin_attributes.has_atomic attrs in
+         begin match is_atomic, mut with
+         | true, Mutable _
+         | false, _  -> ()
+         | true, Immutable ->
+           raise (Error (loc, Atomic_field_must_be_mutable name.txt));
+         end;
          {ld_id = Ident.create_local name.txt;
           ld_name = name;
           ld_uid = Uid.mk ~current_unit:(Env.get_unit_name ());
@@ -4694,6 +4708,10 @@ let report_error ppf = function
   | No_unboxed_version p ->
       fprintf ppf "@[The type %a@ has no unboxed version.@]"
         (Style.as_inline_code Printtyp.path) p
+  | Atomic_field_must_be_mutable name ->
+      fprintf ppf
+        "@[The label %a must be mutable to be declared atomic.@]"
+        Style.inline_code name
 
 let () =
   Location.register_error_of_exn
diff --git a/typing/typedecl.mli b/typing/typedecl.mli
@@ -177,6 +177,7 @@ type error =
   | Unsafe_mode_crossing_on_invalid_type_kind
   | Illegal_baggage of jkind_l
   | No_unboxed_version of Path.t
+  | Atomic_field_must_be_mutable of string
 
 exception Error of Location.t * error
 
diff --git a/typing/types.ml b/typing/types.ml
diff --git a/typing/types.mli b/typing/types.mli
diff --git a/typing/untypeast.ml b/typing/untypeast.ml
