track next node and avoid scanning into it

Author: gabritto

internal/astnav/tokens.go

@@ -40,6 +40,18 @@ func GetTokenAtPosition(sourceFile *ast.SourceFile, position int) *ast.Node {
 	return getTokenAtPosition(sourceFile, position, true /*allowPositionInLeadingTrivia*/, nil)
 }
 
+func shouldVisitNode(node *ast.Node, allowReparsed bool) bool {
+	// Skip reparsed nodes unless:
+	// 1. The node itself is AsExpression or SatisfiesExpression, OR
+	// 2. We're already inside an AsExpression or SatisfiesExpression (allowReparsed=true)
+	// These are special cases where reparsed nodes from JSDoc type assertions
+	// should still be navigable to reach identifiers.
+	isSpecialReparsed := node.Flags&ast.NodeFlagsReparsed != 0 &&
+		(node.Kind == ast.KindAsExpression || node.Kind == ast.KindSatisfiesExpression)
+
+	return node.Flags&ast.NodeFlagsReparsed == 0 || isSpecialReparsed || allowReparsed
+}
+
 func getTokenAtPosition(
 	sourceFile *ast.SourceFile,
 	position int,
@@ -65,6 +77,9 @@ func getTokenAtPosition(
 	// `left` tracks the lower boundary of the node/token that could be returned,
 	// and is eventually the scanner's start position, if the scanner is used.
 	left := 0
+	// `nodeAfterLeft` tracks the first node we visit after visiting the node that advances `left`.
+	// When scanning in between nodes for token, we should only scan up to the start of `nodeAfterLeft`.
+	var nodeAfterLeft *ast.Node
 	// `allowReparsed` is set when we're navigating inside an AsExpression or
 	// SatisfiesExpression, which allows visiting their reparsed children to reach
 	// the actual identifier from JSDoc type assertions.
@@ -91,16 +106,14 @@ func getTokenAtPosition(
 	visitNode := func(node *ast.Node, _ *ast.NodeVisitor) *ast.Node {
 		// We can't abort visiting children, so once a match is found, we set `next`
 		// and do nothing on subsequent visits.
-		if node != nil && next == nil {
-			// Skip reparsed nodes unless:
-			// 1. The node itself is AsExpression or SatisfiesExpression, OR
-			// 2. We're already inside an AsExpression or SatisfiesExpression (allowReparsed=true)
-			// These are special cases where reparsed nodes from JSDoc type assertions
-			// should still be navigable to reach identifiers.
-			isSpecialReparsed := node.Flags&ast.NodeFlagsReparsed != 0 &&
-				(node.Kind == ast.KindAsExpression || node.Kind == ast.KindSatisfiesExpression)
-
-			if node.Flags&ast.NodeFlagsReparsed == 0 || isSpecialReparsed || allowReparsed {
+		if node == nil {
+			return nil
+		}
+		if nodeAfterLeft == nil {
+			nodeAfterLeft = node
+		}
+		if next == nil {
+			if shouldVisitNode(node, allowReparsed) {
 				result := testNode(node)
 				switch result {
 				case -1:
@@ -111,6 +124,7 @@ func getTokenAtPosition(
 						// all its leading trivia in its position.
 						left = node.End()
 					}
+					nodeAfterLeft = nil
 				case 0:
 					next = node
 				}
@@ -120,12 +134,25 @@ func getTokenAtPosition(
 	}
 
 	visitNodeList := func(nodeList *ast.NodeList, _ *ast.NodeVisitor) *ast.NodeList {
-		if nodeList != nil && len(nodeList.Nodes) > 0 && next == nil {
+		if nodeList == nil || len(nodeList.Nodes) == 0 {
+			return nodeList
+		}
+		if nodeAfterLeft == nil {
+			for _, node := range nodeList.Nodes {
+				if shouldVisitNode(node, allowReparsed) {
+					nodeAfterLeft = node
+					break
+				}
+			}
+		}
+		if next == nil {
 			if nodeList.End() == position && includePrecedingTokenAtEndPosition != nil {
 				left = nodeList.End()
+				nodeAfterLeft = nil
 				prevSubtree = nodeList.Nodes[len(nodeList.Nodes)-1]
 			} else if nodeList.End() <= position {
 				left = nodeList.End()
+				nodeAfterLeft = nil
 			} else if nodeList.Pos() <= position {
 				nodes := nodeList.Nodes
 				index, match := core.BinarySearchUniqueFunc(nodes, func(middle int, node *ast.Node) int {
@@ -135,6 +162,13 @@ func getTokenAtPosition(
 					cmp := testNode(node)
 					if cmp < 0 {
 						left = node.End()
+						nodeAfterLeft = nil
+						for i := middle + 1; i < len(nodes); i++ {
+							if shouldVisitNode(nodes[i], allowReparsed) {
+								nodeAfterLeft = nodes[i]
+								break
+							}
+						}
 					}
 					return cmp
 				})
@@ -147,6 +181,11 @@ func getTokenAtPosition(
 						cmp := testNode(node)
 						if cmp < 0 {
 							left = node.End()
+							if middle+1 < len(nodes) {
+								nodeAfterLeft = nodes[middle+1]
+							} else {
+								nodeAfterLeft = nil
+							}
 						}
 						return cmp
 					})
@@ -182,12 +221,33 @@ func getTokenAtPosition(
 				return current
 			}
 			scanner := scanner.GetScannerForSourceFile(sourceFile, left)
-			for left < current.End() {
+			end := current.End()
+			// We should only scan up to the start of the next node in the AST after the node ending at position `left`.
+			// It is necessary to enforce this invariant in cases where `position` occurs in between two node/tokens,
+			// such that we would not find a token in the loop below before we reach the next node.
+			// We can fall into this case when `allowPositionInLeadingTrivia` is false and `position` is in a leading trivia,
+			// or when `position` would be in the leading trivia of a node but this node is inside JSDoc:
+			// ```
+			// /**
+			//  * @type {{
+			//  */*$*/ identifier: boolean;
+			//  * }}
+			//  */
+			// ```
+			// The position of marker '$' falls in between the asterisk token and the identifier token, but is not
+			// part of the leading trivia for `identifier`.
+			if nodeAfterLeft != nil {
+				end = nodeAfterLeft.Pos()
+			}
+			for left < end {
 				token := scanner.Token()
 				tokenFullStart := scanner.TokenFullStart()
 				tokenStart := core.IfElse(allowPositionInLeadingTrivia, tokenFullStart, scanner.TokenStart())
 				tokenEnd := scanner.TokenEnd()
 				flags := scanner.TokenFlags()
+				if tokenEnd > end {
+					break
+				}
 				if tokenStart <= position && (position < tokenEnd) {
 					if token == ast.KindIdentifier || !ast.IsTokenKind(token) {
 						if ast.IsJSDocKind(current.Kind) {
@@ -210,6 +270,7 @@ func getTokenAtPosition(
 		}
 		current = next
 		left = current.Pos()
+		nodeAfterLeft = nil
 		next = nil
 		// When navigating into AsExpression or SatisfiesExpression, allow visiting
 		// their reparsed children to reach identifiers from JSDoc type assertions.

internal/fourslash/tests/completionsJSDocTrivia_test.go

@@ -0,0 +1,35 @@
+package fourslash_test
+
+import (
+	"testing"
+
+	"github.com/microsoft/typescript-go/internal/fourslash"
+	. "github.com/microsoft/typescript-go/internal/fourslash/tests/util"
+	"github.com/microsoft/typescript-go/internal/testutil"
+)
+
+func TestCompletionsJSDocTrivia(t *testing.T) {
+	t.Parallel()
+	defer testutil.RecoverAndFail(t, "Panic on fourslash test")
+
+	const content = `// @noLib: true
+/**
+ * @type {{
+ * 'string-property': boolean;
+ */*$*/ identifierProperty: boolean;
+ * }}
+ */
+var someVariable;`
+
+	f, done := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	defer done()
+	f.GoToMarker(t, "$")
+	f.VerifyCompletions(t, nil, &fourslash.CompletionsExpectedList{
+		IsIncomplete: false,
+		ItemDefaults: &fourslash.CompletionsExpectedItemDefaults{
+			CommitCharacters: &[]string{".", ",", ";"},
+			EditRange:        Ignored,
+		},
+		Items: &fourslash.CompletionsExpectedItems{},
+	})
+}