feat(allocator): TakeIn trait with AllocatorAccessor (#11201)

closes #10314

Author: Boshen

crates/oxc_allocator/src/accessor.rs

@@ -0,0 +1,14 @@
+use crate::Allocator;
+
+/// Accessor for getting the underlying allocator.
+pub trait AllocatorAccessor<'a> {
+    /// Get the underlying allocator.
+    fn allocator(self) -> &'a Allocator;
+}
+
+impl<'a> AllocatorAccessor<'a> for &'a Allocator {
+    #[inline]
+    fn allocator(self) -> &'a Allocator {
+        self
+    }
+}

crates/oxc_allocator/src/lib.rs

@@ -21,6 +21,7 @@
 
 #![warn(missing_docs)]
 
+mod accessor;
 mod address;
 mod alloc;
 mod allocator;
@@ -36,6 +37,7 @@ mod take_in;
 mod vec;
 mod vec2;
 
+pub use accessor::AllocatorAccessor;
 pub use address::{Address, GetAddress};
 pub use allocator::Allocator;
 pub use boxed::Box;

crates/oxc_allocator/src/take_in.rs

@@ -1,12 +1,13 @@
 use std::{cell::Cell, mem, num};
 
-use crate::{Allocator, Box, Vec};
+use crate::{Allocator, AllocatorAccessor, Box, Vec};
 
 /// A trait to replace an existing AST node with a dummy.
 pub trait TakeIn<'a>: Dummy<'a> {
     /// Replace node with a dummy.
     #[must_use]
-    fn take_in(&mut self, allocator: &'a Allocator) -> Self {
+    fn take_in<A: AllocatorAccessor<'a>>(&mut self, allocator_accessor: A) -> Self {
+        let allocator = allocator_accessor.allocator();
         let dummy = Dummy::dummy(allocator);
         mem::replace(self, dummy)
     }

crates/oxc_ast/src/ast_builder_impl.rs

@@ -2,7 +2,7 @@
 
 use std::borrow::Cow;
 
-use oxc_allocator::{Allocator, Box, FromIn, IntoIn, String, Vec};
+use oxc_allocator::{Allocator, AllocatorAccessor, Box, FromIn, IntoIn, String, Vec};
 use oxc_span::{Atom, SPAN, Span};
 use oxc_syntax::{number::NumberBase, operator::UnaryOperator, scope::ScopeId};
 
@@ -19,6 +19,13 @@ impl<'a, T> FromIn<'a, NONE> for Option<Box<'a, T>> {
     }
 }
 
+impl<'a> AllocatorAccessor<'a> for AstBuilder<'a> {
+    #[inline]
+    fn allocator(self) -> &'a Allocator {
+        self.allocator
+    }
+}
+
 impl<'a> AstBuilder<'a> {
     /// Create a new AST builder that will allocate nodes in the given allocator.
     #[inline]

crates/oxc_isolated_declarations/src/module.rs

@@ -187,7 +187,7 @@ impl<'a> IsolatedDeclarations<'a> {
         stmts.iter_mut().for_each(|stmt| {
             if let Statement::ExportNamedDeclaration(decl) = stmt {
                 if let Some(declaration) = &mut decl.declaration {
-                    *stmt = Statement::from(declaration.take_in(self.ast.allocator));
+                    *stmt = Statement::from(declaration.take_in(self.ast));
                 }
             }
         });

crates/oxc_minifier/src/peephole/convert_to_dotted_properties.rs

@@ -22,7 +22,7 @@ impl<'a> LatePeepholeOptimizations {
             *expr =
                 MemberExpression::StaticMemberExpression(ctx.ast.alloc_static_member_expression(
                     e.span,
-                    e.object.take_in(ctx.ast.allocator),
+                    e.object.take_in(ctx.ast),
                     property,
                     e.optional,
                 ));

crates/oxc_minifier/src/peephole/fold_constants.rs

@@ -132,7 +132,7 @@ impl<'a> PeepholeOptimizations {
             // (TRUE || x) => TRUE (also, (3 || x) => 3)
             // (FALSE && x) => FALSE
             if if lval { op.is_or() } else { op.is_and() } {
-                return Some(logical_expr.left.take_in(ctx.ast.allocator));
+                return Some(logical_expr.left.take_in(ctx.ast));
             } else if !left.may_have_side_effects(&ctx) {
                 let should_keep_indirect_access =
                     Self::should_keep_indirect_access(&logical_expr.right, ctx);
@@ -147,19 +147,19 @@ impl<'a> PeepholeOptimizations {
                                 None,
                                 NumberBase::Decimal,
                             ),
-                            logical_expr.right.take_in(ctx.ast.allocator),
+                            logical_expr.right.take_in(ctx.ast),
                         ]),
                     ));
                 }
                 // (FALSE || x) => x
                 // (TRUE && x) => x
-                return Some(logical_expr.right.take_in(ctx.ast.allocator));
+                return Some(logical_expr.right.take_in(ctx.ast));
             }
             // Left side may have side effects, but we know its boolean value.
             // e.g. true_with_sideeffects || foo() => true_with_sideeffects, foo()
             // or: false_with_sideeffects && foo() => false_with_sideeffects, foo()
-            let left = logical_expr.left.take_in(ctx.ast.allocator);
-            let right = logical_expr.right.take_in(ctx.ast.allocator);
+            let left = logical_expr.left.take_in(ctx.ast);
+            let right = logical_expr.right.take_in(ctx.ast);
             let vec = ctx.ast.vec_from_array([left, right]);
             let sequence_expr = ctx.ast.expression_sequence(logical_expr.span, vec);
             return Some(sequence_expr);
