Support recommended completion for union

Author: Andy Hanson

src/harness/fourslash.ts

@@ -908,8 +908,8 @@ namespace FourSlash {
         }
 
         private verifyCompletionEntry(actual: ts.CompletionEntry, expected: FourSlashInterface.ExpectedCompletionEntry) {
-            const { insertText, replacementSpan, hasAction, kind, text, documentation, sourceDisplay } = typeof expected === "string"
-                ? { insertText: undefined, replacementSpan: undefined, hasAction: undefined, kind: undefined, text: undefined, documentation: undefined, sourceDisplay: undefined }
+            const { insertText, replacementSpan, hasAction, isRecommended, kind, text, documentation, sourceDisplay } = typeof expected === "string"
+                ? { insertText: undefined, replacementSpan: undefined, hasAction: undefined, isRecommended: undefined, kind: undefined, text: undefined, documentation: undefined, sourceDisplay: undefined }
                 : expected;
 
             if (actual.insertText !== insertText) {
@@ -926,6 +926,7 @@ namespace FourSlash {
             if (kind !== undefined) assert.equal(actual.kind, kind);
 
             assert.equal(actual.hasAction, hasAction);
+            assert.equal(actual.isRecommended, isRecommended);
 
             if (text) {
                 const actualDetails = this.getCompletionEntryDetails(actual.name, actual.source);
@@ -4747,6 +4748,7 @@ namespace FourSlashInterface {
         readonly insertText?: string,
         readonly replacementSpan?: FourSlash.Range,
         readonly hasAction?: boolean, // If not specified, will assert that this is false.
+        readonly isRecommended?: boolean; // If not specified, will assert that this is false.
         readonly kind?: string, // If not specified, won't assert about this
         readonly text: string;
         readonly documentation: string;

src/services/completions.ts

@@ -681,12 +681,15 @@ namespace ts.Completions {
     }
 
     function getRecommendedCompletion(currentToken: Node, position: number, sourceFile: SourceFile, checker: TypeChecker): Symbol | undefined {
-        const ty = getContextualType(currentToken, position, sourceFile, 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;
+        const contextualType = getContextualType(currentToken, position, sourceFile, checker);
+        // For a union `A | B`, just recommend the first one arbitrarily.
+        return firstDefined(contextualType && (contextualType.isUnion() ? contextualType.types : [contextualType]), type => {
+            const symbol = type && type.symbol;
+            // Don't include make a recommended completion for an abstract class
+            return symbol && (symbol.flags & SymbolFlags.EnumMember || symbol.flags & SymbolFlags.Class && !isAbstractConstructorSymbol(symbol))
+                ? getFirstSymbolInChain(symbol, currentToken, checker)
+                : undefined;
+        });
     }
 
     function getContextualType(currentToken: Node, position: number, sourceFile: SourceFile, checker: TypeChecker): Type | undefined {

tests/cases/fourslash/completionsRecommended_union.ts

@@ -0,0 +1,22 @@
+/// <reference path="fourslash.ts" />
+
+// @strictNullChecks: true
+
+////const enum E { A = "A", B = "B" }
+////const enum E2 { X = "X", Y = "Y" }
+////const e: E | undefined = /*a*/
+////const e2: E | E2 = /*b*/
+
+verify.completions(
+    {
+        marker: "a",
+        includes: { name: "E", isRecommended: true },
+        isNewIdentifierLocation: true,
+    },
+    {
+        marker: "b",
+        // Arbitrarily chooses one to be recommended
+        includes: [{ name: "E", isRecommended: true, }, { name: "E2" }],
+        isNewIdentifierLocation: true,
+    },
+);

tests/cases/fourslash/fourslash.ts

@@ -537,6 +537,7 @@ declare namespace FourSlashInterface {
         readonly insertText?: string,
         readonly replacementSpan?: Range,
         readonly hasAction?: boolean,
+        readonly isRecommended?: boolean,
         readonly kind?: string,
 
         // details