[ty] Shrink size of AstNodeRef (#20028)

## Summary

Removes the `module_ptr` field from `AstNodeRef` in release mode, and
change `NodeIndex` to a `NonZeroU32` to reduce the size of
`Option<AstNodeRef<_>>` fields.

I believe CI runs in debug mode, so this won't show up in the memory
report, but this reduces memory by ~2% in release mode.

crates/ruff_python_ast/src/node_index.rs

@@ -1,3 +1,4 @@
+use std::num::NonZeroU32;
 use std::sync::atomic::{AtomicU32, Ordering};
 
 /// An AST node that has an index.
@@ -16,64 +17,82 @@ where
 }
 
 /// A unique index for a node within an AST.
-///
-/// This type is interiorly mutable to allow assigning node indices
-/// on-demand after parsing.
-#[derive(Default)]
+#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
 #[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
-pub struct AtomicNodeIndex(AtomicU32);
-
-impl AtomicNodeIndex {
-    /// Returns a placeholder `AtomicNodeIndex`.
-    pub const fn dummy() -> AtomicNodeIndex {
-        AtomicNodeIndex(AtomicU32::new(u32::MAX))
-    }
-
-    /// Load the current value of the `AtomicNodeIndex`.
-    pub fn load(&self) -> NodeIndex {
-        NodeIndex(self.0.load(Ordering::Relaxed))
-    }
-
-    /// Set the value of the `AtomicNodeIndex`.
-    pub fn set(&self, value: u32) {
-        self.0.store(value, Ordering::Relaxed);
-    }
-}
-
-/// A unique index for a node within an AST.
-#[derive(PartialEq, Eq, Debug, PartialOrd, Ord, Clone, Copy, Hash)]
-#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
-pub struct NodeIndex(u32);
+pub struct NodeIndex(NonZeroU32);
 
 impl NodeIndex {
-    pub fn as_usize(self) -> usize {
-        self.0 as _
-    }
+    /// A placeholder `NodeIndex`.
+    pub const NONE: NodeIndex = NodeIndex(NonZeroU32::new(NodeIndex::_NONE).unwrap());
 
-    pub fn as_u32(self) -> u32 {
-        self.0
+    // Note that the index `u32::MAX` is reserved for the `NonZeroU32` niche, and
+    // this placeholder also reserves the second highest index.
+    const _NONE: u32 = u32::MAX - 1;
+
+    /// Returns the index as a `u32`. or `None` for `NodeIndex::NONE`.
+    pub fn as_u32(self) -> Option<u32> {
+        if self == NodeIndex::NONE {
+            None
+        } else {
+            Some(self.0.get() - 1)
+        }
     }
 }
 
 impl From<u32> for NodeIndex {
     fn from(value: u32) -> Self {
-        NodeIndex(value)
+        match NonZeroU32::new(value + 1).map(NodeIndex) {
+            None | Some(NodeIndex::NONE) => panic!("exceeded maximum `NodeIndex`"),
+            Some(index) => index,
+        }
     }
 }
 
-impl From<u32> for AtomicNodeIndex {
-    fn from(value: u32) -> Self {
-        AtomicNodeIndex(AtomicU32::from(value))
+impl std::fmt::Debug for NodeIndex {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if *self == Self::NONE {
+            f.debug_tuple("NodeIndex(None)").finish()
+        } else {
+            f.debug_tuple("NodeIndex").field(&self.0).finish()
+        }
+    }
+}
+
+/// A unique index for a node within an AST.
+///
+/// This type is interiorly mutable to allow assigning node indices
+/// on-demand after parsing.
+#[cfg_attr(feature = "get-size", derive(get_size2::GetSize))]
+pub struct AtomicNodeIndex(AtomicU32);
+
+#[allow(clippy::declare_interior_mutable_const)]
+impl AtomicNodeIndex {
+    /// A placeholder `AtomicNodeIndex`.
+    pub const NONE: AtomicNodeIndex = AtomicNodeIndex(AtomicU32::new(NodeIndex::_NONE));
+
+    /// Load the current value of the `AtomicNodeIndex`.
+    pub fn load(&self) -> NodeIndex {
+        let index = NonZeroU32::new(self.0.load(Ordering::Relaxed))
+            .expect("value stored was a valid `NodeIndex`");
+
+        NodeIndex(index)
+    }
+
+    /// Set the value of the `AtomicNodeIndex`.
+    pub fn set(&self, index: NodeIndex) {
+        self.0.store(index.0.get(), Ordering::Relaxed);
+    }
+}
+
+impl Default for AtomicNodeIndex {
+    fn default() -> Self {
+        Self::NONE
     }
 }
 
 impl std::fmt::Debug for AtomicNodeIndex {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        if *self == AtomicNodeIndex::dummy() {
-            f.debug_tuple("AtomicNodeIndex").finish_non_exhaustive()
-        } else {
-            f.debug_tuple("AtomicNodeIndex").field(&self.0).finish()
-        }
+        std::fmt::Debug::fmt(&self.load(), f)
     }
 }
 
@@ -108,3 +127,26 @@ impl Clone for AtomicNodeIndex {
         Self(AtomicU32::from(self.0.load(Ordering::Relaxed)))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{AtomicNodeIndex, NodeIndex};
+
+    #[test]
+    fn test_node_index() {
+        let index = AtomicNodeIndex::NONE;
+
+        assert_eq!(index.load(), NodeIndex::NONE);
+        assert_eq!(format!("{index:?}"), "NodeIndex(None)");
+
+        index.set(NodeIndex::from(1));
+        assert_eq!(index.load(), NodeIndex::from(1));
+        assert_eq!(index.load().as_u32(), Some(1));
+
+        let index = NodeIndex::from(0);
+        assert_eq!(index.as_u32(), Some(0));
+
+        let index = NodeIndex::NONE;
+        assert_eq!(index.as_u32(), None);
+    }
+}