feat(allocator): add Address::from_ref method

crates/oxc_allocator/src/address.rs

@@ -1,4 +1,4 @@
-use std::ptr;
+use std::ptr::{NonNull, addr_of};
 
 use crate::Box;
 
@@ -72,6 +72,79 @@ impl Address {
     pub fn from_ptr<T>(p: *const T) -> Self {
         Self(p as usize)
     }
+
+    /// Get the memory address of a reference to an AST node in arena.
+    ///
+    /// **This method is an escape hatch only.**
+    /// Prefer using `GetAddress::address` instead, because it is more likely to produce a stable `Address`.
+    ///
+    /// 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.
+    ///
+    /// But 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.
+    ///
+    /// The reference must point to an AST node in the arena (not on the stack), or the returned `Address`
+    /// will be meaningless. Be careful not to pass a double-reference to `from_ref`, or the resulting `Address`
+    /// will point to the reference itself, instead of the thing being referenced.
+    ///
+    /// ```ignore
+    /// impl<'a> Visit<'a> for MyVisitor {
+    ///     fn visit_identifier_reference(&mut self, ident: &IdentifierReference<'a>) {
+    ///         // Correct - `address` is address of the `IdentifierReference`
+    ///         let address = Address::from_ref(ident);
+    ///         // WRONG - `address` is address of `&IdentifierReference` reference itself, which is on the stack
+    ///         let address = Address::from_ref(&ident);
+    ///     }
+    /// }
+    /// ```
+    ///
+    /// # Example
+    ///
+    /// Demonstration of the difference between `Address::from_ref` and `GetAddress::address`:
+    ///
+    /// ```ignore
+    /// use oxc_allocator::{Address, GetAddress, Vec};
+    /// use oxc_span::SPAN;
+    ///
+    /// // Create a `Vec<Statement>` containing a single `BlockStatement`
+    /// let mut stmts = Vec::with_capacity_in(1, &allocator);
+    /// stmts.push(ast_builder.statement_block(SPAN, Vec::new_in(&allocator)));
+    ///
+    /// let block_address = stmts[0].address();
+    /// let stmt_address = Address::from_ref(&stmts[0]);
+    ///
+    /// // Add another `Statement` to the `Vec`.
+    /// // This causes the `Vec` to grow and reallocate.
+    /// stmts.push(ast_builder.statement_empty(SPAN));
+    ///
+    /// let block_address_after_push = stmts[0].address();
+    /// let stmt_address_after_push = Address::from_ref(&stmts[0]);
+    ///
+    /// // Address of the `BlockStatement` is unchanged
+    /// // (because the `Box`'s pointer still points to same memory location)
+    /// assert!(block_address_after_push == block_address);
+    /// // Address of the `Statement` has changed
+    /// // (because the `Vec` reallocated, so its contents have moved in memory)
+    /// assert!(stmt_address_after_push != stmt_address);
+    ///
+    /// // Insert a new `Statement` at start of the `Vec`.
+    /// // The `BlockStatement` is now at index 1.
+    /// stmts.insert(0, ast_builder.statement_empty(SPAN));
+    ///
+    /// let block_address_after_insert = stmts[1].address();
+    /// let stmt_address_after_insert = Address::from_ref(&stmts[1]);
+    ///
+    /// // Address of the `BlockStatement` is still unchanged
+    /// assert!(block_address_after_insert == block_address_after_push);
+    /// // Address of the `Statement` has changed again
+    /// assert!(stmt_address_after_insert != stmt_address_after_push);
+    /// ```
+    #[inline]
+    pub fn from_ref<T>(r: &T) -> Self {
+        let p = NonNull::from_ref(r);
+        Self(p.addr().get())
+    }
 }
 
 /// Trait for getting the memory address of an AST node.
@@ -87,7 +160,7 @@ impl<T> GetAddress for Box<'_, T> {
     /// so this address acts as a unique identifier for the duration of the arena's existence.
     #[inline]
     fn address(&self) -> Address {
-        Address::from_ptr(ptr::addr_of!(**self))
+        Address::from_ptr(addr_of!(**self))
     }
 }
 

crates/oxc_ast/src/ast_kind_impl.rs

@@ -277,7 +277,7 @@ impl<'a> AstKind<'a> {
             // - not the default value e.g. `({ assignee = ident } = obj)`.
             AstKind::AssignmentTargetPropertyIdentifier(assign_target) => {
                 let binding = &assign_target.binding;
-                Address::from_ptr(binding) == self.address()
+                Address::from_ref(binding) == self.address()
             }
             // `({ prop: ident } = obj)`
             // Only match if ident is the assignee
@@ -695,9 +695,9 @@ impl GetAddress for MemberExpressionKind<'_> {
     #[inline] // This should boil down to a single instruction
     fn address(&self) -> Address {
         match *self {
-            Self::Computed(member_expr) => Address::from_ptr(member_expr),
-            Self::Static(member_expr) => Address::from_ptr(member_expr),
-            Self::PrivateField(member_expr) => Address::from_ptr(member_expr),
+            Self::Computed(member_expr) => Address::from_ref(member_expr),
+            Self::Static(member_expr) => Address::from_ref(member_expr),
+            Self::PrivateField(member_expr) => Address::from_ref(member_expr),
         }
     }
 }
@@ -751,12 +751,12 @@ impl GetAddress for ModuleDeclarationKind<'_> {
     #[inline] // This should boil down to a single instruction
     fn address(&self) -> Address {
         match *self {
-            Self::Import(decl) => Address::from_ptr(decl),
-            Self::ExportAll(decl) => Address::from_ptr(decl),
-            Self::ExportNamed(decl) => Address::from_ptr(decl),
-            Self::ExportDefault(decl) => Address::from_ptr(decl),
-            Self::TSExportAssignment(decl) => Address::from_ptr(decl),
-            Self::TSNamespaceExport(decl) => Address::from_ptr(decl),
+            Self::Import(decl) => Address::from_ref(decl),
+            Self::ExportAll(decl) => Address::from_ref(decl),
+            Self::ExportNamed(decl) => Address::from_ref(decl),
+            Self::ExportDefault(decl) => Address::from_ref(decl),
+            Self::TSExportAssignment(decl) => Address::from_ref(decl),
+            Self::TSNamespaceExport(decl) => Address::from_ref(decl),
         }
     }
 }
@@ -784,8 +784,8 @@ impl GetAddress for PropertyKeyKind<'_> {
     #[inline] // This should boil down to a single instruction
     fn address(&self) -> Address {
         match *self {
-            Self::Static(ident) => Address::from_ptr(ident),
-            Self::Private(ident) => Address::from_ptr(ident),
+            Self::Static(ident) => Address::from_ref(ident),
+            Self::Private(ident) => Address::from_ref(ident),
         }
     }
 }