Support completions contextual types in more places (#20768)

* Support completions contextual types in more places

* Adjust formatting

Author: Andy

Diff for: src/compiler/checker.ts

@@ -13943,7 +13943,7 @@ namespace ts {
         //   the contextual type of an initializer expression is the type annotation of the containing declaration, if present.
         function getContextualTypeForInitializerExpression(node: Expression): Type {
             const declaration = <VariableLikeDeclaration>node.parent;
-            if (hasInitializer(declaration) && node === declaration.initializer || node.kind === SyntaxKind.EqualsToken) {
+            if (hasInitializer(declaration) && node === declaration.initializer) {
                 const typeNode = getEffectiveTypeAnnotationNode(declaration);
                 if (typeNode) {
                     return getTypeFromTypeNode(typeNode);
@@ -14075,12 +14075,6 @@ namespace ts {
                 case SyntaxKind.AmpersandAmpersandToken:
                 case SyntaxKind.CommaToken:
                     return node === right ? getContextualType(binaryExpression) : undefined;
-                case SyntaxKind.EqualsEqualsEqualsToken:
-                case SyntaxKind.EqualsEqualsToken:
-                case SyntaxKind.ExclamationEqualsEqualsToken:
-                case SyntaxKind.ExclamationEqualsToken:
-                    // For completions after `x === `
-                    return node === operatorToken ? getTypeOfExpression(binaryExpression.left) : undefined;
                 default:
                     return undefined;
             }
@@ -14296,12 +14290,8 @@ namespace ts {
                     return getContextualTypeForReturnExpression(node);
                 case SyntaxKind.YieldExpression:
                     return getContextualTypeForYieldOperand(<YieldExpression>parent);
+                    case SyntaxKind.CallExpression:
                 case SyntaxKind.NewExpression:
-                    if (node.kind === SyntaxKind.NewKeyword) { // for completions after `new `
-                        return getContextualType(parent as NewExpression);
-                    }
-                    // falls through
-                case SyntaxKind.CallExpression:
                     return getContextualTypeForArgument(<CallExpression | NewExpression>parent, node);
                 case SyntaxKind.TypeAssertionExpression:
                 case SyntaxKind.AsExpression:
@@ -14336,12 +14326,6 @@ namespace ts {
                 case SyntaxKind.JsxOpeningElement:
                 case SyntaxKind.JsxSelfClosingElement:
                     return getAttributesTypeFromJsxOpeningLikeElement(<JsxOpeningLikeElement>parent);
-                case SyntaxKind.CaseClause: {
-                    if (node.kind === SyntaxKind.CaseKeyword) { // for completions after `case `
-                        const switchStatement = (parent as CaseClause).parent.parent;
-                        return getTypeOfExpression(switchStatement.expression);
-                    }
-                }
             }
             return undefined;
         }

Diff for: src/harness/fourslash.ts

@@ -441,12 +441,11 @@ namespace FourSlash {
             this.goToPosition(marker.position);
         }
 
-        public goToEachMarker(action: () => void) {
-            const markers = this.getMarkers();
+        public goToEachMarker(markers: ReadonlyArray<Marker>, action: (marker: FourSlash.Marker, index: number) => void) {
             assert(markers.length);
-            for (const marker of markers) {
-                this.goToMarker(marker);
-                action();
+            for (let i = 0; i < markers.length; i++) {
+                this.goToMarker(markers[i]);
+                action(markers[i], i);
             }
         }
 
@@ -3764,8 +3763,11 @@ namespace FourSlashInterface {
             this.state.goToMarker(name);
         }
 
-        public eachMarker(action: () => void) {
-            this.state.goToEachMarker(action);
+        public eachMarker(markers: ReadonlyArray<string>, action: (marker: FourSlash.Marker, index: number) => void): void;
+        public eachMarker(action: (marker: FourSlash.Marker, index: number) => void): void;
+        public eachMarker(a: ReadonlyArray<string> | ((marker: FourSlash.Marker, index: number) => void), b?: (marker: FourSlash.Marker, index: number) => void): void  {
+            const markers = typeof a === "function" ? this.state.getMarkers() : a.map(m => this.state.getMarkerByName(m));
+            this.state.goToEachMarker(markers, typeof a === "function" ? a : b);
         }
 
         public rangeStart(range: FourSlash.Range) {

Diff for: src/services/completions.ts

@@ -238,7 +238,7 @@ namespace ts.Completions {
 
     function getStringLiteralCompletionEntries(sourceFile: SourceFile, position: number, typeChecker: TypeChecker, compilerOptions: CompilerOptions, host: LanguageServiceHost, log: Log): CompletionInfo | undefined {
         const node = findPrecedingToken(position, sourceFile);
-        if (!node || (node.kind !== SyntaxKind.StringLiteral && node.kind !== SyntaxKind.NoSubstitutionTemplateLiteral)) {
+        if (!node || !isStringLiteral(node) && !isNoSubstitutionTemplateLiteral(node)) {
             return undefined;
         }
 
@@ -277,21 +277,9 @@ namespace ts.Completions {
             //      import x = require("/*completion position*/");
             //      var y = require("/*completion position*/");
             //      export * from "/*completion position*/";
-            const entries = PathCompletions.getStringLiteralCompletionsFromModuleNames(<StringLiteral>node, compilerOptions, host, typeChecker);
+            const entries = PathCompletions.getStringLiteralCompletionsFromModuleNames(node, compilerOptions, host, typeChecker);
             return pathCompletionsInfo(entries);
         }
-        else if (isEqualityExpression(node.parent)) {
-            // Get completions from the type of the other operand
-            // i.e. switch (a) {
-            //         case '/*completion position*/'
-            //      }
-            return getStringLiteralCompletionEntriesFromType(typeChecker.getTypeAtLocation(node.parent.left === node ? node.parent.right : node.parent.left), typeChecker);
-        }
-        else if (isCaseOrDefaultClause(node.parent)) {
-            // Get completions from the type of the switch expression
-            // i.e. x === '/*completion position'
-            return getStringLiteralCompletionEntriesFromType(typeChecker.getTypeAtLocation((<SwitchStatement>node.parent.parent.parent).expression), typeChecker);
-        }
         else {
             const argumentInfo = SignatureHelp.getImmediatelyContainingArgumentInfo(node, position, sourceFile);
             if (argumentInfo) {
@@ -303,7 +291,7 @@ namespace ts.Completions {
 
             // Get completion for string literal from string literal type
             // i.e. var x: "hi" | "hello" = "/*completion position*/"
-            return getStringLiteralCompletionEntriesFromType(typeChecker.getContextualType(<LiteralExpression>node), typeChecker);
+            return getStringLiteralCompletionEntriesFromType(getContextualTypeFromParent(node, typeChecker), typeChecker);
         }
     }
 
@@ -602,15 +590,57 @@ namespace ts.Completions {
     }
     type Request = { kind: "JsDocTagName" } | { kind: "JsDocTag" } | { kind: "JsDocParameterName", tag: JSDocParameterTag };
 
-    function getRecommendedCompletion(currentToken: Node, checker: TypeChecker/*, symbolToOriginInfoMap: SymbolOriginInfoMap*/): Symbol | undefined {
-        const ty = checker.getContextualType(currentToken as Expression);
+    function getRecommendedCompletion(currentToken: Node, checker: TypeChecker): Symbol | undefined {
+        const ty = getContextualType(currentToken, checker);
         const symbol = ty && ty.symbol;
         // Don't include make a recommended completion for an abstract class
         return symbol && (symbol.flags & SymbolFlags.Enum || symbol.flags & SymbolFlags.Class && !isAbstractConstructorSymbol(symbol))
             ? getFirstSymbolInChain(symbol, currentToken, checker)
             : undefined;
     }
 
+    function getContextualType(currentToken: Node, checker: ts.TypeChecker): Type | undefined {
+        const { parent } = currentToken;
+        switch (currentToken.kind) {
+            case ts.SyntaxKind.Identifier:
+                return getContextualTypeFromParent(currentToken as ts.Identifier, checker);
+            case ts.SyntaxKind.EqualsToken:
+                return ts.isVariableDeclaration(parent) ? checker.getContextualType(parent.initializer) :
+                    ts.isBinaryExpression(parent) ? checker.getTypeAtLocation(parent.left) : undefined;
+            case ts.SyntaxKind.NewKeyword:
+                return checker.getContextualType(parent as ts.Expression);
+            case ts.SyntaxKind.CaseKeyword:
+                return getSwitchedType(cast(currentToken.parent, isCaseClause), checker);
+            default:
+                return isEqualityOperatorKind(currentToken.kind) && ts.isBinaryExpression(parent) && isEqualityOperatorKind(parent.operatorToken.kind)
+                    // completion at `x ===/**/` should be for the right side
+                    ? checker.getTypeAtLocation(parent.left)
+                    : checker.getContextualType(currentToken as ts.Expression);
+        }
+    }
+
+    function getContextualTypeFromParent(node: ts.Expression, checker: ts.TypeChecker): Type | undefined {
+        const { parent } = node;
+        switch (parent.kind) {
+            case ts.SyntaxKind.NewExpression:
+                return checker.getContextualType(parent as ts.NewExpression);
+            case ts.SyntaxKind.BinaryExpression: {
+                const { left, operatorToken, right } = parent as ts.BinaryExpression;
+                return isEqualityOperatorKind(operatorToken.kind)
+                    ? checker.getTypeAtLocation(node === right ? left : right)
+                    : checker.getContextualType(node);
+            }
+            case ts.SyntaxKind.CaseClause:
+                return (parent as ts.CaseClause).expression === node ? getSwitchedType(parent as ts.CaseClause, checker) : undefined;
+            default:
+                return checker.getContextualType(node);
+        }
+    }
+
+    function getSwitchedType(caseClause: ts.CaseClause, checker: ts.TypeChecker): ts.Type {
+        return checker.getTypeAtLocation(caseClause.parent.parent.expression);
+    }
+
     function getFirstSymbolInChain(symbol: Symbol, enclosingDeclaration: Node, checker: TypeChecker): Symbol | undefined {
         const chain = checker.getAccessibleSymbolChain(symbol, enclosingDeclaration, /*meaning*/ SymbolFlags.All, /*useOnlyExternalAliasing*/ false);
         if (chain) return first(chain);
@@ -851,7 +881,7 @@ namespace ts.Completions {
 
         log("getCompletionData: Semantic work: " + (timestamp() - semanticStart));
 
-        const recommendedCompletion = getRecommendedCompletion(previousToken, typeChecker);
+        const recommendedCompletion = previousToken && getRecommendedCompletion(previousToken, typeChecker);
         return { symbols, isGlobalCompletion, isMemberCompletion, allowStringLiteral, isNewIdentifierLocation, location, isRightOfDot: (isRightOfDot || isRightOfOpenTag), request, keywordFilters, symbolToOriginInfoMap, recommendedCompletion, previousToken };
 
         type JSDocTagWithTypeExpression = JSDocParameterTag | JSDocPropertyTag | JSDocReturnTag | JSDocTypeTag | JSDocTypedefTag;
@@ -2081,15 +2111,16 @@ namespace ts.Completions {
         return isConstructorParameterCompletionKeyword(stringToToken(text));
     }
 
-    function isEqualityExpression(node: Node): node is BinaryExpression {
-        return isBinaryExpression(node) && isEqualityOperatorKind(node.operatorToken.kind);
-    }
-
-    function isEqualityOperatorKind(kind: SyntaxKind) {
-        return kind === SyntaxKind.EqualsEqualsToken ||
-            kind === SyntaxKind.ExclamationEqualsToken ||
-            kind === SyntaxKind.EqualsEqualsEqualsToken ||
-            kind === SyntaxKind.ExclamationEqualsEqualsToken;
+    function isEqualityOperatorKind(kind: ts.SyntaxKind): kind is EqualityOperator {
+        switch (kind) {
+            case ts.SyntaxKind.EqualsEqualsEqualsToken:
+            case ts.SyntaxKind.EqualsEqualsToken:
+            case ts.SyntaxKind.ExclamationEqualsEqualsToken:
+            case ts.SyntaxKind.ExclamationEqualsToken:
+                return true;
+            default:
+                return false;
+        }
     }
 
     /** Get the corresponding JSDocTag node if the position is in a jsDoc comment */

Diff for: src/services/pathCompletions.ts

@@ -1,6 +1,6 @@
 /* @internal */
 namespace ts.Completions.PathCompletions {
-    export function getStringLiteralCompletionsFromModuleNames(node: StringLiteral, compilerOptions: CompilerOptions, host: LanguageServiceHost, typeChecker: TypeChecker): CompletionEntry[] {
+    export function getStringLiteralCompletionsFromModuleNames(node: LiteralExpression, compilerOptions: CompilerOptions, host: LanguageServiceHost, typeChecker: TypeChecker): CompletionEntry[] {
         const literalValue = normalizeSlashes(node.text);
 
         const scriptPath = node.getSourceFile().path;

Diff for: tests/cases/fourslash/completionsRecommended_equals.ts

@@ -1,8 +1,10 @@
 /// <reference path="fourslash.ts" />
 
-////enum E {}
-////declare const e: E;
-////e === /**/
+////enum Enu {}
+////declare const e: Enu;
+////e === /*a*/;
+////e === E/*b*/
 
-goTo.marker();
-verify.completionListContains("E", "enum E", "", "enum", undefined, undefined, { isRecommended: true });
+goTo.eachMarker(["a", "b"], () => {
+    verify.completionListContains("Enu", "enum Enu", "", "enum", undefined, undefined, { isRecommended: true });
+});

Diff for: tests/cases/fourslash/completionsRecommended_import.ts

@@ -3,25 +3,36 @@
 // @noLib: true
 
 // @Filename: /a.ts
-////export class C {}
-////export function f(c: C) {}
+////export class Cls {}
+////export function f(c: Cls) {}
 
 // @Filename: /b.ts
 ////import { f } from "./a";
-// Here we will recommend a new import of 'C'
-////f(new /*b*/);
+// Here we will recommend a new import of 'Cls'
+////f(new C/*b0*/);
+////f(new /*b1*/);
 
 // @Filename: /c.ts
-////import * as a from "./a";
-// Here we will recommend 'a' because it contains 'C'.
-////a.f(new /*c*/);
+////import * as alpha from "./a";
+// Here we will recommend 'alpha' because it contains 'Cls'.
+////alpha.f(new al/*c0*/);
+////alpha.f(new /*c1*/);
 
-goTo.marker("b");
-verify.completionListContains({ name: "C", source: "/a" }, "class C", "", "class", undefined, /*hasAction*/ true, {
-    includeExternalModuleExports: true,
-    isRecommended: true,
-    sourceDisplay: "./a",
+goTo.eachMarker(["b0", "b1"], (_, idx) => {
+    verify.completionListContains(
+        { name: "Cls", source: "/a" },
+        idx === 0 ? "constructor Cls(): Cls" : "class Cls",
+        "",
+        "class",
+        undefined,
+        /*hasAction*/ true, {
+        includeExternalModuleExports: true,
+        isRecommended: true,
+        sourceDisplay: "./a",
+    });
+});
+
+goTo.eachMarker(["c0", "c1"], (_, idx) => {
+    verify.completionListContains("alpha", "import alpha", "", "alias", undefined, undefined, { isRecommended: true })
 });
 
-goTo.marker("c");
-verify.completionListContains("a", "import a", "", "alias", undefined, undefined, { isRecommended: true });

Diff for: tests/cases/fourslash/completionsRecommended_local.ts

@@ -1,18 +1,37 @@
 /// <reference path="fourslash.ts" />
 
-////enum E {}
-////class C {}
-////abstract class A {}
-////const e: E = /*e*/
-////const c: C = new /*c*/
-////const a: A = new /*a*/
+////enum Enu {}
+////class Cls {}
+////abstract class Abs {}
+////const e: Enu = E/*e0*/;
+////const e: Enu = /*e1*/;
+////const c: Cls = new C/*c0*/;
+////const c: Cls = new /*c1*/;
+////const a: Abs = new A/*a0*/;
+////const a: Abs = new /*a1*/;
 
-goTo.marker("e");
-verify.completionListContains("E", "enum E", "", "enum", undefined, undefined, { isRecommended: true });
+// Also works on mutations
+////let enu: Enu;
+////enu = E/*let0*/;
+////enu = E/*let1*/;
 
-goTo.marker("c");
-verify.completionListContains("C", "class C", "", "class", undefined, undefined, { isRecommended: true });
+goTo.eachMarker(["e0"], () => {//, "e1", "let0", "let1"
+    verify.completionListContains("Enu", "enum Enu", "", "enum", undefined, undefined, { isRecommended: true });
+});
 
-goTo.marker("a");
-// Not recommended, because it's an abstract class
-verify.completionListContains("A", "class A", "", "class");
+goTo.eachMarker(["c0", "c1"], (_, idx) => {
+    verify.completionListContains(
+        "Cls",
+        idx === 0 ? "constructor Cls(): Cls" : "class Cls",
+        "",
+        "class",
+        undefined,
+        undefined, {
+        isRecommended: true,
+    });
+});
+
+goTo.eachMarker(["a0", "a1"], (_, idx) => {
+    // Not recommended, because it's an abstract class
+    verify.completionListContains("Abs", idx == 0 ? "constructor Abs(): Abs" : "class Abs", "", "class");
+});

Diff for: tests/cases/fourslash/completionsRecommended_namespace.ts

@@ -3,31 +3,37 @@
 // @noLib: true
 
 // @Filename: /a.ts
-////export namespace N {
+////export namespace Name {
 ////    export class C {}
 ////}
-////export function f(c: N.C) {}
-////f(new /*a*/);
+////export function f(c: Name.C) {}
+////f(new N/*a0*/);
+////f(new /*a1*/);
 
 // @Filename: /b.ts
 ////import { f } from "./a";
-// Here we will recommend a new import of 'N'
-////f(new /*b*/);
+// Here we will recommend a new import of 'Name'
+////f(new N/*b0*/);
+////f(new /*b1*/);
 
 // @Filename: /c.ts
-////import * as a from "./a";
-// Here we will recommend 'a' because it contains 'N' which contains 'C'.
-////a.f(new /*c*/);
+////import * as alpha from "./a";
+// Here we will recommend 'a' because it contains 'Name' which contains 'C'.
+////alpha.f(new a/*c0*/);
+////alpha.f(new /*c1*/);
 
-goTo.marker("a");
-verify.completionListContains("N", "namespace N", "", "module", undefined, undefined, { isRecommended: true });
+goTo.eachMarker(["a0", "a1"], () => {
+    verify.completionListContains("Name", "namespace Name", "", "module", undefined, undefined, { isRecommended: true });
+});
 
-goTo.marker("b");
-verify.completionListContains({ name: "N", source: "/a" }, "namespace N", "", "module", undefined, /*hasAction*/ true, {
-    includeExternalModuleExports: true,
-    isRecommended: true,
-    sourceDisplay: "./a",
+goTo.eachMarker(["b0", "b1"], () => {
+    verify.completionListContains({ name: "Name", source: "/a" }, "namespace Name", "", "module", undefined, /*hasAction*/ true, {
+        includeExternalModuleExports: true,
+        isRecommended: true,
+        sourceDisplay: "./a",
+    });
 });
 
-goTo.marker("c");
-verify.completionListContains("a", "import a", "", "alias", undefined, undefined, { isRecommended: true });
+goTo.eachMarker(["c0", "c1"], () => {
+    verify.completionListContains("alpha", "import alpha", "", "alias", undefined, undefined, { isRecommended: true });
+});