Reject visitor removal of required children

tools/scripts/tsc/generate-go-ast.ts

@@ -996,7 +996,7 @@ function generateVisitEachChild(w: CodeWriter, node: NodeType) {
         if (m.isChild() && type.kind === "list" && type.listKind === "raw") {
             const localName = m.goParamName();
             rawListLocals.set(m.name, localName);
-            w.write(`${localName} := core.SameMap(node.${m.name}, func(n *Node) *Node { return v.visitNode(n) })`);
+            w.write(`${localName} := core.SameMap(node.${m.name}, func(n *Node) *Node { return requiredChild(n, v.visitNode(n)) })`);
         }
     }
 
@@ -1014,17 +1014,13 @@ function generateVisitEachChild(w: CodeWriter, node: NodeType) {
         const type = m.type;
         const listKind = m.listKind;
 
-        if (m.visit) {
-            return `v.visit${api.capitalize(m.visit)}(${access})`;
-        }
-        if (listKind === "ModifierList") {
-            return `v.visitModifiers(${access})`;
-        }
-        if (listKind === "NodeList") {
-            return `v.visitNodes(${access})`;
-        }
-        if (type.kind === "primitive" && type.name === "Token") return `v.visitToken(${access})`;
-        return `v.visitNode(${access})`;
+        const visit = m.visit ? `visit${api.capitalize(m.visit)}`
+            : listKind === "ModifierList" ? "visitModifiers"
+            : listKind === "NodeList" ? "visitNodes"
+            : type.kind === "primitive" && type.name === "Token" ? "visitToken"
+            : "visitNode";
+        const result = `v.${visit}(${access})`;
+        return m.optional ? result : `requiredChild(${access}, ${result})`;
     }).join(", ");
 
     w.write(`return v.Factory.Update${node.name}(node, ${args})`);

tsc/internal/ast/ast.go

