Set parent pointers as nodes are ready in parser, rather than all at once as parse is complete

internal/ast/ast.go

@@ -2,6 +2,7 @@ package ast
 
 import (
 	"fmt"
+	"iter"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -224,6 +225,7 @@ func (n *Node) AsNode() *Node                             { return n }
 func (n *Node) Pos() int                                  { return n.Loc.Pos() }
 func (n *Node) End() int                                  { return n.Loc.End() }
 func (n *Node) ForEachChild(v Visitor) bool               { return n.data.ForEachChild(v) }
+func (n *Node) IterChildren() iter.Seq[*Node]             { return n.data.IterChildren() }
 func (n *Node) Clone(f NodeFactoryCoercible) *Node        { return n.data.Clone(f) }
 func (n *Node) VisitEachChild(v *NodeVisitor) *Node       { return n.data.VisitEachChild(v) }
 func (n *Node) Name() *DeclarationName                    { return n.data.Name() }
@@ -1655,6 +1657,7 @@ func (n *Node) AsSyntaxList() *SyntaxList {
 type nodeData interface {
 	AsNode() *Node
 	ForEachChild(v Visitor) bool
+	IterChildren() iter.Seq[*Node]
 	VisitEachChild(v *NodeVisitor) *Node
 	Clone(v NodeFactoryCoercible) *Node
 	Name() *DeclarationName
@@ -1681,8 +1684,21 @@ type NodeDefault struct {
 	Node
 }
 
-func (node *NodeDefault) AsNode() *Node                                     { return &node.Node }
-func (node *NodeDefault) ForEachChild(v Visitor) bool                       { return false }
+func invert(yield func(v *Node) bool) Visitor {
+	return func(n *Node) bool {
+		return !yield(n)
+	}
+}
+
+func (node *NodeDefault) AsNode() *Node               { return &node.Node }
+func (node *NodeDefault) ForEachChild(v Visitor) bool { return false }
+func (node *NodeDefault) forEachChildIter(yield func(v *Node) bool) {
+	node.data.ForEachChild(invert(yield)) // `true` is return early for a ts visitor, `false` is return early for a go iterator yield function
+}
+
+func (node *NodeDefault) IterChildren() iter.Seq[*Node] {
+	return node.forEachChildIter
+}
 func (node *NodeDefault) VisitEachChild(v *NodeVisitor) *Node               { return node.AsNode() }
 func (node *NodeDefault) Clone(v NodeFactoryCoercible) *Node                { return nil }
 func (node *NodeDefault) Name() *DeclarationName                            { return nil }

internal/checker/checker.go

@@ -30198,7 +30198,7 @@ func (c *Checker) getSymbolOfNameOrPropertyAccessExpression(name *ast.Node) *ast
 		return c.getSymbolOfNode(name.Parent)
 	}
 
-	if name.Parent.Kind == ast.KindExportAssignment && ast.IsEntityNameExpression(name) {
+	if (name.Parent.Kind == ast.KindExportAssignment || name.Parent.Kind == ast.KindJSExportAssignment) && ast.IsEntityNameExpression(name) {
 		// Even an entity name expression that doesn't resolve as an entityname may still typecheck as a property access expression
 		success := c.resolveEntityName(
 			name,

internal/compiler/fileloader.go

@@ -561,13 +561,13 @@ func getDefaultResolutionModeForFile(fileName string, meta ast.SourceFileMetaDat
 }
 
 func getModeForUsageLocation(fileName string, meta ast.SourceFileMetaData, usage *ast.StringLiteralLike, options *core.CompilerOptions) core.ResolutionMode {
-	if ast.IsImportDeclaration(usage.Parent) || ast.IsExportDeclaration(usage.Parent) || ast.IsJSDocImportTag(usage.Parent) {
+	if ast.IsImportDeclaration(usage.Parent) || usage.Parent.Kind == ast.KindJSImportDeclaration || ast.IsExportDeclaration(usage.Parent) || ast.IsJSDocImportTag(usage.Parent) {
 		isTypeOnly := ast.IsExclusivelyTypeOnlyImportOrExport(usage.Parent)
 		if isTypeOnly {
 			var override core.ResolutionMode
 			var ok bool
 			switch usage.Parent.Kind {
-			case ast.KindImportDeclaration:
+			case ast.KindImportDeclaration, ast.KindJSImportDeclaration:
 				override, ok = usage.Parent.AsImportDeclaration().Attributes.GetResolutionModeOverride()
 			case ast.KindExportDeclaration:
 				override, ok = usage.Parent.AsExportDeclaration().Attributes.GetResolutionModeOverride()

internal/parser/jsdoc.go

@@ -44,6 +44,7 @@ func (p *Parser) withJSDoc(node *ast.Node, hasJSDoc bool) []*ast.Node {
 	pos := node.Pos()
 	for _, comment := range ranges {
 		if parsed := p.parseJSDocComment(node, comment.Pos(), comment.End(), pos); parsed != nil {
+			parsed.Parent = node
 			jsdoc = append(jsdoc, parsed)
 			pos = parsed.End()
 		}
@@ -977,6 +978,9 @@ func (p *Parser) parseTypedefTag(start int, tagName *ast.IdentifierNode, indent
 
 	typedefTag := p.factory.NewJSDocTypedefTag(tagName, typeExpression, fullName, comment)
 	p.finishNodeWithEnd(typedefTag, start, end)
+	if typeExpression != nil {
+		typeExpression.Parent = typedefTag // forcibly overwrite parent potentially set by inner type expression parse
+	}
 	return typedefTag
 }
 

internal/parser/parser.go

@@ -78,13 +78,22 @@ type Parser struct {
 	jsdocTagCommentsSpace   []string
 	reparseList             []*ast.Node
 	commonJSModuleIndicator *ast.Node
+
+	currentParent        *ast.Node
+	setParentFromContext ast.Visitor
+}
+
+func newParser() *Parser {
+	res := &Parser{}
+	res.initializeClosures()
+	return res
 }
 
 var viableKeywordSuggestions = scanner.GetViableKeywordSuggestions()
 
 var parserPool = sync.Pool{
 	New: func() any {
-		return &Parser{}
+		return newParser()
 	},
 }
 
@@ -93,7 +102,7 @@ func getParser() *Parser {
 }
 
 func putParser(p *Parser) {
-	*p = Parser{scanner: p.scanner}
+	*p = Parser{scanner: p.scanner, setParentFromContext: p.setParentFromContext}
 	parserPool.Put(p)
 }
 
@@ -108,6 +117,13 @@ func ParseSourceFile(opts ast.SourceFileParseOptions, sourceText string, scriptK
 	return p.parseSourceFileWorker()
 }
 
+func (p *Parser) initializeClosures() {
+	p.setParentFromContext = func(n *ast.Node) bool {
+		n.Parent = p.currentParent
+		return false
+	}
+}
+
 func (p *Parser) parseJSONText() *ast.SourceFile {
 	pos := p.nodePos()
 	var statements *ast.NodeList
@@ -354,14 +370,6 @@ func (p *Parser) finishSourceFile(result *ast.SourceFile, isDeclarationFile bool
 	result.IdentifierCount = p.identifierCount
 	result.SetJSDocCache(p.jsdocCache)
 
-	ast.SetParentInChildren(result.AsNode())
-	for parent, children := range p.jsdocCache {
-		for _, child := range children {
-			child.Parent = parent
-			ast.SetParentInChildren(child)
-		}
-	}
-
 	ast.SetExternalModuleIndicator(result, p.opts.ExternalModuleIndicatorOptions)
 }
 
@@ -467,7 +475,11 @@ func (p *Parser) reparseTopLevelAwait(sourceFile *ast.SourceFile) *ast.Node {
 		}
 	}
 
-	return p.factory.NewSourceFile(sourceFile.ParseOptions(), p.sourceText, p.newNodeList(sourceFile.Statements.Loc, statements))
+	result := p.factory.NewSourceFile(sourceFile.ParseOptions(), p.sourceText, p.newNodeList(sourceFile.Statements.Loc, statements))
+	for _, s := range statements {
+		s.Parent = result.AsNode() // force (re)set parent to reparsed source file
+	}
+	return result
 }
 
 func (p *Parser) parseListIndex(kind ParsingContext, parseElement func(p *Parser, index int) *ast.Node) *ast.NodeList {
@@ -3572,6 +3584,7 @@ func (p *Parser) parseTupleElementType() *ast.TypeNode {
 		node := p.factory.NewOptionalTypeNode(typeNode.Type())
 		node.Flags = typeNode.Flags
 		node.Loc = typeNode.Loc
+		typeNode.Type().Parent = node
 		return node
 	}
 	return typeNode
@@ -4693,6 +4706,16 @@ func (p *Parser) parseJsxElementOrSelfClosingElementOrFragment(inExpressionConte
 			p.finishNodeWithEnd(newClosingElement, end, end)
 			newLast := p.factory.NewJsxElement(lastChild.AsJsxElement().OpeningElement, lastChild.AsJsxElement().Children, newClosingElement)
 			p.finishNodeWithEnd(newLast, lastChild.AsJsxElement().OpeningElement.Pos(), end)
+			// force reset parent pointers from discarded parse result
+			if lastChild.AsJsxElement().OpeningElement != nil {
+				lastChild.AsJsxElement().OpeningElement.Parent = newLast
+			}
+			if lastChild.AsJsxElement().Children != nil {
+				for _, c := range lastChild.AsJsxElement().Children.Nodes {
+					c.Parent = newLast
+				}
+			}
+			newClosingElement.Parent = newLast
 			children = p.newNodeList(core.NewTextRange(children.Pos(), newLast.End()), append(children.Nodes[0:len(children.Nodes)-1], newLast))
 			closingElement = lastChild.AsJsxElement().ClosingElement
 		} else {
@@ -4709,6 +4732,7 @@ func (p *Parser) parseJsxElementOrSelfClosingElementOrFragment(inExpressionConte
 		}
 		result = p.factory.NewJsxElement(opening, children, closingElement)
 		p.finishNode(result, pos)
+		closingElement.Parent = result // force reset parent pointers from possibly discarded parse result
 	case ast.KindJsxOpeningFragment:
 		result = p.factory.NewJsxFragment(opening, p.parseJsxChildren(opening), p.parseJsxClosingFragment(inExpressionContext))
 		p.finishNode(result, pos)
@@ -5316,7 +5340,9 @@ func (p *Parser) parseMemberExpressionRest(pos int, expression *ast.Expression,
 		if p.isTemplateStartOfTaggedTemplate() {
 			// Absorb type arguments into TemplateExpression when preceding expression is ExpressionWithTypeArguments
 			if questionDotToken == nil && ast.IsExpressionWithTypeArguments(expression) {
-				expression = p.parseTaggedTemplateRest(pos, expression.AsExpressionWithTypeArguments().Expression, questionDotToken, expression.AsExpressionWithTypeArguments().TypeArguments)
+				original := expression.AsExpressionWithTypeArguments()
+				expression = p.parseTaggedTemplateRest(pos, original.Expression, questionDotToken, original.TypeArguments)
+				p.unparseExpressionWithTypeArguments(original.Expression, original.TypeArguments, expression)
 			} else {
 				expression = p.parseTaggedTemplateRest(pos, expression, questionDotToken, nil /*typeArguments*/)
 			}
@@ -5431,10 +5457,12 @@ func (p *Parser) parseCallExpressionRest(pos int, expression *ast.Expression) *a
 				typeArguments = expression.AsExpressionWithTypeArguments().TypeArguments
 				expression = expression.AsExpressionWithTypeArguments().Expression
 			}
+			inner := expression
 			argumentList := p.parseArgumentList()
 			isOptionalChain := questionDotToken != nil || p.tryReparseOptionalChain(expression)
 			expression = p.factory.NewCallExpression(expression, questionDotToken, typeArguments, argumentList, core.IfElse(isOptionalChain, ast.NodeFlagsOptionalChain, ast.NodeFlagsNone))
 			p.finishNode(expression, pos)
+			p.unparseExpressionWithTypeArguments(inner, typeArguments, expression)
 			continue
 		}
 		if questionDotToken != nil {
@@ -5715,6 +5743,18 @@ func (p *Parser) parseDecoratedExpression() *ast.Expression {
 	return result
 }
 
+func (p *Parser) unparseExpressionWithTypeArguments(expression *ast.Node, typeArguments *ast.NodeList, result *ast.Node) {
+	// force overwrite the `.Parent` of the expression and type arguments to erase the fact that they may have originally been parsed as an ExpressionWithTypeArguments and be parented to such
+	if expression != nil {
+		expression.Parent = result
+	}
+	if typeArguments != nil {
+		for _, a := range typeArguments.Nodes {
+			a.Parent = result
+		}
+	}
+}
+
 func (p *Parser) parseNewExpressionOrNewDotTarget() *ast.Node {
 	pos := p.nodePos()
 	p.parseExpected(ast.KindNewKeyword)
@@ -5741,6 +5781,7 @@ func (p *Parser) parseNewExpressionOrNewDotTarget() *ast.Node {
 	}
 	result := p.factory.NewNewExpression(expression, typeArguments, argumentList)
 	p.finishNode(result, pos)
+	p.unparseExpressionWithTypeArguments(expression, typeArguments, result)
 	return result
 }
 
@@ -5889,6 +5930,13 @@ func (p *Parser) finishNodeWithEnd(node *ast.Node, pos int, end int) {
 		node.Flags |= ast.NodeFlagsThisNodeHasError
 		p.hasParseError = false
 	}
+	p.overrideParentInImmediateChildren(node)
+}
+
+func (p *Parser) overrideParentInImmediateChildren(node *ast.Node) {
+	p.currentParent = node
+	node.ForEachChild(p.setParentFromContext)
+	p.currentParent = nil
 }
 
 func (p *Parser) nextTokenIsSlash() bool {