fix forEachChildRecursively (microsoft#42300)

Author: ajafff

src/compiler/parser.ts

@@ -558,61 +558,54 @@ namespace ts {
      * and while doing so, handles traversing the structure without relying on the callstack to encode the tree structure.
      */
     export function forEachChildRecursively<T>(rootNode: Node, cbNode: (node: Node, parent: Node) => T | "skip" | undefined, cbNodes?: (nodes: NodeArray<Node>, parent: Node) => T | "skip" | undefined): T | undefined {
-
-        const stack: Node[] = [rootNode];
-        while (stack.length) {
-            const parent = stack.pop()!;
-            const res = visitAllPossibleChildren(parent, gatherPossibleChildren(parent));
-            if (res) {
-                return res;
-            }
-        }
-
-        return;
-
-        function gatherPossibleChildren(node: Node) {
-            const children: (Node | NodeArray<Node>)[] = [];
-            forEachChild(node, addWorkItem, addWorkItem); // By using a stack above and `unshift` here, we emulate a depth-first preorder traversal
-            return children;
-
-            function addWorkItem(n: Node | NodeArray<Node>) {
-                children.unshift(n);
-            }
-        }
-
-        function visitAllPossibleChildren(parent: Node, children: readonly (Node | NodeArray<Node>)[]) {
-            for (const child of children) {
-                if (isArray(child)) {
-                    if (cbNodes) {
-                        const res = cbNodes(child, parent);
-                        if (res) {
-                            if (res === "skip") continue;
-                            return res;
-                        }
-                    }
-
-                    for (let i = child.length - 1; i >= 0; i--) {
-                        const realChild = child[i];
-                        const res = cbNode(realChild, parent);
-                        if (res) {
-                            if (res === "skip") continue;
-                            return res;
-                        }
-                        stack.push(realChild);
-                    }
-                }
-                else {
-                    stack.push(child);
-                    const res = cbNode(child, parent);
+        const queue: (Node | NodeArray<Node>)[] = gatherPossibleChildren(rootNode);
+        const parents: Node[] = []; // tracks parent references for elements in queue
+        while (parents.length < queue.length) {
+            parents.push(rootNode);
+        }
+        while (queue.length !== 0) {
+            const current = queue.pop()!;
+            const parent = parents.pop()!;
+            if (isArray(current)) {
+                if (cbNodes) {
+                    const res = cbNodes(current, parent);
                     if (res) {
                         if (res === "skip") continue;
                         return res;
                     }
                 }
+                for (let i = current.length - 1; i >= 0; --i) {
+                    queue.push(current[i]);
+                    parents.push(parent);
+                }
+            }
+            else {
+                const res = cbNode(current, parent);
+                if (res) {
+                    if (res === "skip") continue;
+                    return res;
+                }
+                if (current.kind >= SyntaxKind.FirstNode) {
+                    // add children in reverse order to the queue, so popping gives the first child
+                    for (const child of gatherPossibleChildren(current)) {
+                        queue.push(child);
+                        parents.push(current);
+                    }
+                }
             }
         }
     }
 
+    function gatherPossibleChildren(node: Node) {
+        const children: (Node | NodeArray<Node>)[] = [];
+        forEachChild(node, addWorkItem, addWorkItem); // By using a stack above and `unshift` here, we emulate a depth-first preorder traversal
+        return children;
+
+        function addWorkItem(n: Node | NodeArray<Node>) {
+            children.unshift(n);
+        }
+    }
+
     export function createSourceFile(fileName: string, sourceText: string, languageVersion: ScriptTarget, setParentNodes = false, scriptKind?: ScriptKind): SourceFile {
         tracing.push(tracing.Phase.Parse, "createSourceFile", { path: fileName }, /*separateBeginAndEnd*/ true);
         performance.mark("beforeParse");