feat(allocator, ast): introduce UnstableAddress trait

crates/oxc_allocator/src/address.rs

@@ -19,7 +19,9 @@ impl Address {
     /// Get the memory address of a pointer 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`].
+    /// Prefer using [`GetAddress::address`] or [`UnstableAddress::unstable_address`] instead,
+    /// because they are more likely to produce a stable [`Address`].
+    /// (Yes even `unstable_address` is more likely to produce a stable `Address` than this function!)
     ///
     /// 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.
@@ -97,7 +99,9 @@ impl Address {
     /// 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`].
+    /// Prefer using [`GetAddress::address`] or [`UnstableAddress::unstable_address`] instead,
+    /// because they are more likely to produce a stable `Address`.
+    /// (Yes even `unstable_address` is more likely to produce a stable `Address` than this function!)
     ///
     /// 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.
@@ -196,3 +200,104 @@ impl GetAddress for Address {
         *self
     }
 }
+
+/// Trait for getting the memory address of an AST node which is not necessarily stable.
+///
+/// See [`UnstableAddress::unstable_address`] for more details.
+///
+/// This trait is implemented for all AST struct types.
+///
+/// *DO NOT* implement this trait on any other type.
+pub trait UnstableAddress {
+    /// Get the memory [`Address`] of a reference to an AST node in arena, which is not necessarily stable.
+    ///
+    /// # Stable addresses
+    ///
+    /// It's ideal to obtain an `Address` for an AST node which is guaranteed to be stable for the life of the AST.
+    ///
+    /// Then you can reliably compare two `Address`es to determine if they refer to the same node,
+    /// without any risk of the result being wrong because one or other of the nodes has moved in memory.
+    ///
+    /// Types which have a stable address:
+    /// * [`Box<T>`]
+    /// * AST enums where all variants are `Box`es (e.g. `Statement`, `Expression`)
+    ///
+    /// Some other types guarantee stability for a shorter period of time:
+    /// * `oxc_ast::AstKind` - guaranteed stable address while `Semantic` is alive.
+    /// * `oxc_traverse::Ancestor` - guaranteed stable address while traversing descendents of the ancestor.
+    ///
+    /// The above types all implement [`GetAddress::address`]. If you have access to one of these types,
+    /// it's better to use `GetAddress::address` instead of this method.
+    ///
+    /// # Why this method exists
+    ///
+    /// Sometimes you only have access to a reference to an AST node, but you know from context that the node
+    /// will not move in memory during the time you need the `Address` to remain stable.
+    ///
+    /// You can use this method in such cases, but you need to be careful. For correct behavior, you must ensure
+    /// yourself that the node will not move in memory during the time you need the `Address` to remain accurate.
+    ///
+    /// When using this method, it is recommended to make a comment at the call site, explaining how you can prove
+    /// that the `Address` will remain accurate (the AST node will not move) for the time period that you need it to be.
+    ///
+    /// # When a type does not move in memory
+    ///
+    /// If the AST is immutable (e.g. in `Visit` trait), then any node in the AST is statically positioned in memory.
+    /// Therefore, in the linter, any reference to an AST node is guaranteed to have a stable `Address`.
+    ///
+    /// If an AST node is in `Vec`, then it'll remain at same memory address, but only as long as the `Vec` does
+    /// not reallocate (e.g. by `Vec::push`, `Vec::extend`).
+    ///
+    /// This method will return a stable `Address` for any AST node in such circumstances.
+    ///
+    /// # Common pitfalls
+    ///
+    /// ## References to AST nodes on the stack
+    ///
+    /// Ensure that the reference passed to this method is to a node which is in the arena, *not* on the stack.
+    ///
+    /// ```ignore
+    /// let binary_expr: BinaryExpression<'a> = get_owned_binary_expression_somehow();
+    /// // WRONG: `binary_expr` is on the stack, so the `Address` will be meaningless
+    /// let address = binary_expr.unstable_address();
+    /// // More correct: `binary_expr` is in the arena.
+    /// // Will have a stable `Address` as long as `vec` does not reallocate.
+    /// let mut vec = Vec::new_in(&allocator);
+    /// vec.push(binary_expr);
+    /// let address = vec[0].unstable_address();
+    /// ```
+    ///
+    /// ## AST nodes in `Vec`s
+    ///
+    /// ```ignore
+    /// let mut vec: &mut Vec<BinaryExpression<'a>> = get_vec_somehow();
+    ///
+    /// let address = vec[0].unstable_address();
+    /// vec.push(get_owned_binary_expression_somehow());
+    /// let address_after_push = vec[0].unstable_address();
+    ///
+    /// // This assertion may or may not pass, depending on whether `push` caused the `Vec` to reallocate.
+    /// // This depends on whether `vec` had spare capacity or not, prior to the `push` call.
+    /// assert!(address_after_push == address);
+    /// ```
+    ///
+    /// # Guardrails
+    ///
+    /// This method is less error-prone than [`Address::from_ptr`] and [`Address::from_ref`],
+    /// because it provides a few guardrails:
+    ///
+    /// * [`UnstableAddress`] is only implemented on AST struct types, so you can't call it on a type which
+    ///   it doesn't make sense to get the `Address` of.
+    ///
+    /// * You don't need to worry about passing it a double-reference (`&&T`), because this method will automatically
+    ///   dereference as required.
+    ///
+    /// Even with these guardrails, usage of this method still requires care, for the reasons discussed above.
+    ///
+    /// [`Box<T>`]: crate::Box
+    #[inline(always)] // Because it's a no-op
+    fn unstable_address(&self) -> Address {
+        let p = NonNull::from_ref(self);
+        Address(p.addr().get())
+    }
+}

crates/oxc_allocator/src/lib.rs

@@ -62,7 +62,7 @@ mod vec;
 mod vec2;
 
 pub use accessor::AllocatorAccessor;
-pub use address::{Address, GetAddress};
+pub use address::{Address, GetAddress, UnstableAddress};
 pub use allocator::Allocator;
 #[cfg(feature = "bitset")]
 pub use bitset::BitSet;