feat(traverse): implement GetAddress for Ancestor

crates/oxc_allocator/src/address.rs

@@ -3,13 +3,30 @@ use std::ptr;
 use crate::Box;
 
 /// Memory address of an AST node in arena.
-///
-/// `Address` is generated from a `Box<T>`.
-/// AST nodes in a `Box` in an arena are guaranteed to never move in memory,
-/// so this address acts as a unique identifier for the duration of the arena's existence.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct Address(usize);
 
+impl Address {
+    /// Dummy address.
+    ///
+    /// Never equal to any real `Address`, but is equal to itself.
+    pub const DUMMY: Self = Self(0);
+
+    /// Get the memory address of a pointer to an AST node in arena.
+    ///
+    /// The pointer must point to an AST node in the arena (not on the stack),
+    /// or the returned `Address` will be meaningless.
+    ///
+    /// If the AST node is in a `Box`, the address is guaranteed to be a unique identifier
+    /// for the duration of the arena's existence.
+    /// If the node is in a `Vec`, then the `Address` may not remain accurate if the `Vec`
+    /// is resized or has elements added or removed before this node.
+    #[inline]
+    pub fn from_ptr<T>(p: *const T) -> Self {
+        Self(p as usize)
+    }
+}
+
 /// Trait for getting the memory address of an AST node.
 pub trait GetAddress {
     /// Get the memory address of a value allocated in the arena.
@@ -18,8 +35,11 @@ pub trait GetAddress {
 
 impl<'a, T> GetAddress for Box<'a, T> {
     /// Get the memory address of a value allocated in the arena.
+    ///
+    /// AST nodes in a `Box` in an arena are guaranteed to never move in memory,
+    /// so this address acts as a unique identifier for the duration of the arena's existence.
     #[inline]
     fn address(&self) -> Address {
-        Address(ptr::addr_of!(**self) as usize)
+        Address::from_ptr(ptr::addr_of!(**self))
     }
 }

crates/oxc_traverse/scripts/lib/ancestor.mjs

@@ -9,6 +9,7 @@ export default function generateAncestorsCode(types) {
     ancestorEnumVariants = '',
     isFunctions = '',
     ancestorTypes = '',
+    addressMatchArms = '',
     discriminant = 1;
   for (const type of Object.values(types)) {
     if (type.kind === 'enum') continue;
@@ -62,6 +63,13 @@ export default function generateAncestorsCode(types) {
         impl${lifetimes} ${structName} {
           ${methodsCode}
         }
+
+        impl${lifetimes} GetAddress for ${structName} {
+          #[inline]
+          fn address(&self) -> Address {
+            Address::from_ptr(self.0)
+          }
+        }
       `;
 
       const variantName = `${type.name}${fieldNameCamel}`;
@@ -75,6 +83,8 @@ export default function generateAncestorsCode(types) {
         (variantNamesForEnums[fieldTypeName] || (variantNamesForEnums[fieldTypeName] = []))
           .push(variantName);
       }
+
+      addressMatchArms += `Self::${variantName}(a) => a.address(),\n`;
     }
 
     if (variantNames.length > 0) {
@@ -114,7 +124,7 @@ export default function generateAncestorsCode(types) {
 
     use memoffset::offset_of;
 
-    use oxc_allocator::{Box, Vec};
+    use oxc_allocator::{Address, Box, GetAddress, Vec};
     use oxc_ast::ast::*;
     use oxc_syntax::scope::ScopeId;
 
@@ -157,6 +167,18 @@ export default function generateAncestorsCode(types) {
       ${isFunctions}
     }
 
+    impl<'a, 't> GetAddress for Ancestor<'a, 't> {
+      /// Get memory address of node represented by \`Ancestor\` in the arena.
+      // Compiler should reduce this down to only a couple of assembly operations.
+      #[inline]
+      fn address(&self) -> Address {
+        match self {
+          Self::None => Address::DUMMY,
+          ${addressMatchArms}
+        }
+      }
+    }
+
     ${ancestorTypes}
   `;
 }