perf(transformer): speed up searching for position to insert new statements

crates/oxc_transformer/src/common/top_level_statements.rs

@@ -41,16 +41,58 @@ impl<'a, 'ctx> Traverse<'a> for TopLevelStatements<'a, 'ctx> {
         }
 
         // Insert statements after any existing `import` statements
-        let index = program
-            .body
-            .iter()
-            .rposition(|stmt| matches!(stmt, Statement::ImportDeclaration(_)))
-            .map_or(0, |i| i + 1);
+        let insert_index = if self.ctx.source_type.is_module() {
+            find_insertion_index(&program.body)
+        } else {
+            // Scripts can't have `import` statements, so no need to search
+            0
+        };
 
-        program.body.splice(index..index, stmts.drain(..));
+        program.body.splice(insert_index..insert_index, stmts.drain(..));
     }
 }
 
+/// Find index to insert statements at.
+///
+/// We want to insert after any `import` statements.
+///
+/// We could search from the end of the file backwards until we hit an `import` statement, but in
+/// a large file, that's a lot of statements to search through. So instead, search from the *start*
+/// for first statement which is *not* an `import`.
+/// Usually the correct insertion point is before that statement.
+///
+/// But there is one annoying Babel test that has a non-`import` statement followed by `import`s,
+/// and it expects new statements to be inserted after the last of those `import`s.
+/// `babel-plugin-transform-react-jsx/test/fixtures/autoImport/after-polyfills-2`
+/// To pass that test, we search again if the first statement is not an `import`.
+///
+/// TODO(improve-on-babel): Insertion position is not important. We only do this to pass Babel's tests.
+/// Remove this once we don't have to match Babel's output exactly, and just insert at the start.
+fn find_insertion_index(stmts: &[Statement]) -> usize {
+    let Some(first_stmt) = stmts.first() else {
+        // No statements. Insert at start.
+        return 0;
+    };
+
+    let search_start_index = if matches!(first_stmt, Statement::ImportDeclaration(_)) {
+        // First statement is `import`. Search for more `import`s after this.
+        1
+    } else if !matches!(stmts.get(1), Some(Statement::ImportDeclaration(_))) {
+        // Either there's only 1 statement (a non-`import`), or first 2 statements are both not `import`.
+        // Insert at the start.
+        return 0;
+    } else {
+        // Non-`import`, followed by `import`. Search for more `import`s after this.
+        2
+    };
+
+    // Find first non-`import` after this
+    return stmts[search_start_index..]
+        .iter()
+        .position(|stmt| !matches!(stmt, Statement::ImportDeclaration(_)))
+        .map_or_else(|| stmts.len(), |index| search_start_index + index);
+}
+
 /// Store for statements to be added at top of program
 pub struct TopLevelStatementsStore<'a> {
     stmts: RefCell<Vec<Statement<'a>>>,

tasks/coverage/snapshots/semantic_typescript.snap

@@ -21130,7 +21130,7 @@ after transform: ScopeId(0): [ScopeId(1), ScopeId(2)]
 rebuilt        : ScopeId(0): [ScopeId(1)]
 Symbol reference IDs mismatch:
 after transform: SymbolId(0): [ReferenceId(0)]
-rebuilt        : SymbolId(0): []
+rebuilt        : SymbolId(2): []
 
 tasks/coverage/typescript/tests/cases/compiler/jsxFactoryAndJsxFragmentFactory.tsx
 semantic error: Bindings mismatch:
@@ -21370,7 +21370,7 @@ rebuilt        : []
 tasks/coverage/typescript/tests/cases/compiler/jsxPartialSpread.tsx
 semantic error: Symbol reference IDs mismatch:
 after transform: SymbolId(0): [ReferenceId(2), ReferenceId(3)]
-rebuilt        : SymbolId(0): [ReferenceId(3)]
+rebuilt        : SymbolId(2): [ReferenceId(3)]
 Unresolved references mismatch:
 after transform: ["Parameters", "Partial"]
 rebuilt        : []