feat(ast_tools): Support NonZero* primitives

tasks/ast_tools/src/generators/assert_layouts.rs

@@ -4,7 +4,10 @@
 //! Memory layouts are different on 64-bit and 32-bit platforms.
 //! Calculate each separately, and generate assertions for each.
 
-use std::cmp::{max, min};
+use std::{
+    cmp::{max, min},
+    num,
+};
 
 use proc_macro2::TokenStream;
 use quote::quote;
@@ -346,6 +349,11 @@ fn calculate_layout_for_primitive(primitive_def: &PrimitiveDef) -> Layout {
         layout_64: PlatformLayout::from_size_align(8, 8),
         layout_32: PlatformLayout::from_size_align(4, 4),
     };
+    // `NonZeroUsize` and `NonZeroIsize` are pointer-sized, with a single niche
+    let non_zero_usize_layout = Layout {
+        layout_64: PlatformLayout::from_size_align_niche(8, 8, Niche::new(0, 8, true, 1)),
+        layout_32: PlatformLayout::from_size_align_niche(4, 4, Niche::new(0, 4, true, 1)),
+    };
 
     #[expect(clippy::match_same_arms)]
     match primitive_def.name() {
@@ -370,6 +378,22 @@ fn calculate_layout_for_primitive(primitive_def: &PrimitiveDef) -> Layout {
         "f64" => Layout::from_type::<f64>(),
         "&str" => str_layout.clone(),
         "Atom" => str_layout,
+        "NonZeroU8" => Layout::from_type_with_niche_for_zero::<num::NonZeroU8>(),
+        "NonZeroU16" => Layout::from_type_with_niche_for_zero::<num::NonZeroU16>(),
+        "NonZeroU32" => Layout::from_type_with_niche_for_zero::<num::NonZeroU32>(),
+        "NonZeroU64" => Layout::from_type_with_niche_for_zero::<num::NonZeroU64>(),
+        "NonZeroU128" => {
+            panic!("Cannot calculate alignment for `NonZeroU128`. It differs depending on Rust version.")
+        }
+        "NonZeroUsize" => non_zero_usize_layout.clone(),
+        "NonZeroI8" => Layout::from_type_with_niche_for_zero::<num::NonZeroI8>(),
+        "NonZeroI16" => Layout::from_type_with_niche_for_zero::<num::NonZeroI16>(),
+        "NonZeroI32" => Layout::from_type_with_niche_for_zero::<num::NonZeroI32>(),
+        "NonZeroI64" => Layout::from_type_with_niche_for_zero::<num::NonZeroI64>(),
+        "NonZeroI128" => {
+            panic!("Cannot calculate alignment for `NonZeroI128`. It differs depending on Rust version.")
+        }
+        "NonZeroIsize" => non_zero_usize_layout,
         "ScopeId" => semantic_id_layout.clone(),
         "SymbolId" => semantic_id_layout.clone(),
         "ReferenceId" => semantic_id_layout,

tasks/ast_tools/src/parse/parse.rs

@@ -121,6 +121,18 @@ impl<'c> Parser<'c> {
             "isize" => primitive("isize"),
             "f32" => primitive("f32"),
             "f64" => primitive("f64"),
+            "NonZeroU8" => primitive("NonZeroU8"),
+            "NonZeroU16" => primitive("NonZeroU16"),
+            "NonZeroU32" => primitive("NonZeroU32"),
+            "NonZeroU64" => primitive("NonZeroU64"),
+            "NonZeroU128" => primitive("NonZeroU128"),
+            "NonZeroUsize" => primitive("NonZeroUsize"),
+            "NonZeroI8" => primitive("NonZeroI8"),
+            "NonZeroI16" => primitive("NonZeroI16"),
+            "NonZeroI32" => primitive("NonZeroI32"),
+            "NonZeroI64" => primitive("NonZeroI64"),
+            "NonZeroI128" => primitive("NonZeroI128"),
+            "NonZeroIsize" => primitive("NonZeroIsize"),
             "&str" => primitive("&str"),
             "Atom" => primitive("Atom"),
             "ScopeId" => primitive("ScopeId"),

tasks/ast_tools/src/schema/extensions/layout.rs

@@ -18,6 +18,16 @@ impl Layout {
         )
     }
 
+    /// Create [`Layout`] from a Rust type.
+    pub fn from_type_with_niche_for_zero<T>() -> Self {
+        let size = u32::try_from(size_of::<T>()).unwrap();
+        Self::from_size_align_niche(
+            size,
+            u32::try_from(align_of::<T>()).unwrap(),
+            Niche::new(0, size, true, 1),
+        )
+    }
+
     /// Create [`Layout`] from `size` and `align` pair, with no niche.
     ///
     /// Layout is same for both 64-bit and 32-bit platforms.