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

Code before:
            direction: Direction::After,
        });
    }

    /// Insert multiple statements after the statement at the given address
    pub fn insert_many_after(&mut self, address: *const Stmt, stmts: Vec<Stmt>) {
        let entry = self.stmts.entry(address).or_default();
        for stmt in stmts {
            entry.push(AdjacentStmt {
                stmt,
                direction: Direction::After,
            });
        }
    }

    /// Get all statements to be inserted at the given address
    fn take_stmts(&mut self, address: *const Stmt) -> Option<Vec<AdjacentStmt>> {
        self.stmts.remove(&address)
    }
}

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() {

Review comment (security):
The comment says "we cannot inject statements after the arrow body", but the conversion is done to enable injecting statements before the return statement (inside the block), not after. Consider updating the comment to clarify that this conversion enables statement injection within the arrow function body.
```suggestion
        // to enable injecting statements within the arrow function body (before the return statement).
```

Code after:
            stmt,
            direction: Direction::After,
        });
    }

    /// Insert multiple statements after the statement at the given address
    pub fn insert_many_after(&mut self, address: *const Stmt, stmts: Vec<Stmt>) {
        let entry = self.stmts.entry(address).or_default();
        for stmt in stmts {
            entry.push(AdjacentStmt {
                stmt,
                direction: Direction::After,
            });
        }
    }

    /// Get all statements to be inserted at the given address
    fn take_stmts(&mut self, address: *const Stmt) -> Option<Vec<AdjacentStmt>> {
        self.stmts.remove(&address)
    }
}

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());
                    }
                }
            }
        }
    }