feat(ast, semantic): implement GetAddress for AstKind and AstNode (#11758)

Implement `GetAddress` on `AstKind` and `AstNode`. This should make it a bit easier to migrate away from comparing spans to find out if 2 AST nodes are the same.

This was unblocked by #11684 (thank you @ulrichstark).

Author: overlookmotel

crates/oxc_ast/src/generated/ast_kind.rs

@@ -6,6 +6,7 @@
 
 use std::ptr;
 
+use oxc_allocator::{Address, GetAddress};
 use oxc_span::{GetSpan, Span};
 
 use crate::ast::*;
@@ -598,6 +599,194 @@ impl GetSpan for AstKind<'_> {
     }
 }
 
+impl GetAddress for AstKind<'_> {
+    fn address(&self) -> Address {
+        match *self {
+            Self::Program(it) => Address::from_ptr(it),
+            Self::IdentifierName(it) => Address::from_ptr(it),
+            Self::IdentifierReference(it) => Address::from_ptr(it),
+            Self::BindingIdentifier(it) => Address::from_ptr(it),
+            Self::LabelIdentifier(it) => Address::from_ptr(it),
+            Self::ThisExpression(it) => Address::from_ptr(it),
+            Self::ArrayExpression(it) => Address::from_ptr(it),
+            Self::Elision(it) => Address::from_ptr(it),
+            Self::ObjectExpression(it) => Address::from_ptr(it),
+            Self::ObjectProperty(it) => Address::from_ptr(it),
+            Self::PropertyKey(it) => it.address(),
+            Self::TemplateLiteral(it) => Address::from_ptr(it),
+            Self::TaggedTemplateExpression(it) => Address::from_ptr(it),
+            Self::MemberExpression(it) => it.address(),
+            Self::CallExpression(it) => Address::from_ptr(it),
+            Self::NewExpression(it) => Address::from_ptr(it),
+            Self::MetaProperty(it) => Address::from_ptr(it),
+            Self::SpreadElement(it) => Address::from_ptr(it),
+            Self::Argument(it) => it.address(),
+            Self::UpdateExpression(it) => Address::from_ptr(it),
+            Self::UnaryExpression(it) => Address::from_ptr(it),
+            Self::BinaryExpression(it) => Address::from_ptr(it),
+            Self::PrivateInExpression(it) => Address::from_ptr(it),
+            Self::LogicalExpression(it) => Address::from_ptr(it),
+            Self::ConditionalExpression(it) => Address::from_ptr(it),
+            Self::AssignmentExpression(it) => Address::from_ptr(it),
+            Self::AssignmentTarget(it) => it.address(),
+            Self::SimpleAssignmentTarget(it) => it.address(),
+            Self::AssignmentTargetPattern(it) => it.address(),
+            Self::ArrayAssignmentTarget(it) => Address::from_ptr(it),
+            Self::ObjectAssignmentTarget(it) => Address::from_ptr(it),
+            Self::AssignmentTargetWithDefault(it) => Address::from_ptr(it),
+            Self::SequenceExpression(it) => Address::from_ptr(it),
+            Self::Super(it) => Address::from_ptr(it),
+            Self::AwaitExpression(it) => Address::from_ptr(it),
+            Self::ChainExpression(it) => Address::from_ptr(it),
+            Self::ParenthesizedExpression(it) => Address::from_ptr(it),
+            Self::Directive(it) => Address::from_ptr(it),
+            Self::Hashbang(it) => Address::from_ptr(it),
+            Self::BlockStatement(it) => Address::from_ptr(it),
+            Self::VariableDeclaration(it) => Address::from_ptr(it),
+            Self::VariableDeclarator(it) => Address::from_ptr(it),
+            Self::EmptyStatement(it) => Address::from_ptr(it),
+            Self::ExpressionStatement(it) => Address::from_ptr(it),
+            Self::IfStatement(it) => Address::from_ptr(it),
+            Self::DoWhileStatement(it) => Address::from_ptr(it),
+            Self::WhileStatement(it) => Address::from_ptr(it),
+            Self::ForStatement(it) => Address::from_ptr(it),
+            Self::ForInStatement(it) => Address::from_ptr(it),
+            Self::ForOfStatement(it) => Address::from_ptr(it),
+            Self::ContinueStatement(it) => Address::from_ptr(it),
+            Self::BreakStatement(it) => Address::from_ptr(it),
+            Self::ReturnStatement(it) => Address::from_ptr(it),
+            Self::WithStatement(it) => Address::from_ptr(it),
+            Self::SwitchStatement(it) => Address::from_ptr(it),
+            Self::SwitchCase(it) => Address::from_ptr(it),
+            Self::LabeledStatement(it) => Address::from_ptr(it),
+            Self::ThrowStatement(it) => Address::from_ptr(it),
+            Self::TryStatement(it) => Address::from_ptr(it),
+            Self::CatchClause(it) => Address::from_ptr(it),
+            Self::CatchParameter(it) => Address::from_ptr(it),
+            Self::DebuggerStatement(it) => Address::from_ptr(it),
+            Self::AssignmentPattern(it) => Address::from_ptr(it),
+            Self::ObjectPattern(it) => Address::from_ptr(it),
+            Self::ArrayPattern(it) => Address::from_ptr(it),
+            Self::BindingRestElement(it) => Address::from_ptr(it),
+            Self::Function(it) => Address::from_ptr(it),
+            Self::FormalParameters(it) => Address::from_ptr(it),
+            Self::FormalParameter(it) => Address::from_ptr(it),
+            Self::FunctionBody(it) => Address::from_ptr(it),
+            Self::ArrowFunctionExpression(it) => Address::from_ptr(it),
+            Self::YieldExpression(it) => Address::from_ptr(it),
+            Self::Class(it) => Address::from_ptr(it),
+            Self::ClassBody(it) => Address::from_ptr(it),
+            Self::MethodDefinition(it) => Address::from_ptr(it),
+            Self::PropertyDefinition(it) => Address::from_ptr(it),
+            Self::PrivateIdentifier(it) => Address::from_ptr(it),
+            Self::StaticBlock(it) => Address::from_ptr(it),
+            Self::ModuleDeclaration(it) => it.address(),
+            Self::AccessorProperty(it) => Address::from_ptr(it),
+            Self::ImportExpression(it) => Address::from_ptr(it),
+            Self::ImportDeclaration(it) => Address::from_ptr(it),
+            Self::ImportSpecifier(it) => Address::from_ptr(it),
+            Self::ImportDefaultSpecifier(it) => Address::from_ptr(it),
+            Self::ImportNamespaceSpecifier(it) => Address::from_ptr(it),
+            Self::WithClause(it) => Address::from_ptr(it),
+            Self::ImportAttribute(it) => Address::from_ptr(it),
+            Self::ExportNamedDeclaration(it) => Address::from_ptr(it),
+            Self::ExportDefaultDeclaration(it) => Address::from_ptr(it),
+            Self::ExportAllDeclaration(it) => Address::from_ptr(it),
+            Self::ExportSpecifier(it) => Address::from_ptr(it),
+            Self::V8IntrinsicExpression(it) => Address::from_ptr(it),
+            Self::BooleanLiteral(it) => Address::from_ptr(it),
+            Self::NullLiteral(it) => Address::from_ptr(it),
+            Self::NumericLiteral(it) => Address::from_ptr(it),
+            Self::StringLiteral(it) => Address::from_ptr(it),
+            Self::BigIntLiteral(it) => Address::from_ptr(it),
+            Self::RegExpLiteral(it) => Address::from_ptr(it),
+            Self::JSXElement(it) => Address::from_ptr(it),
+            Self::JSXOpeningElement(it) => Address::from_ptr(it),
+            Self::JSXClosingElement(it) => Address::from_ptr(it),
+            Self::JSXFragment(it) => Address::from_ptr(it),
+            Self::JSXOpeningFragment(it) => Address::from_ptr(it),
+            Self::JSXClosingFragment(it) => Address::from_ptr(it),
+            Self::JSXNamespacedName(it) => Address::from_ptr(it),
+            Self::JSXMemberExpression(it) => Address::from_ptr(it),
+            Self::JSXExpressionContainer(it) => Address::from_ptr(it),
+            Self::JSXEmptyExpression(it) => Address::from_ptr(it),
+            Self::JSXAttribute(it) => Address::from_ptr(it),
+            Self::JSXSpreadAttribute(it) => Address::from_ptr(it),
+            Self::JSXIdentifier(it) => Address::from_ptr(it),
+            Self::JSXSpreadChild(it) => Address::from_ptr(it),
+            Self::JSXText(it) => Address::from_ptr(it),
+            Self::TSThisParameter(it) => Address::from_ptr(it),
+            Self::TSEnumDeclaration(it) => Address::from_ptr(it),
+            Self::TSEnumBody(it) => Address::from_ptr(it),
+            Self::TSEnumMember(it) => Address::from_ptr(it),
+            Self::TSTypeAnnotation(it) => Address::from_ptr(it),
+            Self::TSLiteralType(it) => Address::from_ptr(it),
+            Self::TSConditionalType(it) => Address::from_ptr(it),
+            Self::TSUnionType(it) => Address::from_ptr(it),
+            Self::TSIntersectionType(it) => Address::from_ptr(it),
+            Self::TSParenthesizedType(it) => Address::from_ptr(it),
+            Self::TSTypeOperator(it) => Address::from_ptr(it),
+            Self::TSArrayType(it) => Address::from_ptr(it),
+            Self::TSIndexedAccessType(it) => Address::from_ptr(it),
+            Self::TSTupleType(it) => Address::from_ptr(it),
+            Self::TSNamedTupleMember(it) => Address::from_ptr(it),
+            Self::TSOptionalType(it) => Address::from_ptr(it),
+            Self::TSRestType(it) => Address::from_ptr(it),
+            Self::TSAnyKeyword(it) => Address::from_ptr(it),
+            Self::TSStringKeyword(it) => Address::from_ptr(it),
+            Self::TSBooleanKeyword(it) => Address::from_ptr(it),
+            Self::TSNumberKeyword(it) => Address::from_ptr(it),
+            Self::TSNeverKeyword(it) => Address::from_ptr(it),
+            Self::TSIntrinsicKeyword(it) => Address::from_ptr(it),
+            Self::TSUnknownKeyword(it) => Address::from_ptr(it),
+            Self::TSNullKeyword(it) => Address::from_ptr(it),
+            Self::TSUndefinedKeyword(it) => Address::from_ptr(it),
+            Self::TSVoidKeyword(it) => Address::from_ptr(it),
+            Self::TSSymbolKeyword(it) => Address::from_ptr(it),
+            Self::TSThisType(it) => Address::from_ptr(it),
+            Self::TSObjectKeyword(it) => Address::from_ptr(it),
+            Self::TSBigIntKeyword(it) => Address::from_ptr(it),
+            Self::TSTypeReference(it) => Address::from_ptr(it),
+            Self::TSTypeName(it) => it.address(),
+            Self::TSQualifiedName(it) => Address::from_ptr(it),
+            Self::TSTypeParameterInstantiation(it) => Address::from_ptr(it),
+            Self::TSTypeParameter(it) => Address::from_ptr(it),
+            Self::TSTypeParameterDeclaration(it) => Address::from_ptr(it),
+            Self::TSTypeAliasDeclaration(it) => Address::from_ptr(it),
+            Self::TSClassImplements(it) => Address::from_ptr(it),
+            Self::TSInterfaceDeclaration(it) => Address::from_ptr(it),
+            Self::TSPropertySignature(it) => Address::from_ptr(it),
+            Self::TSCallSignatureDeclaration(it) => Address::from_ptr(it),
+            Self::TSMethodSignature(it) => Address::from_ptr(it),
+            Self::TSConstructSignatureDeclaration(it) => Address::from_ptr(it),
+            Self::TSIndexSignatureName(it) => Address::from_ptr(it),
+            Self::TSInterfaceHeritage(it) => Address::from_ptr(it),
+            Self::TSTypePredicate(it) => Address::from_ptr(it),
+            Self::TSModuleDeclaration(it) => Address::from_ptr(it),
+            Self::TSModuleBlock(it) => Address::from_ptr(it),
+            Self::TSTypeLiteral(it) => Address::from_ptr(it),
+            Self::TSInferType(it) => Address::from_ptr(it),
+            Self::TSTypeQuery(it) => Address::from_ptr(it),
+            Self::TSImportType(it) => Address::from_ptr(it),
+            Self::TSMappedType(it) => Address::from_ptr(it),
+            Self::TSTemplateLiteralType(it) => Address::from_ptr(it),
+            Self::TSAsExpression(it) => Address::from_ptr(it),
+            Self::TSSatisfiesExpression(it) => Address::from_ptr(it),
+            Self::TSTypeAssertion(it) => Address::from_ptr(it),
+            Self::TSImportEqualsDeclaration(it) => Address::from_ptr(it),
+            Self::TSExternalModuleReference(it) => Address::from_ptr(it),
+            Self::TSNonNullExpression(it) => Address::from_ptr(it),
+            Self::Decorator(it) => Address::from_ptr(it),
+            Self::TSExportAssignment(it) => Address::from_ptr(it),
+            Self::TSNamespaceExportDeclaration(it) => Address::from_ptr(it),
+            Self::TSInstantiationExpression(it) => Address::from_ptr(it),
+            Self::JSDocNullableType(it) => Address::from_ptr(it),
+            Self::JSDocNonNullableType(it) => Address::from_ptr(it),
+            Self::JSDocUnknownType(it) => Address::from_ptr(it),
+        }
+    }
+}
+
 impl<'a> AstKind<'a> {
     #[inline]
     pub fn as_program(self) -> Option<&'a Program<'a>> {

crates/oxc_semantic/src/node.rs

@@ -1,3 +1,4 @@
+use oxc_allocator::{Address, GetAddress};
 use oxc_ast::AstKind;
 use oxc_cfg::BlockNodeId;
 use oxc_index::IndexVec;
@@ -85,6 +86,13 @@ impl GetSpan for AstNode<'_> {
     }
 }
 