@@ -2718,7 +2718,7 @@ func (node *SourceFile) ForEachChild(v Visitor) bool {
 }
 
 func (node *SourceFile) VisitEachChild(v *NodeVisitor) *Node {
-	return v.Factory.UpdateSourceFile(node, v.visitTopLevelStatements(node.Statements), v.visitToken(node.EndOfFileToken))
+	return v.Factory.UpdateSourceFile(node, requiredChild(node.Statements, v.visitTopLevelStatements(node.Statements)), requiredChild(node.EndOfFileToken, v.visitToken(node.EndOfFileToken)))
 }
 
 func (node *SourceFile) IsJS() bool {
@@ -3113,7 +3113,7 @@ func forEachChild_JSDocParameterOrPropertyTag(node *JSDocParameterOrPropertyTag,
 }
 
 func visitEachChild_JSDocParameterOrPropertyTag(node *JSDocParameterOrPropertyTag, v *NodeVisitor) *Node {
-	return v.Factory.UpdateJSDocParameterOrPropertyTag(node, v.visitNode(node.TagName), v.visitNode(node.name), node.IsBracketed, v.visitNode(node.TypeExpression), node.IsNameFirst, v.visitNodes(node.Comment))
+	return v.Factory.UpdateJSDocParameterOrPropertyTag(node, requiredChild(node.TagName, v.visitNode(node.TagName)), requiredChild(node.name, v.visitNode(node.name)), node.IsBracketed, v.visitNode(node.TypeExpression), node.IsNameFirst, v.visitNodes(node.Comment))
 }
 
 func (f *NodeFactory) ReleaseArenas() {

tsc/internal/ast/deepclone_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/microsoft/TypeScript/tsc/internal/ast"
+	"github.com/microsoft/TypeScript/tsc/internal/testutil"
 	"github.com/microsoft/TypeScript/tsc/internal/testutil/parsetestutil"
 	"gotest.tools/v3/assert"
 )
@@ -22,6 +23,58 @@ func getChildren(node *ast.Node) []*ast.Node {
 	return children
 }
 
+func TestVisitEachChildRequiredChild(t *testing.T) {
+	t.Parallel()
+	factory := ast.NewNodeFactory(ast.NodeFactoryHooks{})
+	for _, test := range []struct {
+		name  string
+		node  *ast.Node
+		kind  ast.Kind
+		hooks ast.NodeVisitorHooks
+	}{
+		{name: "node", node: factory.NewExpressionStatement(factory.NewIdentifier("x")), kind: ast.KindIdentifier},
+		{name: "embedded statement", node: factory.NewIfStatement(factory.NewIdentifier("x"), factory.NewEmptyStatement(), nil), kind: ast.KindEmptyStatement},
+		{name: "raw list", node: factory.NewSyntaxList([]*ast.Node{factory.NewIdentifier("x")}), kind: ast.KindIdentifier},
+		{
+			name: "hook",
+			node: factory.NewExpressionStatement(factory.NewIdentifier("x")),
+			hooks: ast.NodeVisitorHooks{
+				VisitNode: func(node *ast.Node, v *ast.NodeVisitor) *ast.Node { return nil },
+			},
+		},
+		{
+			name: "list hook",
+			node: factory.NewBlock(factory.NewNodeList(nil), false),
+			hooks: ast.NodeVisitorHooks{
+				VisitNodes: func(nodes *ast.NodeList, v *ast.NodeVisitor) *ast.NodeList { return nil },
+			},
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+			visitor := ast.NewNodeVisitor(func(node *ast.Node) *ast.Node {
+				if node.Kind == test.kind {
+					return nil
+				}
+				return node
+			}, factory, test.hooks)
+			testutil.AssertPanics(t, func() {
+				visitor.VisitEachChild(test.node)
+			}, "Visitor returned nil for a required child")
+		})
+	}
+}
+
+func TestVisitEachChildElision(t *testing.T) {
+	t.Parallel()
+	factory := ast.NewNodeFactory(ast.NodeFactoryHooks{})
+	node := factory.NewReturnStatement(factory.NewIdentifier("x"))
+	visitor := ast.NewNodeVisitor(func(node *ast.Node) *ast.Node { return nil }, factory, ast.NodeVisitorHooks{})
+	assert.Assert(t, visitor.VisitEachChild(node).AsReturnStatement().Expression == nil)
+	block := factory.NewBlock(factory.NewNodeList([]*ast.Node{node}), false)
+	assert.Equal(t, len(visitor.VisitEachChild(block).AsBlock().Statements.Nodes), 0)
+}
+
 func TestDeepCloneNodeSanityCheck(t *testing.T) {
 	t.Parallel()
 	data := []struct {

tsc/internal/ast/visitor.go

@@ -183,6 +183,7 @@ func (v *NodeVisitor) VisitSlice(nodes []*Node) (result []*Node, changed bool) {
 }
 
 // Visits each child of a Node, possibly returning a new Node of the same kind in its place.
+// Panics if visiting a non-nil required child returns nil.
 func (v *NodeVisitor) VisitEachChild(node *Node) *Node {
 	if node == nil || v.Visit == nil {
 		return node
@@ -191,6 +192,13 @@ func (v *NodeVisitor) VisitEachChild(node *Node) *Node {
 	return node.VisitEachChild(v)
 }
 
+func requiredChild[T any](original *T, child *T) *T {
+	if original != nil && child == nil {
+		panic("Visitor returned nil for a required child")
+	}
+	return child
+}
+
 func (v *NodeVisitor) visitNode(node *Node) *Node {
 	if v.Hooks.VisitNode != nil {
 		return v.Hooks.VisitNode(node, v)

tsc/internal/astnav/tokens.go

@@ -103,7 +103,7 @@ func getTokenAtPosition(
 		// We can't abort visiting children, so once a match is found, we set `next`
 		// and do nothing on subsequent visits.
 		if node == nil || node.Flags&ast.NodeFlagsReparsed != 0 {
-			return nil
+			return node
 		}
 		if nodeAfterLeft == nil {
 			nodeAfterLeft = node

tsc/internal/transformers/declarations/transform.go

@@ -2110,13 +2110,13 @@ func (tx *DeclarationTransformer) visitThisPropertyAssignments(node *ast.Node) *
 	thisContainer := ast.GetThisContainer(node, false, false)
 	thisTarget = thisContainer.Parent
 	if thisTarget == nil {
-		return nil // thisContainer was source file, can't have expando-this
+		return node // thisContainer was source file, can't have expando-this
 	}
 	if ast.HasStaticModifier(thisContainer) || ast.IsClassStaticBlockDeclaration(thisContainer) {
 		isStatic = true
 	}
 	if thisTarget != tx.enclosingDeclaration {
-		return nil // stop searching within new `this` contexts
+		return node // stop searching within new `this` contexts
 	}
 caseBlock:
 	switch ast.GetAssignmentDeclarationKind(node) {

tsc/internal/transformers/tstransforms/typeeraser.go

@@ -92,6 +92,11 @@ func (tx *TypeEraserTransformer) visit(node *ast.Node) *ast.Node {
 		ast.KindIndexedAccessType,
 		ast.KindMappedType,
 		ast.KindLiteralType,
+		ast.KindImportType,
+		ast.KindJSDocNullableType,
+		ast.KindJSDocNonNullableType,
+		ast.KindJSDocOptionalType,
+		ast.KindJSDocVariadicType,
 		// TypeScript index signatures are elided.
 		ast.KindIndexSignature:
 		return nil
@@ -394,6 +399,10 @@ func (tx *TypeEraserTransformer) visit(node *ast.Node) *ast.Node {
 		}
 		return node
 
+	case ast.KindExportAssignment:
+		n := node.AsExportAssignment()
+		return tx.Factory().UpdateExportAssignment(n, tx.Visitor().VisitModifiers(n.Modifiers()), n.IsExportEquals, nil, tx.Visitor().VisitNode(n.Expression))
+
 	case ast.KindEnumDeclaration:
 		if !ast.IsEnumConst(node) {
 			return tx.Visitor().VisitEachChild(node)

tsc/internal/transformers/tstransforms/typeeraser_test.go

@@ -3,14 +3,63 @@ package tstransforms_test
 import (
 	"testing"
 
+	"github.com/microsoft/TypeScript/tsc/internal/ast"
 	"github.com/microsoft/TypeScript/tsc/internal/core"
+	"github.com/microsoft/TypeScript/tsc/internal/parser"
 	"github.com/microsoft/TypeScript/tsc/internal/printer"
 	"github.com/microsoft/TypeScript/tsc/internal/testutil/emittestutil"
 	"github.com/microsoft/TypeScript/tsc/internal/testutil/parsetestutil"
 	"github.com/microsoft/TypeScript/tsc/internal/transformers"
 	"github.com/microsoft/TypeScript/tsc/internal/transformers/tstransforms"
+	"gotest.tools/v3/assert"
 )
 
+func TestTypeEraserJSDocExportAssignment(t *testing.T) {
+	t.Parallel()
+	for _, typeText := range []string{`import("./a").Foo`, `typeof import("./a")`} {
+		t.Run(typeText, func(t *testing.T) {
+			t.Parallel()
+			file := parser.ParseSourceFile(ast.SourceFileParseOptions{
+				FileName: "/main.js",
+			}, "/** @type {"+typeText+"} */\nexport default foo;", core.ScriptKindJS)
+			parsetestutil.CheckDiagnostics(t, file)
+			assert.Assert(t, file.Statements.Nodes[0].AsExportAssignment().Type != nil)
+			transformed := tstransforms.NewTypeEraserTransformer(&transformers.TransformOptions{
+				CompilerOptions: &core.CompilerOptions{},
+				Context:         printer.NewEmitContext(),
+			}).TransformSourceFile(file)
+			assert.Assert(t, transformed.Statements.Nodes[0].AsExportAssignment().Type == nil)
+			emittestutil.CheckEmit(t, nil, transformed, "/** @type {"+typeText+"} */\nexport default foo;")
+		})
+	}
+}
+
+func TestTypeEraserJSDocWrappedTypes(t *testing.T) {
+	t.Parallel()
+	for _, test := range []struct {
+		name  string
+		input string
+	}{
+		{name: "nullable property", input: "const obj = {\n    /** @type {?number} */\n    p: null\n};"},
+		{name: "non-nullable shorthand", input: "const p = 1;\nconst obj = {\n    /** @type {!number} */\n    p\n};"},
+		{name: "optional assignment", input: "const obj = {};\n/** @type {number=} */\nobj.p = null;"},
+		{name: "variadic property", input: "const obj = {\n    /** @type {...number} */\n    p: []\n};"},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+			file := parser.ParseSourceFile(ast.SourceFileParseOptions{
+				FileName: "/main.js",
+			}, test.input, core.ScriptKindJS)
+			parsetestutil.CheckDiagnostics(t, file)
+			transformed := tstransforms.NewTypeEraserTransformer(&transformers.TransformOptions{
+				CompilerOptions: &core.CompilerOptions{},
+				Context:         printer.NewEmitContext(),
+			}).TransformSourceFile(file)
+			emittestutil.CheckEmit(t, nil, transformed, test.input)
+		})
+	}
+}
+
 func TestTypeEraser(t *testing.T) {
 	t.Parallel()
 	data := []struct {