refactor(es/utils): Unify prepend_stmts (#9493)

Author: magic-akari

.changeset/rare-queens-change.md

@@ -0,0 +1,7 @@
+---
+swc_ecma_transforms_typescript: patch
+swc_ecma_utils: patch
+swc_core: patch
+---
+
+refactor(es/utils): Unify `prepend_stmts`

crates/swc_ecma_transforms_typescript/src/transform.rs

@@ -13,7 +13,7 @@ use swc_ecma_ast::*;
 use swc_ecma_utils::{
     alias_ident_for, constructor::inject_after_super, find_pat_ids, is_literal, member_expr,
     private_ident, quote_ident, quote_str, stack_size::maybe_grow_default, ExprFactory, QueryRef,
-    RefRewriter,
+    RefRewriter, StmtLikeInjector,
 };
 use swc_ecma_visit::{as_folder, noop_visit_mut_type, Fold, VisitMut, VisitMutWith};
 
@@ -22,7 +22,7 @@ use crate::{
     ts_enum::{EnumValueComputer, InlineEnum, TsEnumRecord, TsEnumRecordKey, TsEnumRecordValue},
     utils::{
         assign_value_to_this_private_prop, assign_value_to_this_prop, AsCollapsibleDecl,
-        AsEnumOrModule, Factory, VecStmtLike,
+        AsEnumOrModule, Factory,
     },
 };
 
@@ -1162,7 +1162,7 @@ impl Transform {
         }
 
         if should_inject {
-            n.inject_after_directive([
+            n.prepend_stmts([
                 // import { createRequire } from "module";
                 ImportDecl {
                     span: DUMMY_SP,

crates/swc_ecma_transforms_typescript/src/utils.rs

@@ -139,32 +139,3 @@ impl Factory {
         }
     }
 }
-
-pub(crate) trait VecStmtLike {
-    type StmtLike;
-    fn inject_after_directive<T>(&mut self, replace_with: T)
-    where
-        T: IntoIterator<Item = Self::StmtLike>;
-}
-
-impl VecStmtLike for Vec<ModuleItem> {
-    type StmtLike = ModuleItem;
-
-    fn inject_after_directive<T>(&mut self, replace_with: T)
-    where
-        T: IntoIterator<Item = Self::StmtLike>,
-    {
-        let directive_pos = self
-            .iter()
-            .position(|stmt| match stmt {
-                // search first none directive
-                ModuleItem::Stmt(Stmt::Expr(ExprStmt { expr, .. })) => {
-                    !matches!(**expr, Expr::Lit(Lit::Str(..)))
-                }
-                _ => true,
-            })
-            .unwrap_or_default();
-
-        self.splice(directive_pos..directive_pos, replace_with);
-    }
-}

crates/swc_ecma_utils/src/lib.rs

@@ -336,6 +336,45 @@ impl StmtLike for ModuleItem {
     }
 }
 
+/// Prepends statements after directive statements.
+pub trait StmtLikeInjector<S>
+where
+    S: StmtLike,
+{
+    fn prepend_stmt(&mut self, insert_with: S);
+    fn prepend_stmts<I>(&mut self, insert_with: I)
+    where
+        I: IntoIterator<Item = S>;
+}
+
+impl<S> StmtLikeInjector<S> for Vec<S>
+where
+    S: StmtLike,
+{
+    /// Note: If there is no directive, use `insert` instead.
+    fn prepend_stmt(&mut self, insert_with: S) {
+        let directive_pos = self
+            .iter()
+            .position(|stmt| !stmt.as_stmt().map_or(false, is_maybe_branch_directive))
+            .unwrap_or(self.len());
+
+        self.insert(directive_pos, insert_with);
+    }
+
+    /// Note: If there is no directive, use `splice` instead.
+    fn prepend_stmts<I>(&mut self, insert_with: I)
+    where
+        I: IntoIterator<Item = S>,
+    {
+        let directive_pos = self
+            .iter()
+            .position(|stmt| !stmt.as_stmt().map_or(false, is_maybe_branch_directive))
+            .unwrap_or(self.len());
+
+        self.splice(directive_pos..directive_pos, insert_with);
+    }
+}
+
 pub type BoolValue = Value<bool>;
 
 pub trait IsEmpty {
@@ -2218,17 +2257,9 @@ pub fn opt_chain_test(
 }
 
 /// inject `branch` after directives
+#[inline]
 pub fn prepend_stmt<T: StmtLike>(stmts: &mut Vec<T>, stmt: T) {
-    let idx = stmts
-        .iter()
-        .position(|item| {
-            item.as_stmt()
-                .map(|s| !is_maybe_branch_directive(s))
-                .unwrap_or(true)
-        })
-        .unwrap_or(stmts.len());
-
-    stmts.insert(idx, stmt);
+    stmts.prepend_stmt(stmt);
 }
 
 /// If the stmt is maybe a directive like `"use strict";`
@@ -2240,34 +2271,9 @@ pub fn is_maybe_branch_directive(stmt: &Stmt) -> bool {
 }
 
 /// inject `stmts` after directives
+#[inline]
 pub fn prepend_stmts<T: StmtLike>(to: &mut Vec<T>, stmts: impl ExactSizeIterator<Item = T>) {
-    let idx = to
-        .iter()
-        .position(|item| {
-            if let Some(&Stmt::Expr(ExprStmt { ref expr, .. })) = item.as_stmt() {
-                match &**expr {
-                    Expr::Lit(Lit::Str(..)) => return false,
-                    Expr::Call(expr) => match expr.callee {
-                        Callee::Super(_) | Callee::Import(_) => return false,
-                        Callee::Expr(_) => {}
-                    },
-                    _ => {}
-                }
-            }
-
-            true
-        })
-        .unwrap_or(to.len());
-
-    let mut buf = Vec::with_capacity(to.len() + stmts.len());
-    // TODO: Optimize (maybe unsafe)
-
-    buf.extend(to.drain(..idx));
-    buf.extend(stmts);
-    buf.append(to);
-    debug_assert!(to.is_empty());
-
-    *to = buf
+    to.prepend_stmts(stmts);
 }
 
 pub trait IsDirective {