From ba7c27840733ff33baf3a603b5ac13a49c40c43f Mon Sep 17 00:00:00 2001
From: Quramy <yosuke.kurami@gmail.com>
Date: Wed, 13 Mar 2024 23:50:14 +0900
Subject: [PATCH] chore: Allow findAllNodes callback to return generic AST
 nodes

---
 src/ts-ast-util/utilily-functions.test.ts | 14 +++++++++++++-
 src/ts-ast-util/utilily-functions.ts      | 14 +++++++++-----
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/ts-ast-util/utilily-functions.test.ts b/src/ts-ast-util/utilily-functions.test.ts
index c8251f281..14ffad70b 100644
--- a/src/ts-ast-util/utilily-functions.test.ts
+++ b/src/ts-ast-util/utilily-functions.test.ts
@@ -30,7 +30,7 @@ describe(isTagged, () => {
 });
 
 describe(findAllNodes, () => {
-  it('findAllNodes should return nodes which match given condition', () => {
+  it('should return nodes which match given condition', () => {
     // prettier-ignore
     const text = 'const a = `AAA`;' + '\n'
              + 'const b = `BBB`;';
@@ -39,6 +39,18 @@ describe(findAllNodes, () => {
     expect(actual.length).toBe(2);
     expect(actual.map(n => n.getText())).toEqual(['`AAA`', '`BBB`']);
   });
+
+  it('should accepts callback which returns ts.Node', () => {
+    const text = `
+      const a = fn(\`AAA\`);
+      const b = fn('BBB');
+    `;
+    const s = ts.createSourceFile('input.ts', text, ts.ScriptTarget.Latest);
+    const actual = findAllNodes(s, node => (ts.isCallExpression(node) ? node.arguments[0] : undefined));
+    expect(actual.length).toBe(2);
+    expect(ts.isNoSubstitutionTemplateLiteral(actual[0])).toBeTruthy();
+    expect(ts.isStringLiteral(actual[1])).toBeTruthy();
+  });
 });
 
 describe(isImportDeclarationWithCondition, () => {
diff --git a/src/ts-ast-util/utilily-functions.ts b/src/ts-ast-util/utilily-functions.ts
index 17f9d4d28..57840ac4f 100644
--- a/src/ts-ast-util/utilily-functions.ts
+++ b/src/ts-ast-util/utilily-functions.ts
@@ -48,18 +48,22 @@ export function findNode(sourceFile: ts.SourceFile, position: number): ts.Node |
   return find(sourceFile);
 }
 
-export function findAllNodes(sourceFile: ts.SourceFile, cond: (n: ts.Node) => boolean): ts.Node[] {
-  const result: ts.Node[] = [];
+export function findAllNodes<S extends ts.Node>(
+  sourceFile: ts.SourceFile,
+  cond: (n: ts.Node) => S | boolean | undefined,
+): S[] {
+  const result: (S | ts.Node)[] = [];
   function find(node: ts.Node) {
-    if (cond(node)) {
-      result.push(node);
+    const hit = cond(node);
+    if (hit) {
+      result.push(hit === true ? node : hit);
       return;
     } else {
       ts.forEachChild(node, find);
     }
   }
   find(sourceFile);
-  return result;
+  return result as S[];
 }
 
 export function hasTagged(node: ts.Node | undefined, condition: TagCondition) {