Fix contextual type for array completions (#2309)

Author: gabritto

internal/checker/services.go

@@ -671,16 +671,34 @@ func (c *Checker) GetContextualDeclarationsForObjectLiteralElement(objectLiteral
 	return result
 }
 
-// GetContextualTypeForArrayElement returns the contextual type for an element at the given index
+// GetContextualTypeForArrayLiteralAtPosition returns the contextual type for an element at the given position
 // in an array with the given contextual type.
-func (c *Checker) GetContextualTypeForArrayElement(contextualArrayType *Type, elementIndex int) *Type {
+func (c *Checker) GetContextualTypeForArrayLiteralAtPosition(contextualArrayType *Type, arrayLiteral *ast.Node, position int) *Type {
 	if contextualArrayType == nil {
 		return nil
 	}
-	// Pass -1 for length, firstSpreadIndex, and lastSpreadIndex since we don't have
-	// access to the actual array literal. This falls back to getting the iterated type
-	// or checking numeric properties, which is appropriate for completion contexts.
-	return c.getContextualTypeForElementExpression(contextualArrayType, elementIndex, -1, -1, -1)
+	firstSpreadIndex, lastSpreadIndex := -1, -1
+	elementIndex := 0
+	elements := arrayLiteral.Elements()
+	for i, elem := range elements {
+		if elem.Pos() < position {
+			elementIndex++
+		}
+		if ast.IsSpreadElement(elem) {
+			if firstSpreadIndex == -1 {
+				firstSpreadIndex = i
+			}
+			lastSpreadIndex = i
+		}
+	}
+	// The array may be incomplete, so we don't know its final length.
+	return c.getContextualTypeForElementExpression(
+		contextualArrayType,
+		elementIndex,
+		-1, /*length*/
+		firstSpreadIndex,
+		lastSpreadIndex,
+	)
 }
 
 var knownGenericTypeNames = map[string]struct{}{

internal/fourslash/tests/argumentCompletions_test.go

@@ -0,0 +1,41 @@
+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 TestArgumentCompletions(t *testing.T) {
+	t.Parallel()
+
+	defer testutil.RecoverAndFail(t, "Panic on fourslash test")
+	const content = `
+function foo(a: "a", b: "b") {}
+foo("a", /*1*/);
+
+
+const t3 = ['x', 'y', 'z'] as const;
+const x: [string, string, string, 'a' | 'b'] = [...t3, /*2*/];
+`
+	f, done := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	defer done()
+	f.VerifyCompletions(t, "1", &fourslash.CompletionsExpectedList{
+		ItemDefaults: &fourslash.CompletionsExpectedItemDefaults{
+			CommitCharacters: &DefaultCommitCharacters,
+		},
+		Items: &fourslash.CompletionsExpectedItems{
+			Includes: []fourslash.CompletionsExpectedItem{`"b"`},
+		},
+	})
+	f.VerifyCompletions(t, "2", &fourslash.CompletionsExpectedList{
+		ItemDefaults: &fourslash.CompletionsExpectedItemDefaults{
+			CommitCharacters: &DefaultCommitCharacters,
+		},
+		Items: &fourslash.CompletionsExpectedItems{
+			Includes: []fourslash.CompletionsExpectedItem{`"b"`},
+		},
+	})
+}

internal/ls/completions.go

@@ -2977,27 +2977,7 @@ func getContextualType(previousToken *ast.Node, position int, file *ast.SourceFi
 			contextualArrayType := typeChecker.GetContextualType(parent, checker.ContextFlagsNone)
 			if contextualArrayType != nil {
 				// Get the type for the first element (index 0)
-				return typeChecker.GetContextualTypeForArrayElement(contextualArrayType, 0)
-			}
-		}
-		return nil
-	case ast.KindCommaToken:
-		// When completing after `,` in an array literal (e.g., `[x, /*here*/]`),
-		// we should provide contextual type for the element after the comma
-		if ast.IsArrayLiteralExpression(parent) {
-			contextualArrayType := typeChecker.GetContextualType(parent, checker.ContextFlagsNone)
-			if contextualArrayType != nil {
-				// Count how many elements come before the cursor position
-				arrayLiteral := parent.AsArrayLiteralExpression()
-				elementIndex := 0
-				for _, elem := range arrayLiteral.Elements.Nodes {
-					if elem.Pos() < position {
-						elementIndex++
-					} else {
-						break
-					}
-				}
-				return typeChecker.GetContextualTypeForArrayElement(contextualArrayType, elementIndex)
+				return typeChecker.GetContextualTypeForArrayLiteralAtPosition(contextualArrayType, parent, position)
 			}
 		}
 		return nil
@@ -3023,22 +3003,30 @@ func getContextualType(previousToken *ast.Node, position int, file *ast.SourceFi
 		if ast.IsConditionalExpression(parent) {
 			return getContextualTypeForConditionalExpression(parent, position, file, typeChecker)
 		}
-		// Fall through to default for other colon contexts (object literals, etc.)
-		fallthrough
-	default:
-		argInfo := getArgumentInfoForCompletions(previousToken, position, file, typeChecker)
-		if argInfo != nil {
-			return typeChecker.GetContextualTypeForArgumentAtIndex(argInfo.invocation, argInfo.argumentIndex)
-		} else if isEqualityOperatorKind(previousToken.Kind) && ast.IsBinaryExpression(parent) && isEqualityOperatorKind(parent.AsBinaryExpression().OperatorToken.Kind) {
-			// completion at `x ===/**/`
-			return typeChecker.GetTypeAtLocation(parent.AsBinaryExpression().Left)
-		} else {
-			contextualType := typeChecker.GetContextualType(previousToken, checker.ContextFlagsCompletions)
-			if contextualType != nil {
-				return contextualType
+	case ast.KindCommaToken:
+		// When completing after `,` in an array literal (e.g., `[x, /*here*/]`),
+		// we should provide contextual type for the element after the comma.
+		if ast.IsArrayLiteralExpression(parent) {
+			contextualArrayType := typeChecker.GetContextualType(parent, checker.ContextFlagsNone)
+			if contextualArrayType != nil {
+				return typeChecker.GetContextualTypeForArrayLiteralAtPosition(contextualArrayType, parent, position)
 			}
-			return typeChecker.GetContextualType(previousToken, checker.ContextFlagsNone)
+			return nil
+		}
+	}
+	// Default case: see if we're in an argument position.
+	argInfo := getArgumentInfoForCompletions(previousToken, position, file, typeChecker)
+	if argInfo != nil {
+		return typeChecker.GetContextualTypeForArgumentAtIndex(argInfo.invocation, argInfo.argumentIndex)
+	} else if isEqualityOperatorKind(previousToken.Kind) && ast.IsBinaryExpression(parent) && isEqualityOperatorKind(parent.AsBinaryExpression().OperatorToken.Kind) {
+		// completion at `x ===/**/`
+		return typeChecker.GetTypeAtLocation(parent.AsBinaryExpression().Left)
+	} else {
+		contextualType := typeChecker.GetContextualType(previousToken, checker.ContextFlagsCompletions)
+		if contextualType != nil {
+			return contextualType
 		}
+		return typeChecker.GetContextualType(previousToken, checker.ContextFlagsNone)
 	}
 }