feat(ast_tools): support Atomic* types in structs

Author: overlookmotel

tasks/ast_tools/src/generators/assert_layouts.rs

@@ -8,6 +8,7 @@ use std::{
     borrow::Cow,
     cmp::{Ordering, max, min},
     num,
+    sync::atomic,
 };
 
 use phf_codegen::Map as PhfMapGen;
@@ -472,6 +473,20 @@ fn calculate_layout_for_primitive(primitive_def: &PrimitiveDef) -> Layout {
             )
         }
         "NonZeroIsize" => non_zero_usize_layout,
+        // Unlike `bool`, `AtomicBool` does not have any niches
+        "AtomicBool" => Layout::from_type::<atomic::AtomicBool>(),
+        "AtomicU8" => Layout::from_type::<atomic::AtomicU8>(),
+        "AtomicU16" => Layout::from_type::<atomic::AtomicU16>(),
+        "AtomicU32" => Layout::from_type::<atomic::AtomicU32>(),
+        "AtomicU64" => Layout::from_type::<atomic::AtomicU64>(),
+        "AtomicUsize" => usize_layout,
+        "AtomicI8" => Layout::from_type::<atomic::AtomicI8>(),
+        "AtomicI16" => Layout::from_type::<atomic::AtomicI16>(),
+        "AtomicI32" => Layout::from_type::<atomic::AtomicI32>(),
+        "AtomicI64" => Layout::from_type::<atomic::AtomicI64>(),
+        "AtomicIsize" => usize_layout,
+        // `AtomicPtr` has no niche - like `*mut T`, not `NonNull<T>`
+        "AtomicPtr" => usize_layout,
         "PointerAlign" => Layout {
             layout_64: PlatformLayout::from_size_align(0, 8),
             layout_32: PlatformLayout::from_size_align(0, 4),

tasks/ast_tools/src/generators/raw_transfer.rs

@@ -514,8 +514,9 @@ fn generate_primitive(primitive_def: &PrimitiveDef, code: &mut String, schema: &
         ",
         "f64" => "return float64[pos >> 3];",
         "&str" => STR_DESERIALIZER_BODY,
-        // Reuse deserializers for zeroed types
+        // Reuse deserializers for zeroed and atomic types
         type_name if type_name.starts_with("NonZero") => return,
+        type_name if type_name.starts_with("Atomic") => return,
         type_name => panic!("Cannot generate deserializer for primitive `{type_name}`"),
     };
 
@@ -960,6 +961,9 @@ impl DeserializeFunctionName for PrimitiveDef {
         } else if let Some(type_name) = type_name.strip_prefix("NonZero") {
             // Use zeroed type's deserializer for `NonZero*` types
             Cow::Borrowed(type_name)
+        } else if let Some(type_name) = type_name.strip_prefix("Atomic") {
+            // Use standard type's deserializer for `Atomic*` types
+            Cow::Borrowed(type_name)
         } else {
             upper_case_first(type_name)
         }

tasks/ast_tools/src/generators/raw_transfer_lazy.rs

@@ -925,8 +925,9 @@ fn generate_primitive(primitive_def: &PrimitiveDef, state: &mut State, schema: &
         ",
         "f64" => "return ast.buffer.float64[pos >> 3];",
         "&str" => STR_DESERIALIZER_BODY,
-        // Reuse constructors for zeroed types
+        // Reuse constructors for zeroed and atomic types
         type_name if type_name.starts_with("NonZero") => return,
+        type_name if type_name.starts_with("Atomic") => return,
         type_name => panic!("Cannot generate constructor for primitive `{type_name}`"),
     };
 
@@ -1248,6 +1249,9 @@ impl FunctionNames for PrimitiveDef {
         } else if let Some(type_name) = type_name.strip_prefix("NonZero") {
             // Use zeroed type's constructor for `NonZero*` types
             Cow::Borrowed(type_name)
+        } else if let Some(type_name) = type_name.strip_prefix("Atomic") {
+            // Use standard type's constructor for `Atomic*` types
+            Cow::Borrowed(type_name)
         } else {
             upper_case_first(type_name)
         }

tasks/ast_tools/src/parse/parse.rs

@@ -221,6 +221,20 @@ impl<'c> Parser<'c> {
             "NonZeroI64" => primitive("NonZeroI64"),
             "NonZeroI128" => primitive("NonZeroI128"),
             "NonZeroIsize" => primitive("NonZeroIsize"),
+            "AtomicBool" => primitive("AtomicBool"),
+            "AtomicU8" => primitive("AtomicU8"),
+            "AtomicU16" => primitive("AtomicU16"),
+            "AtomicU32" => primitive("AtomicU32"),
+            "AtomicU64" => primitive("AtomicU64"),
+            "AtomicU128" => primitive("AtomicU128"),
+            "AtomicUsize" => primitive("AtomicUsize"),
+            "AtomicI8" => primitive("AtomicI8"),
+            "AtomicI16" => primitive("AtomicI16"),
+            "AtomicI32" => primitive("AtomicI32"),
+            "AtomicI64" => primitive("AtomicI64"),
+            "AtomicI128" => primitive("AtomicI128"),
+            "AtomicIsize" => primitive("AtomicIsize"),
+            "AtomicPtr" => primitive("AtomicPtr"),
             "&str" => primitive("&str"),
             "Atom" => primitive("Atom"),
             // TODO: Remove the need for this by adding