Repository: swc-project/swc
File: crates/swc_ecma_transformer/src/common/statement_injector.rs (Rust)

Code before:
impl VisitMutHook<TraverseCtx> for StmtInjector {
    fn exit_module_items(&mut self, node: &mut Vec<ModuleItem>, ctx: &mut TraverseCtx) {
        // First pass: collect all (index, adjacent_stmts) pairs while addresses are
        // valid
        let mut insertions = Vec::new();
        for (i, item) in node.iter().enumerate() {
            // Only process ModuleItem::Stmt variants
            if let ModuleItem::Stmt(stmt) = item {
                let address = stmt as *const Stmt;
                if let Some(adjacent_stmts) = ctx.statement_injector.take_stmts(address) {
                    insertions.push((i, adjacent_stmts));
                }
            }
        }

        // Second pass: process in reverse order to avoid index invalidation
        for (i, adjacent_stmts) in insertions.into_iter().rev() {
            let mut before_stmts = Vec::new();
            let mut after_stmts = Vec::new();

            // Separate statements by direction
            for adjacent in adjacent_stmts {
                match adjacent.direction {
                    Direction::Before => before_stmts.push(adjacent.stmt),
                    Direction::After => after_stmts.push(adjacent.stmt),
                }
            }

            // Insert statements after (insert first since we're going backwards)
            if !after_stmts.is_empty() {
                // Insert all after statements at position i + 1
                for (offset, stmt) in after_stmts.into_iter().enumerate() {
                    node.insert(i + 1 + offset, ModuleItem::Stmt(stmt));
                }
            }

            // Insert statements before
            if !before_stmts.is_empty() {
                // Insert all before statements at position i
                for (offset, stmt) in before_stmts.into_iter().enumerate() {
                    node.insert(i + offset, ModuleItem::Stmt(stmt));
                }
            }
        }
    }

    fn exit_stmts(&mut self, stmts: &mut Vec<Stmt>, ctx: &mut TraverseCtx) {
        // First pass: collect all (index, adjacent_stmts) pairs while addresses are
        // valid
        let mut insertions = Vec::new();
        for (i, stmt) in stmts.iter().enumerate() {

Review comment (security):
The `unwrap()` call assumes that `return_stmt.arg` is always `Some`, but this may not be true if statements were injected into the arrow function. If the arrow body was modified between `enter_arrow_expr` and `exit_arrow_expr` (e.g., if statement injection occurred), the return statement might not contain the original expression, potentially causing a panic.
```suggestion
                        if let Some(arg) = return_stmt.arg.take() {
                            *node.body = BlockStmtOrExpr::Expr(arg);
                        }
```

Code after:

impl VisitMutHook<TraverseCtx> for StmtInjector {
    fn enter_arrow_expr(&mut self, node: &mut ArrowExpr, _: &mut TraverseCtx) {
        // We convert the arrow body to a block statement with a return statement
        // because we cannot inject statements after the arrow body.
        if let BlockStmtOrExpr::Expr(expr) = &mut *node.body {
            *node.body = BlockStmtOrExpr::BlockStmt(BlockStmt {
                span: DUMMY_SP,
                stmts: vec![Stmt::Return(ReturnStmt {
                    span: DUMMY_SP,
                    arg: Some(expr.take()),
                })],
                ..Default::default()
            });
        }
    }

    fn exit_arrow_expr(&mut self, node: &mut ArrowExpr, _: &mut TraverseCtx) {
        if let BlockStmtOrExpr::BlockStmt(block_stmt) = &mut *node.body {
            // If the block statement is something we inserted, we need to
            // convert it back to an expression.
            if block_stmt.span == DUMMY_SP && block_stmt.stmts.len() == 1 {
                if let Stmt::Return(return_stmt) = &mut block_stmt.stmts[0] {
                    if return_stmt.span == DUMMY_SP {
                        *node.body = BlockStmtOrExpr::Expr(return_stmt.arg.take().unwrap());
                    }
                }
            }
        }
    }

    fn exit_module_items(&mut self, node: &mut Vec<ModuleItem>, ctx: &mut TraverseCtx) {
        // First pass: collect all (index, adjacent_stmts) pairs while addresses are
        // valid
        let mut insertions = Vec::new();
        for (i, item) in node.iter().enumerate() {
            // Only process ModuleItem::Stmt variants
            if let ModuleItem::Stmt(stmt) = item {
                let address = stmt as *const Stmt;
                if let Some(adjacent_stmts) = ctx.statement_injector.take_stmts(address) {
                    insertions.push((i, adjacent_stmts));
                }
            }
        }

        // Second pass: process in reverse order to avoid index invalidation
        for (i, adjacent_stmts) in insertions.into_iter().rev() {
            let mut before_stmts = Vec::new();
            let mut after_stmts = Vec::new();

            // Separate statements by direction