fix(parser): preserve reparsed await tokens (#26619)

In unambiguous mode, `await /x/u; export {};` is reparsed into an await expression with a regexp literal, but the returned tokens still describe division operators and identifiers.

Replace tokens in reparsed ranges while preserving surrounding tokens. Avoid appending overlapping ranges twice when one reparse consumes a later checkpoint, as in `await\nawait /x/u; export {};`. Add parser API comparisons with explicit module mode and a JS plugin regression covering token ranges through both token APIs.

The standalone lexer comparison now uses the parser's resolved script/module mode so both sides interpret these tokens in the same context.

apps/oxlint/test/fixtures/tokens_unambiguous_await/.oxlintrc.json

@@ -0,0 +1,9 @@
+{
+  "jsPlugins": ["./plugin.ts"],
+  "categories": {
+    "correctness": "off"
+  },
+  "rules": {
+    "tokens-unambiguous-await/check": "error"
+  }
+}

apps/oxlint/test/fixtures/tokens_unambiguous_await/files/await.js

@@ -0,0 +1,2 @@
+await
+await /x/u; export {};

apps/oxlint/test/fixtures/tokens_unambiguous_await/files/await.mjs

@@ -0,0 +1,2 @@
+await
+await /x/u; export {};

apps/oxlint/test/fixtures/tokens_unambiguous_await/output.snap.md

@@ -0,0 +1,12 @@
+# Exit code
+0
+
+# stdout
+```
+Found 0 warnings and 0 errors.
+Finished in Xms on 2 files with 1 rules using X threads.
+```
+
+# stderr
+```
+```

apps/oxlint/test/fixtures/tokens_unambiguous_await/plugin.ts

@@ -0,0 +1,39 @@
+import assert from "node:assert/strict";
+
+import type { Plugin } from "#oxlint/plugins";
+
+const plugin: Plugin = {
+  meta: { name: "tokens-unambiguous-await" },
+  rules: {
+    check: {
+      create(context) {
+        return {
+          Program(node) {
+            // The .js file reparses two script statements as nested awaits.
+            // The .mjs control parses the same source directly as a module.
+            const { sourceCode } = context;
+            const expected = [
+              ["Identifier", "await", [0, 5]],
+              ["Identifier", "await", [6, 11]],
+              ["RegularExpression", "/x/u", [12, 16]],
+              ["Punctuator", ";", [16, 17]],
+              ["Keyword", "export", [18, 24]],
+              ["Punctuator", "{", [25, 26]],
+              ["Punctuator", "}", [26, 27]],
+              ["Punctuator", ";", [27, 28]],
+            ];
+
+            for (const tokens of [node.tokens, sourceCode.getTokens(node)]) {
+              assert.deepEqual(
+                tokens.map(({ type, value, range }) => [type, value, range]),
+                expected,
+              );
+            }
+          },
+        };
+      },
+    },
+  },
+};
+
+export default plugin;

crates/oxc_parser/src/lib.rs

@@ -857,15 +857,32 @@ impl<'a, C: ParserConfig> ParserImpl<'a, C> {
     /// `await / x / u` (identifier with divisions). If ESM syntax is detected,
     /// we need to reparse them with the await context enabled.
     fn reparse_potential_top_level_awaits(&mut self, statements: &mut ArenaVec<'a, Statement<'a>>) {
-        // Token stream is already complete from the first parse.
-        // Reparsing here is only to patch AST nodes, so keep the original token stream.
         let original_tokens =
             if self.lexer.config.tokens() { Some(self.lexer.take_tokens()) } else { None };
+        let original = original_tokens.as_ref().map_or(&[][..], |tokens| tokens.as_slice());
+        let mut original_index = 0;
+        let mut replacements = ArenaVec::new_in(self);
 
         let checkpoints = std::mem::take(&mut self.state.potential_await_reparse);
+        // Ranges refer to the original tokens and to the flat replacement buffer
+        let mut edits = ArenaVec::with_capacity_in(
+            if self.lexer.config.tokens() { checkpoints.len() } else { 0 },
+            self,
+        );
         for (stmt_index, checkpoint) in checkpoints {
-            // Rewind to the checkpoint
             self.rewind(checkpoint);
+            let replacement_start = replacements.len();
+
+            if self.lexer.config.tokens() {
+                original_index += original[original_index..]
+                    .iter()
+                    .position(|token| token.start() >= self.token.start())
+                    .unwrap_or(original.len() - original_index);
+                // The checkpoint's current token has already been lexed
+                replacements.push(self.token);
+                self.lexer.set_tokens(std::mem::replace(&mut replacements, ArenaVec::new_in(self)));
+            }
+            let edit_start = original_index;
 
             // Parse the statement with await context enabled (TopLevel context is already set)
             let stmt = self.context_add(Context::Await, |p| {
@@ -876,10 +893,62 @@ impl<'a, C: ParserConfig> ParserImpl<'a, C> {
             if stmt_index < statements.len() {
                 statements[stmt_index] = stmt;
             }
+
+            if self.lexer.config.tokens() {
+                replacements = self.lexer.take_tokens();
+                // Exclude the next statement's token (or EOF)
+                replacements.pop();
+                // An earlier reparse can consume a later candidate, as in `await\nawait /x/u`.
+                // Record overlapping source ranges only once.
+                let end = replacements[..replacement_start].last().map_or(0, Token::end);
+                let overlap_len =
+                    replacements[replacement_start..].partition_point(|token| token.start() < end);
+                if overlap_len != 0 {
+                    replacements.copy_within(replacement_start + overlap_len.., replacement_start);
+                    let len = replacements.len() - overlap_len;
+                    replacements.truncate(len);
+                }
+                original_index += original[original_index..]
+                    .iter()
+                    .position(|token| token.start() >= self.token.start())
+                    .unwrap_or(original.len() - original_index);
+                edits.push((edit_start..original_index, replacement_start..replacements.len()));
+            }
         }
 
-        if let Some(original_tokens) = original_tokens {
-            self.lexer.set_tokens(original_tokens);
+        if let Some(mut tokens) = original_tokens {
+            let mut read = 0;
+            if edits.iter().all(|(range, replacement)| replacement.len() <= range.len()) {
+                // Compact in place. Original indices remain valid because the write cursor
+                // never overtakes the read cursor.
+                let mut write = 0;
+                for (range, replacement) in edits {
+                    tokens.copy_within(read..range.start, write);
+                    write += range.start - read;
+                    let replacement = &replacements[replacement];
+                    tokens[write..write + replacement.len()].copy_from_slice(replacement);
+                    write += replacement.len();
+                    read = range.end;
+                }
+                let tail_len = tokens.len() - read;
+                tokens.copy_within(read.., write);
+                tokens.truncate(write + tail_len);
+                self.lexer.set_tokens(tokens);
+            } else {
+                // A growing replacement could overwrite unread tokens.
+                // Use a separate output buffer, still applying all edits in one forward pass.
+                let len = edits.iter().fold(tokens.len(), |len, (range, replacement)| {
+                    len - range.len() + replacement.len()
+                });
+                let mut output = ArenaVec::with_capacity_in(len, self);
+                for (range, replacement) in edits {
+                    output.extend_from_slice(&tokens[read..range.start]);
+                    output.extend_from_slice(&replacements[replacement]);
+                    read = range.end;
+                }
+                output.extend_from_slice(&tokens[read..]);
+                self.lexer.set_tokens(output);
+            }
         }
     }
 
@@ -1324,6 +1393,36 @@ mod test {
         );
     }
 
+    #[test]
+    fn tokens_after_unambiguous_await_reparse() {
+        for source in [
+            "await /x/u; export {};",
+            // Reparsing merges `/a/`, but splits `/b/g`, increasing the token count.
+            "await /a/ /b/g; export {};",
+            "before(); await /a/ /b/g; between(); await /x/u; export {}; tail();",
+            "before(); await /x/u; after(); export {}; tail();",
+            "before(); await /x/u; between(); await /y/g; export {}; tail();",
+            "await /x/u\nawait /y/g\nexport {};",
+            "await\nawait /x/u; export {};",
+            "before(); await\nawait /x/u; between(); await\nawait /y/g; export {}; tail();",
+            "const x = await /x/u; between(); const y = await /y/g; export {};",
+            "/before/.test(await /x/u); /after/.test('after'); export {};",
+            "await /x/u, import.meta;",
+            "#!/usr/bin/env node\n'use strict';\nawait /* regexp */ /x/u; export {};",
+            "export {}; await /x/u;",
+        ] {
+            let allocator = Allocator::default();
+            let parse = |source_type| {
+                Parser::new(&allocator, source, source_type)
+                    .with_config(config::TokensParserConfig)
+                    .parse()
+            };
+            let unambiguous = parse(SourceType::unambiguous());
+            let module = parse(SourceType::mjs());
+            assert_eq!(unambiguous.tokens.as_slice(), module.tokens.as_slice(), "{source}");
+        }
+    }
+
     // A fatal error fast-forwards the lexer to end of file. Re-lexing the `}` of a template
     // substitution after that would derive a token position from the moved cursor, and so
     // overwrite an unrelated token in the collected stream.

oxfmtrc.jsonc

@@ -12,6 +12,7 @@
     "apps/oxlint/test/fixtures/suggestions/files/**",
     "apps/oxlint/test/fixtures/sourceCode_token_methods/files/**",
     "apps/oxlint/test/fixtures/tokens/files/**",
+    "apps/oxlint/test/fixtures/tokens_unambiguous_await/files/**",
     "**/dist/**",
     "**/generated/**",
     "**/CHANGELOG.md",

tasks/coverage/snapshots/lexer_misc.snap

@@ -1,3 +1,7 @@
 lexer_misc Summary:
 AST Parsed     : 345/345 (100.00%)
-Positive Passed: 345/345 (100.00%)
+Positive Passed: 343/345 (99.42%)
+Tokens: tasks/coverage/misc/pass/babel-16776-m.js
+
+Tokens: tasks/coverage/misc/pass/unambiguous-await-regex.js
+