From 4605d89064e6a43bb156c1885e5ae7760aa7598d Mon Sep 17 00:00:00 2001
From: Daniel Rosenwasser <DanielRosenwasser@users.noreply.github.com>
Date: Mon, 22 Aug 2022 17:30:44 -0700
Subject: [PATCH] Use a mapped type to enforce type-safety on `forEachChild`.
 (#50409)

---
 src/compiler/parser.ts        |  3 +-
 src/compiler/types.ts         | 52 ++++++++++++++++++++++++++++++++++-
 src/compiler/visitorPublic.ts |  2 +-
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/src/compiler/parser.ts b/src/compiler/parser.ts
index 1b82ae4d9a335..93939b6f4f2d1 100644
--- a/src/compiler/parser.ts
+++ b/src/compiler/parser.ts
@@ -96,7 +96,8 @@ namespace ts {
     }
 
     type ForEachChildFunction<TNode> = <T>(node: TNode, cbNode: (node: Node) => T | undefined, cbNodes?: (nodes: NodeArray<Node>) => T | undefined) => T | undefined;
-    const forEachChildTable = {
+    type ForEachChildTable = { [TNode in ForEachChildNodes as TNode["kind"]]: ForEachChildFunction<TNode> };
+    const forEachChildTable: ForEachChildTable = {
         [SyntaxKind.QualifiedName]: function forEachChildInQualifiedName<T>(node: QualifiedName, cbNode: (node: Node) => T | undefined, _cbNodes?: (nodes: NodeArray<Node>) => T | undefined): T | undefined {
             return visitNode(cbNode, node.left) ||
                 visitNode(cbNode, node.right);
diff --git a/src/compiler/types.ts b/src/compiler/types.ts
index 45b66261eb46e..461a750028356 100644
--- a/src/compiler/types.ts
+++ b/src/compiler/types.ts
@@ -885,9 +885,59 @@ namespace ts {
         /* @internal */ jsDocCache?: readonly JSDocTag[];     // Cache for getJSDocTags
     }
 
+    // Ideally, `ForEachChildNodes` and `VisitEachChildNodes` would not differ.
+    // However, `forEachChild` currently processes JSDoc comment syntax and missing declarations more thoroughly.
+    // On the other hand, `visitEachChild` actually processes `Identifier`s (which really *shouldn't* have children,
+    // but are constructed as if they could for faked-up `QualifiedName`s in the language service.)
+
     /* @internal */
-    export type HasChildren =
+    export type ForEachChildNodes =
+        | HasChildren
+        | MissingDeclaration
+        | JSDocTypeExpression
+        | JSDocNonNullableType
+        | JSDocNullableType
+        | JSDocOptionalType
+        | JSDocVariadicType
+        | JSDocFunctionType
+        | JSDoc
+        | JSDocSeeTag
+        | JSDocNameReference
+        | JSDocMemberName
+        | JSDocParameterTag
+        | JSDocPropertyTag
+        | JSDocAuthorTag
+        | JSDocImplementsTag
+        | JSDocAugmentsTag
+        | JSDocTemplateTag
+        | JSDocTypedefTag
+        | JSDocCallbackTag
+        | JSDocReturnTag
+        | JSDocTypeTag
+        | JSDocThisTag
+        | JSDocEnumTag
+        | JSDocSignature
+        | JSDocLink
+        | JSDocLinkCode
+        | JSDocLinkPlain
+        | JSDocTypeLiteral
+        | JSDocUnknownTag
+        | JSDocClassTag
+        | JSDocPublicTag
+        | JSDocPrivateTag
+        | JSDocProtectedTag
+        | JSDocReadonlyTag
+        | JSDocDeprecatedTag
+        ;
+
+    /* @internal */
+    export type VisitEachChildNodes =
+        | HasChildren
         | Identifier
+        ;
+
+    /* @internal */
+    export type HasChildren =
         | QualifiedName
         | ComputedPropertyName
         | TypeParameterDeclaration
diff --git a/src/compiler/visitorPublic.ts b/src/compiler/visitorPublic.ts
index 2541b488bfd50..29a112b062294 100644
--- a/src/compiler/visitorPublic.ts
+++ b/src/compiler/visitorPublic.ts
@@ -403,7 +403,7 @@ namespace ts {
     //  }
     //
     // This is then used as the expected type for `visitEachChildTable`.
-    type VisitEachChildTable = { [TNode in HasChildren as TNode["kind"]]: VisitEachChildFunction<TNode> };
+    type VisitEachChildTable = { [TNode in VisitEachChildNodes as TNode["kind"]]: VisitEachChildFunction<TNode> };
 
     // NOTE: Before you can add a new method to `visitEachChildTable`, you must first ensure the `Node` subtype you
     //       wish to add is defined in the `HasChildren` union in types.ts.