[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/helpers.rs

@@ -1489,7 +1489,7 @@ pub fn pep_604_optional(expr: &Expr) -> Expr {
         op: Operator::BitOr,
         right: Box::new(Expr::NoneLiteral(ast::ExprNoneLiteral::default())),
         range: TextRange::default(),
-        node_index: AtomicNodeIndex::dummy(),
+        node_index: AtomicNodeIndex::NONE,
     }
     .into()
 }
@@ -1501,7 +1501,7 @@ pub fn pep_604_union(elts: &[Expr]) -> Expr {
             elts: vec![],
             ctx: ExprContext::Load,
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             parenthesized: true,
         }),
         [Expr::Tuple(ast::ExprTuple { elts, .. })] => pep_604_union(elts),
@@ -1511,7 +1511,7 @@ pub fn pep_604_union(elts: &[Expr]) -> Expr {
             op: Operator::BitOr,
             right: Box::new(pep_604_union(std::slice::from_ref(elt))),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
         }),
     }
 }
@@ -1522,13 +1522,13 @@ pub fn typing_optional(elt: Expr, binding: Name) -> Expr {
         value: Box::new(Expr::Name(ast::ExprName {
             id: binding,
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             ctx: ExprContext::Load,
         })),
         slice: Box::new(elt),
         ctx: ExprContext::Load,
         range: TextRange::default(),
-        node_index: AtomicNodeIndex::dummy(),
+        node_index: AtomicNodeIndex::NONE,
     })
 }
 
@@ -1541,19 +1541,19 @@ pub fn typing_union(elts: &[Expr], binding: Name) -> Expr {
         value: Box::new(Expr::Name(ast::ExprName {
             id: binding,
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             ctx: ExprContext::Load,
         })),
         slice: Box::new(Expr::Tuple(ast::ExprTuple {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             elts: elts.to_vec(),
             ctx: ExprContext::Load,
             parenthesized: false,
         })),
         ctx: ExprContext::Load,
         range: TextRange::default(),
-        node_index: AtomicNodeIndex::dummy(),
+        node_index: AtomicNodeIndex::NONE,
     })
 }
 
@@ -1686,34 +1686,34 @@ mod tests {
         let name = Expr::Name(ExprName {
             id: "x".into(),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             ctx: ExprContext::Load,
         });
         let constant_one = Expr::NumberLiteral(ExprNumberLiteral {
             value: Number::Int(Int::from(1u8)),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
         });
         let constant_two = Expr::NumberLiteral(ExprNumberLiteral {
             value: Number::Int(Int::from(2u8)),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
         });
         let constant_three = Expr::NumberLiteral(ExprNumberLiteral {
             value: Number::Int(Int::from(3u8)),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
         });
         let type_var_one = TypeParam::TypeVar(TypeParamTypeVar {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             bound: Some(Box::new(constant_one.clone())),
             default: None,
             name: Identifier::new("x", TextRange::default()),
         });
         let type_var_two = TypeParam::TypeVar(TypeParamTypeVar {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             bound: None,
             default: Some(Box::new(constant_two.clone())),
             name: Identifier::new("x", TextRange::default()),
@@ -1723,11 +1723,11 @@ mod tests {
             type_params: Some(Box::new(TypeParams {
                 type_params: vec![type_var_one, type_var_two],
                 range: TextRange::default(),
-                node_index: AtomicNodeIndex::dummy(),
+                node_index: AtomicNodeIndex::NONE,
             })),
             value: Box::new(constant_three.clone()),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
         });
         assert!(!any_over_stmt(&type_alias, &|expr| {
             seen.borrow_mut().push(expr.clone());
@@ -1743,7 +1743,7 @@ mod tests {
     fn any_over_type_param_type_var() {
         let type_var_no_bound = TypeParam::TypeVar(TypeParamTypeVar {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             bound: None,
             default: None,
             name: Identifier::new("x", TextRange::default()),
@@ -1753,12 +1753,12 @@ mod tests {
         let constant = Expr::NumberLiteral(ExprNumberLiteral {
             value: Number::Int(Int::ONE),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
         });
 
         let type_var_with_bound = TypeParam::TypeVar(TypeParamTypeVar {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             bound: Some(Box::new(constant.clone())),
             default: None,
             name: Identifier::new("x", TextRange::default()),
@@ -1776,7 +1776,7 @@ mod tests {
 
         let type_var_with_default = TypeParam::TypeVar(TypeParamTypeVar {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             default: Some(Box::new(constant.clone())),
             bound: None,
             name: Identifier::new("x", TextRange::default()),
@@ -1797,7 +1797,7 @@ mod tests {
     fn any_over_type_param_type_var_tuple() {
         let type_var_tuple = TypeParam::TypeVarTuple(TypeParamTypeVarTuple {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             name: Identifier::new("x", TextRange::default()),
             default: None,
         });
@@ -1809,12 +1809,12 @@ mod tests {
         let constant = Expr::NumberLiteral(ExprNumberLiteral {
             value: Number::Int(Int::ONE),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
         });
 
         let type_var_tuple_with_default = TypeParam::TypeVarTuple(TypeParamTypeVarTuple {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             default: Some(Box::new(constant.clone())),
             name: Identifier::new("x", TextRange::default()),
         });
@@ -1834,7 +1834,7 @@ mod tests {
     fn any_over_type_param_param_spec() {
         let type_param_spec = TypeParam::ParamSpec(TypeParamParamSpec {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             name: Identifier::new("x", TextRange::default()),
             default: None,
         });
@@ -1846,12 +1846,12 @@ mod tests {
         let constant = Expr::NumberLiteral(ExprNumberLiteral {
             value: Number::Int(Int::ONE),
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
         });
 
         let param_spec_with_default = TypeParam::TypeVarTuple(TypeParamTypeVarTuple {
             range: TextRange::default(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             default: Some(Box::new(constant.clone())),
             name: Identifier::new("x", TextRange::default()),
         });

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);
+    }
+}

crates/ruff_python_ast/src/nodes.rs

@@ -1602,7 +1602,7 @@ impl StringLiteral {
         Self {
             range,
             value: "".into(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             flags: StringLiteralFlags::empty().with_invalid(),
         }
     }
@@ -1622,7 +1622,7 @@ impl From<StringLiteral> for Expr {
     fn from(payload: StringLiteral) -> Self {
         ExprStringLiteral {
             range: payload.range,
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             value: StringLiteralValue::single(payload),
         }
         .into()
@@ -2001,7 +2001,7 @@ impl BytesLiteral {
         Self {
             range,
             value: Box::new([]),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             flags: BytesLiteralFlags::empty().with_invalid(),
         }
     }
@@ -2011,7 +2011,7 @@ impl From<BytesLiteral> for Expr {
     fn from(payload: BytesLiteral) -> Self {
         ExprBytesLiteral {
             range: payload.range,
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             value: BytesLiteralValue::single(payload),
         }
         .into()
@@ -3545,7 +3545,7 @@ impl Identifier {
     pub fn new(id: impl Into<Name>, range: TextRange) -> Self {
         Self {
             id: id.into(),
-            node_index: AtomicNodeIndex::dummy(),
+            node_index: AtomicNodeIndex::NONE,
             range,
         }
     }