+impl GetAddress for AstNode<'_> {
+    #[inline]
+    fn address(&self) -> Address {
+        self.kind.address()
+    }
+}
+
 /// Untyped AST nodes flattened into an vec
 #[derive(Debug, Default)]
 pub struct AstNodes<'a> {

tasks/ast_tools/src/generators/ast_kind.rs

@@ -5,6 +5,7 @@
 //! * `AstKind::ty` method.
 //! * `AstKind::as_*` methods.
 //! * `GetSpan` impl for `AstKind`.
+//! * `GetAddress` impl for `AstKind`.
 //!
 //! Variants of `AstKind` and `AstType` are created for:
 //!
@@ -108,6 +109,7 @@ impl Generator for AstKindGenerator {
         let mut type_variants = quote!();
         let mut kind_variants = quote!();
         let mut span_match_arms = quote!();
+        let mut address_match_arms = quote!();
         let mut as_methods = quote!();
 
         let mut next_index = 0u16;
@@ -131,6 +133,13 @@ impl Generator for AstKindGenerator {
 
             span_match_arms.extend(quote!( Self::#type_ident(it) => it.span(), ));
 
+            let get_address = match type_def {
+                TypeDef::Struct(_) => quote!(Address::from_ptr(it)),
+                TypeDef::Enum(_) => quote!(it.address()),
+                _ => unreachable!(),
+            };
+            address_match_arms.extend(quote!( Self::#type_ident(it) => #get_address, ));
+
             let as_method_name = format_ident!("as_{}", type_def.snake_name());
             as_methods.extend(quote! {
                 ///@@line_break
@@ -154,6 +163,7 @@ impl Generator for AstKindGenerator {
             use std::ptr;
 
             ///@@line_break
+            use oxc_allocator::{Address, GetAddress};
             use oxc_span::{GetSpan, Span};
 
             ///@@line_break
@@ -196,6 +206,17 @@ impl Generator for AstKindGenerator {
                 }
             }
 
+            ///@@line_break
+            impl GetAddress for AstKind<'_> {
+                // TODO: Once only structs have `AstKind`s (https://github.com/oxc-project/oxc/issues/11490),
+                // mark this method `#[inline]`, because then it'll be boiled down to a single instruction.
+                fn address(&self) -> Address {
+                    match *self {
+                        #address_match_arms
+                    }
+                }
+            }
+
             ///@@line_break
             impl<'a> AstKind<'a> {
                 #as_methods