feat(linter/plugins): implement SourceCode#getTokensAfter() (#15971)

- Part of
#14829 (comment).
- Follow up to #15861.

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: lilnasy

apps/oxlint/src-js/plugins/tokens.ts

@@ -560,14 +560,89 @@ export function getTokenOrCommentAfter(nodeOrToken: NodeOrToken | Comment, skip?
  * @param countOptions? - Options object. Same options as `getFirstTokens()`.
  * @returns Array of `Token`s.
  */
-/* oxlint-disable no-unused-vars */
 export function getTokensAfter(
   nodeOrToken: NodeOrToken | Comment,
   countOptions?: CountOptions | number | FilterFn | null,
 ): Token[] {
-  throw new Error('`sourceCode.getTokensAfter` not implemented yet'); // TODO
+  if (tokens === null) initTokens();
+  debugAssertIsNonNull(tokens);
+  debugAssertIsNonNull(comments);
+
+  const count =
+    typeof countOptions === 'number'
+      ? countOptions
+      : typeof countOptions === 'object' && countOptions !== null
+        ? countOptions.count
+        : null;
+
+  const filter =
+    typeof countOptions === 'function'
+      ? countOptions
+      : typeof countOptions === 'object' && countOptions !== null
+        ? countOptions.filter
+        : null;
+
+  const includeComments =
+    typeof countOptions === 'object' &&
+    countOptions !== null &&
+    'includeComments' in countOptions &&
+    countOptions.includeComments;
+
+  let nodeTokens: Token[];
+  if (includeComments) {
+    if (tokensWithComments === null) {
+      tokensWithComments = [...tokens, ...comments].sort((a, b) => a.range[0] - b.range[0]);
+    }
+    nodeTokens = tokensWithComments;
+  } else {
+    nodeTokens = tokens;
+  }
+
+  const rangeEnd = nodeOrToken.range[1];
+
+  let sliceStart = nodeTokens.length;
+  for (let lo = 0; lo < sliceStart; ) {
+    const mid = (lo + sliceStart) >> 1;
+    if (nodeTokens[mid].range[0] < rangeEnd) {
+      lo = mid + 1;
+    } else {
+      sliceStart = mid;
+    }
+  }
+
+  let nodeTokensAfter: Token[];
+  // Fast path for the common case
+  if (typeof filter !== 'function') {
+    if (typeof count !== 'number') {
+      nodeTokensAfter = nodeTokens.slice(sliceStart);
+    } else {
+      nodeTokensAfter = nodeTokens.slice(sliceStart, sliceStart + count);
+    }
+  } else {
+    if (typeof count !== 'number') {
+      nodeTokensAfter = [];
+      for (let i = sliceStart; i < nodeTokens.length; i++) {
+        const token = nodeTokens[i];
+        if (filter(token)) {
+          nodeTokensAfter.push(token);
+        }
+      }
+    } else {
+      nodeTokensAfter = [];
+      for (let i = sliceStart; i < nodeTokens.length; i++) {
+        const token = nodeTokens[i];
+        if (filter(token)) {
+          nodeTokensAfter.push(token);
+        }
+        if (nodeTokensAfter.length === count) {
+          break;
+        }
+      }
+    }
+  }
+
+  return nodeTokensAfter;
 }
-/* oxlint-enable no-unused-vars */
 
 /**
  * Get all of the tokens between two non-overlapping nodes.

apps/oxlint/test/tokens.test.ts

@@ -20,6 +20,11 @@ vi.mock('../src-js/plugins/source_code.ts', async (importOriginal) => {
     get sourceText() {
       return sourceText;
     },
+    resetSourceAndAst() {
+      // TODO: refactor this quick fix to get the tests working
+      original.resetSourceAndAst();
+      sourceText = '/*A*/var answer/*B*/=/*C*/a/*D*/* b/*E*///F\n    call();\n/*Z*/';
+    },
   };
 });
 
@@ -305,11 +310,91 @@ describe('when calling getTokenAfter', () => {
   getTokenAfter;
 });
 
+// https://github.com/eslint/eslint/blob/v9.39.1/tests/lib/languages/js/source-code/token-store.js#L363-L459
 describe('when calling getTokensAfter', () => {
-  /* oxlint-disable-next-line no-disabled-tests expect-expect */
-  it('is to be implemented');
-  /* oxlint-disable-next-line no-unused-expressions */
-  getTokensAfter;
+  it('should retrieve zero tokens after a node', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, 0).map((token) => token.value),
+      [],
+    );
+  });
+
+  it('should retrieve one token after a node', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, 1).map((token) => token.value),
+      ['='],
+    );
+  });
+
+  it('should retrieve more than one token after a node', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, 2).map((token) => token.value),
+      ['=', 'a'],
+    );
+  });
+
+  it('should retrieve all tokens after a node', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, 9e9).map((token) => token.value),
+      ['=', 'a', '*', 'b', 'call', '(', ')', ';'],
+    );
+  });
+
+  it('should retrieve more than one token after a node with count option', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, { count: 2 }).map((token) => token.value),
+      ['=', 'a'],
+    );
+  });
+
+  it('should retrieve all matched tokens after a node with filter option', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, {
+        filter: (t) => t.type === 'Identifier',
+      }).map((token) => token.value),
+      ['a', 'b', 'call'],
+    );
+  });
+
+  it('should retrieve matched tokens after a node with count and filter options', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, {
+        count: 2,
+        filter: (t) => t.type === 'Identifier',
+      }).map((token) => token.value),
+      ['a', 'b'],
+    );
+  });
+
+  it('should retrieve all tokens and comments after a node with includeComments option', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, {
+        includeComments: true,
+      }).map((token) => token.value),
+      ['B', '=', 'C', 'a', 'D', '*', 'b', 'E', 'F', 'call', '(', ')', ';', 'Z'],
+    );
+  });
+
+  it('should retrieve several tokens and comments after a node with includeComments and count options', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, {
+        includeComments: true,
+        count: 3,
+      }).map((token) => token.value),
+      ['B', '=', 'C'],
+    );
+  });
+
+  it('should retrieve matched tokens and comments after a node with includeComments and count and filter options', () => {
+    assert.deepStrictEqual(
+      getTokensAfter(VariableDeclaratorIdentifier, {
+        includeComments: true,
+        count: 3,
+        filter: (t) => t.type.startsWith('Block'),
+      }).map((token) => token.value),
+      ['B', 'C', 'D'],
+    );
+  });
 });
 
 describe('when calling getFirstTokens', () => {