@@ -174,8 +174,8 @@ impl<'a> PeepholeOptimizations {
                         if !right_boolean && left_child_op.is_or()
                             || right_boolean && left_child_op.is_and()
                         {
-                            let left = left_child.left.take_in(ctx.ast.allocator);
-                            let right = logical_expr.right.take_in(ctx.ast.allocator);
+                            let left = left_child.left.take_in(ctx.ast);
+                            let right = logical_expr.right.take_in(ctx.ast);
                             let logic_expr = ctx.ast.expression_logical(
                                 logical_expr.span,
                                 left,
@@ -204,8 +204,8 @@ impl<'a> PeepholeOptimizations {
                 Some(if left.may_have_side_effects(&ctx) {
                     // e.g. `(a(), null) ?? 1` => `(a(), null, 1)`
                     let expressions = ctx.ast.vec_from_array([
-                        logical_expr.left.take_in(ctx.ast.allocator),
-                        logical_expr.right.take_in(ctx.ast.allocator),
+                        logical_expr.left.take_in(ctx.ast),
+                        logical_expr.right.take_in(ctx.ast),
                     ]);
                     ctx.ast.expression_sequence(logical_expr.span, expressions)
                 } else {
@@ -222,12 +222,12 @@ impl<'a> PeepholeOptimizations {
                                     None,
                                     NumberBase::Decimal,
                                 ),
-                                logical_expr.right.take_in(ctx.ast.allocator),
+                                logical_expr.right.take_in(ctx.ast),
                             ]),
                         ));
                     }
                     // nullish condition => this expression evaluates to the right side.
-                    logical_expr.right.take_in(ctx.ast.allocator)
+                    logical_expr.right.take_in(ctx.ast)
                 })
             }
             ValueType::Number
@@ -248,12 +248,12 @@ impl<'a> PeepholeOptimizations {
                                 None,
                                 NumberBase::Decimal,
                             ),
-                            logical_expr.left.take_in(ctx.ast.allocator),
+                            logical_expr.left.take_in(ctx.ast),
                         ]),
                     ));
                 }
                 // non-nullish condition => this expression evaluates to the left side.
-                Some(logical_expr.left.take_in(ctx.ast.allocator))
+                Some(logical_expr.left.take_in(ctx.ast))
             }
             ValueType::Undetermined => None,
         }
@@ -391,14 +391,14 @@ impl<'a> PeepholeOptimizations {
                     let span = Span::new(left_binary_expr.right.span().start, e.right.span().end);
                     let value = ctx.ast.atom_from_strs_array([&left_str, &right_str]);
                     let right = ctx.ast.expression_string_literal(span, value, None);
-                    let left = left_binary_expr.left.take_in(ctx.ast.allocator);
+                    let left = left_binary_expr.left.take_in(ctx.ast);
                     return Some(ctx.ast.expression_binary(e.span, left, e.operator, right));
                 }
 
                 if let Some(new_right) =
                     Self::try_fold_add_op(&mut left_binary_expr.right, &mut e.right, ctx)
                 {
-                    let left = left_binary_expr.left.take_in(ctx.ast.allocator);
+                    let left = left_binary_expr.left.take_in(ctx.ast);
                     return Some(ctx.ast.expression_binary(e.span, left, e.operator, new_right));
                 }
             }
@@ -437,7 +437,7 @@ impl<'a> PeepholeOptimizations {
                 }
                 left.quasis.extend(right.quasis.drain(1..)); // first quasi is already handled
                 left.expressions.extend(right.expressions.drain(..));
-                return Some(left_expr.take_in(ctx.ast.allocator));
+                return Some(left_expr.take_in(ctx.ast));
             }
 
             // "`${x}y` + 'z'" => "`${x}yz`"
@@ -453,7 +453,7 @@ impl<'a> PeepholeOptimizations {
                     .cooked
                     .map(|cooked| ctx.ast.atom(&(cooked.as_str().to_string() + &right_str)));
                 last_quasi.value.cooked = new_cooked;
-                return Some(left_expr.take_in(ctx.ast.allocator));
+                return Some(left_expr.take_in(ctx.ast));
             }
         } else if let Expression::TemplateLiteral(right) = right_expr {
             // "'x' + `y${z}`" => "`xy${z}`"
@@ -471,17 +471,17 @@ impl<'a> PeepholeOptimizations {
                     .cooked
                     .map(|cooked| ctx.ast.atom(&(left_str.into_owned() + cooked.as_str())));
                 first_quasi.value.cooked = new_cooked;
-                return Some(right_expr.take_in(ctx.ast.allocator));
+                return Some(right_expr.take_in(ctx.ast));
             }
         }
 
         // remove useless `+ ""` (e.g. `typeof foo + ""` -> `typeof foo`)
         if Self::evaluates_to_empty_string(left_expr) && right_expr.value_type(&ctx).is_string() {
-            return Some(right_expr.take_in(ctx.ast.allocator));
+            return Some(right_expr.take_in(ctx.ast));
         } else if Self::evaluates_to_empty_string(right_expr)
             && left_expr.value_type(&ctx).is_string()
         {
-            return Some(left_expr.take_in(ctx.ast.allocator));
+            return Some(left_expr.take_in(ctx.ast));
         }
 
         None
@@ -523,7 +523,7 @@ impl<'a> PeepholeOptimizations {
 
         Some(ctx.ast.expression_binary(
             e.span,
-            expr_to_move.take_in(ctx.ast.allocator),
+            expr_to_move.take_in(ctx.ast),
             op,
             ctx.value_to_expr(Span::new(left.right.span().start, e.right.span().end), v),
         ))