refactor(ast): get rid of unsafe transmutation in VisitMut trait. (#2764)

This will close #2745,

In this PR I attempt to fix this issue using a combination of ideas
discussed in the issue mentioned above, I've created this early draft so
people can pitch in if there is something I should consider doing.

The first goal of this PR is to resolve the issue with the possible
illegal references, As a result of my approach it would also end up with
a bunch of walk_* and walk_*_mut functions to help with the abstraction.
I want to eliminate enter_node and leave_node functions, but I still
haven't started working on it since I first want to familiarize myself
with all of its usage throughout the project. I'm hesitating to do it at
the moment, When we want to do this it would require quite a bit of
refactoring so we should make sure it is probably going to work and end
up being a better implementation.

crates/oxc_ast/src/ast_kind.rs

@@ -3,9 +3,22 @@ use oxc_span::{Atom, GetSpan, Span};
 #[allow(clippy::wildcard_imports)]
 use crate::ast::*;
 
-/// Untyped AST Node Kind
+macro_rules! ast_kinds {
-#[derive(Debug, Clone, Copy)]
+    { $($ident:ident($type:ty),)* } => (
-pub enum AstKind<'a> {
+        #[derive(Debug, Clone, Copy)]
+        pub enum AstType {
+            $($ident,)*
+        }
+
+        /// Untyped AST Node Kind
+        #[derive(Debug, Clone, Copy)]
+        pub enum AstKind<'a> {
+            $($ident($type),)*
+        }
+    )
+}
+
+ast_kinds! {
     Program(&'a Program<'a>),
     Directive(&'a Directive<'a>),
     Hashbang(&'a Hashbang<'a>),

crates/oxc_ast/src/lib.rs

@@ -20,13 +20,13 @@ mod span;
 pub mod syntax_directed_operations;
 mod trivia;
 mod visit;
-mod visit_mut;
+pub mod visit_mut;
 
 pub use num_bigint::BigUint;
 
 pub use crate::{
     ast_builder::AstBuilder,
-    ast_kind::AstKind,
+    ast_kind::{AstKind, AstType},
     trivia::{Comment, CommentKind, Trivias, TriviasMap},
     visit::Visit,
     visit_mut::VisitMut,

crates/oxc_transformer/src/es2015/arrow_functions.rs

@@ -1,7 +1,7 @@
 use std::rc::Rc;
 
 use oxc_allocator::Vec;
-use oxc_ast::{ast::*, AstBuilder, AstKind, VisitMut};
+use oxc_ast::{ast::*, AstBuilder, AstType, VisitMut};
 use oxc_span::{Atom, SPAN};
 use serde::Deserialize;
 
@@ -16,7 +16,7 @@ use crate::TransformTarget;
 /// * <https://github.com/babel/babel/tree/main/packages/babel-plugin-transform-arrow-functions>
 pub struct ArrowFunctions<'a> {
     ast: Rc<AstBuilder<'a>>,
-    nodes: Vec<'a, AstKind<'a>>,
+    nodes: Vec<'a, AstType>,
     uid: usize,
     has_this: bool,
     /// Insert a variable declaration at the top of the BlockStatement
@@ -33,11 +33,11 @@ pub struct ArrowFunctionsOptions {
 }
 
 impl<'a> VisitMut<'a> for ArrowFunctions<'a> {
-    fn enter_node(&mut self, kind: AstKind<'a>) {
+    fn enter_node(&mut self, kind: AstType) {
         self.nodes.push(kind);
     }
 
-    fn leave_node(&mut self, _kind: AstKind<'a>) {
+    fn leave_node(&mut self, _kind: AstType) {
         self.nodes.pop();
     }
 
@@ -45,8 +45,8 @@ impl<'a> VisitMut<'a> for ArrowFunctions<'a> {
         let parent_kind = self.nodes.last().unwrap();
         let parent_parent_kind = self.nodes[self.nodes.len() - 2];
         if ident.name == "this"
-            && (matches!(parent_kind, AstKind::JSXElementName(_))
+            && (matches!(parent_kind, AstType::JSXElementName)
-                || matches!(parent_parent_kind, AstKind::JSXMemberExpression(_)))
+                || matches!(parent_parent_kind, AstType::JSXMemberExpression))
         {
             if !self.has_this {
                 self.has_this = true;

crates/oxc_transformer/src/es2015/new_target.rs

@@ -1,6 +1,6 @@
 use crate::{context::TransformerCtx, TransformOptions, TransformTarget};
 use oxc_allocator::Vec;
-use oxc_ast::{ast::*, AstBuilder, AstKind, VisitMut};
+use oxc_ast::{ast::*, AstBuilder};
 use oxc_diagnostics::miette;
 use oxc_span::{Atom, Span, SPAN};
 use oxc_syntax::operator::BinaryOperator;
@@ -24,18 +24,54 @@ enum NewTargetKind<'a> {
     Function(Option<Atom<'a>>),
 }
 
-impl<'a> VisitMut<'a> for NewTarget<'a> {
+impl<'a> NewTarget<'a> {
-    fn enter_node(&mut self, kind: AstKind<'a>) {
+    pub(crate) fn enter_method_definition(&mut self, def: &MethodDefinition<'a>) {
-        if let Some(kind) = self.get_kind(kind) {
+        let kind = match def.kind {
-            self.kinds.push(kind);
+            MethodDefinitionKind::Get
+            | MethodDefinitionKind::Set
+            | MethodDefinitionKind::Method => NewTargetKind::Method,
+            MethodDefinitionKind::Constructor => NewTargetKind::Constructor,
+        };
+        self.push(kind);
+    }
+
+    pub(crate) fn leave_method_definition(&mut self, _: &MethodDefinition) {
+        self.pop();
+    }
+
+    pub(crate) fn enter_object_property(&mut self, prop: &ObjectProperty<'a>) {
+        if prop.method {
+            self.push(NewTargetKind::Method);
         }
     }
 
-    fn leave_node(&mut self, kind: AstKind<'a>) {
+    pub(crate) fn leave_object_property(&mut self, prop: &ObjectProperty<'a>) {
-        if self.get_kind(kind).is_some() {
+        if prop.method {
-            self.kinds.pop();
+            self.pop();
         }
     }
+
+    pub(crate) fn enter_function(&mut self, func: &Function<'a>) {
+        if let Some(kind) = self.function_new_target_kind(func) {
+            self.push(kind);
+        }
+    }
+
+    pub(crate) fn leave_function(&mut self, func: &Function<'a>) {
+        if self.function_new_target_kind(func).is_some() {
+            self.pop();
+        }
+    }
+
+    fn function_new_target_kind(&self, func: &Function<'a>) -> Option<NewTargetKind<'a>> {
+        // oxc visitor `MethodDefinitionKind` will enter `Function` node, here need to exclude it
+        if let Some(kind) = self.kinds.last() {
+            if !matches!(kind, NewTargetKind::Function(_)) {
+                return None;
+            }
+        }
+        func.id.as_ref().map(|id| NewTargetKind::Function(Some(id.name.clone())))
+    }
 }
 
 impl<'a> NewTarget<'a> {
@@ -52,26 +88,12 @@ impl<'a> NewTarget<'a> {
         })
     }
 
-    fn get_kind(&self, kind: AstKind<'a>) -> Option<NewTargetKind<'a>> {
+    fn push(&mut self, kind: NewTargetKind<'a>) {
-        match kind {
+        self.kinds.push(kind);
-            AstKind::MethodDefinition(def) => match def.kind {
+    }
-                MethodDefinitionKind::Get
+
-                | MethodDefinitionKind::Set
+    fn pop(&mut self) {
-                | MethodDefinitionKind::Method => Some(NewTargetKind::Method),
+        self.kinds.pop();
-                MethodDefinitionKind::Constructor => Some(NewTargetKind::Constructor),
-            },
-            AstKind::ObjectProperty(property) => property.method.then_some(NewTargetKind::Method),
-            AstKind::Function(function) => {
-                // oxc visitor `MethodDefinitionKind` will enter `Function` node, here need to exclude it
-                if let Some(kind) = self.kinds.last() {
-                    if !matches!(kind, NewTargetKind::Function(_)) {
-                        return None;
-                    }
-                }
-                function.id.as_ref().map(|id| NewTargetKind::Function(Some(id.name.clone())))
-            }
-            _ => None,
-        }
     }
 
     fn create_constructor_expr(&self, span: Span) -> Expression<'a> {

crates/oxc_transformer/src/lib.rs

@@ -28,7 +28,7 @@ use std::{cell::RefCell, rc::Rc, sync::Arc};
 
 use es2015::TemplateLiterals;
 use oxc_allocator::Allocator;
-use oxc_ast::{ast::*, AstBuilder, AstKind, VisitMut};
+use oxc_ast::{ast::*, visit_mut::walk_function_mut, AstBuilder, AstType, VisitMut};
 use oxc_diagnostics::Error;
 use oxc_semantic::{ScopeFlags, Semantic};
 use oxc_span::SourceType;
@@ -152,16 +152,8 @@ impl<'a> Transformer<'a> {
 }
 
 impl<'a> VisitMut<'a> for Transformer<'a> {
-    fn enter_node(&mut self, kind: oxc_ast::AstKind<'a>) {
-        self.es2015_new_target.as_mut().map(|t| t.enter_node(kind));
-    }
-
-    fn leave_node(&mut self, kind: oxc_ast::AstKind<'a>) {
-        self.es2015_new_target.as_mut().map(|t| t.leave_node(kind));
-    }
-
     fn visit_program(&mut self, program: &mut Program<'a>) {
-        let kind = AstKind::Program(self.alloc(program));
+        let kind = AstType::Program;
         self.enter_scope({
             let mut flags = ScopeFlags::Top;
             if program.is_strict() {
@@ -185,7 +177,7 @@ impl<'a> VisitMut<'a> for Transformer<'a> {
     }
 
     fn visit_assignment_expression(&mut self, expr: &mut AssignmentExpression<'a>) {
-        let kind = AstKind::AssignmentExpression(self.alloc(expr));
+        let kind = AstType::AssignmentExpression;
         self.enter_node(kind);
 
         self.es2015_function_name.as_mut().map(|t| t.transform_assignment_expression(expr));
@@ -238,7 +230,7 @@ impl<'a> VisitMut<'a> for Transformer<'a> {
     }
 
     fn visit_catch_clause(&mut self, clause: &mut CatchClause<'a>) {
-        let kind = AstKind::CatchClause(self.alloc(clause));
+        let kind = AstType::CatchClause;
         self.enter_scope(ScopeFlags::empty());
         self.enter_node(kind);
 
@@ -253,7 +245,7 @@ impl<'a> VisitMut<'a> for Transformer<'a> {
     }
 
     fn visit_object_expression(&mut self, expr: &mut ObjectExpression<'a>) {
-        let kind = AstKind::ObjectExpression(self.alloc(expr));
+        let kind = AstType::ObjectExpression;
         self.enter_node(kind);
         self.es2015_function_name.as_mut().map(|t| t.transform_object_expression(expr));
         self.es2015_duplicate_keys.as_mut().map(|t| t.transform_object_expression(expr));
@@ -265,8 +257,9 @@ impl<'a> VisitMut<'a> for Transformer<'a> {
     }
 
     fn visit_object_property(&mut self, prop: &mut ObjectProperty<'a>) {
-        let kind = AstKind::ObjectProperty(self.alloc(prop));
+        let kind = AstType::ObjectProperty;
         self.enter_node(kind);
+        self.es2015_new_target.as_mut().map(|t| t.enter_object_property(prop));
 
         self.es2015_shorthand_properties.as_mut().map(|t| t.transform_object_property(prop));
         self.es3_property_literal.as_mut().map(|t| t.transform_object_property(prop));
@@ -276,6 +269,8 @@ impl<'a> VisitMut<'a> for Transformer<'a> {
         if let Some(init) = &mut prop.init {
             self.visit_expression(init);
         }
+
+        self.es2015_new_target.as_mut().map(|t| t.leave_object_property(prop));
         self.leave_node(kind);
     }
 
@@ -288,7 +283,7 @@ impl<'a> VisitMut<'a> for Transformer<'a> {
     }
 
     fn visit_variable_declarator(&mut self, declarator: &mut VariableDeclarator<'a>) {
-        let kind = AstKind::VariableDeclarator(self.alloc(declarator));
+        let kind = AstType::VariableDeclarator;
         self.enter_node(kind);
 
         self.es2015_function_name.as_mut().map(|t| t.transform_variable_declarator(declarator));
@@ -315,8 +310,9 @@ impl<'a> VisitMut<'a> for Transformer<'a> {
     }
 
     fn visit_method_definition(&mut self, def: &mut MethodDefinition<'a>) {
-        let kind = AstKind::MethodDefinition(self.alloc(def));
+        let kind = AstType::MethodDefinition;
         self.enter_node(kind);
+        self.es2015_new_target.as_mut().map(|t| t.enter_method_definition(def));
 
         self.typescript.as_mut().map(|t| t.transform_method_definition(def));
 
@@ -332,6 +328,13 @@ impl<'a> VisitMut<'a> for Transformer<'a> {
         };
         self.visit_property_key(&mut def.key);
         self.visit_function(&mut def.value, Some(flags));
+        self.es2015_new_target.as_mut().map(|t| t.leave_method_definition(def));
         self.leave_node(kind);
     }
+
+    fn visit_function(&mut self, func: &mut Function<'a>, flags: Option<ScopeFlags>) {
+        self.es2015_new_target.as_mut().map(|t| t.enter_function(func));
+        walk_function_mut(self, func, flags);
+        self.es2015_new_target.as_mut().map(|t| t.leave_function(func));
+    }
 }