[CodeCompletion] Explicitly support enum pattern matching

Pattern matching in Swift can either be expression pattern matching by comparing two instances using the `~=` operator or using enum matching by matching the enum case and its associated types (+ tuple pattern matching, but that’s not relevant here). We currenlty only consider the expression pattern matching case for code completion. To provide enum pattern matching results, we thus need to have a `~=` operator between the code completion token and the match expression

For example, when we are completing

```swift
enum MyEnum {
  case myCase(String)
}

switch x {
case .#^COMPLETE^#
}
```
then we are looking up all overloads of `~=` and try to match it to the call arguments `(<Code Completion Type>, MyEnum)`.
The way we currently get `#^COMPLETE^#` to offer members of `MyEnum`, is that we are trying to make sure that the `~=<T: Equatable>(T, T)` operator defined in the standard library is the best solution even though it has fixes associated with it. For that we need to carefully make sure to ignore other, more favourable overloads of `~=` in `filterSolutions` so that `~=<T: Equatable>(T, T)` has the best score.

This poses several problems:
- If the user defines a custom overload of `~=` that we don't prune when filtering solutions (e.g. `func ~=(pattern: OtherType, value: MyEnum) -> Bool`), it gets a better score than `~=<T: Equatable>(T, T)` and thus we only offer members of `OtherType` instead of members from `MyEnum`
- We are also suggesting static members of `MyEnum`, even though we can't pattern match them due to the lack of the `~=` operator.

If we detect that the completion expression is in a pattern matching position, also suggests all enum members of the matched type. This allows us to remove the hack which deliberately ignores certain overloads of `~=` since we no longer rely on `~=<T: Equatable>(T, T)`. It thus provides correct results in both of the above cases.

Fixes rdar://77263334 [SR-14547]

Author: ahoppen

include/swift/AST/KnownIdentifiers.def

@@ -66,6 +66,7 @@ IDENTIFIER(decodeIfPresent)
 IDENTIFIER(Decoder)
 IDENTIFIER(decoder)
 IDENTIFIER_(Differentiation)
+IDENTIFIER_WITH_NAME(PatternMatchVar, "$match")
 IDENTIFIER(dynamicallyCall)
 IDENTIFIER(dynamicMember)
 IDENTIFIER(Element)

include/swift/Sema/CodeCompletionTypeChecking.h

@@ -89,21 +89,26 @@ namespace swift {
   /// formed during expression type-checking.
   class UnresolvedMemberTypeCheckCompletionCallback: public TypeCheckCompletionCallback {
   public:
-    struct Result {
+    struct ExprResult {
       Type ExpectedTy;
       bool IsImplicitSingleExpressionReturn;
     };
 
   private:
     CodeCompletionExpr *CompletionExpr;
-    SmallVector<Result, 4> Results;
+    SmallVector<ExprResult, 4> ExprResults;
+    SmallVector<Type, 1> EnumPatternTypes;
     bool GotCallback = false;
 
   public:
     UnresolvedMemberTypeCheckCompletionCallback(CodeCompletionExpr *CompletionExpr)
     : CompletionExpr(CompletionExpr) {}
 
-    ArrayRef<Result> getResults() const { return Results; }
+    ArrayRef<ExprResult> getExprResults() const { return ExprResults; }
+
+    /// If we are completing in a pattern matching position, the types of all
+    /// enums for whose cases are valid as an \c EnumElementPattern.
+    ArrayRef<Type> getEnumPatternTypes() const { return EnumPatternTypes; }
 
     /// True if at least one solution was passed via the \c sawSolution
     /// callback.

lib/IDE/CodeCompletion.cpp

@@ -3312,8 +3312,6 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
         getSemanticContext(EED, Reason, dynamicLookupInfo),
         expectedTypeContext);
     Builder.setAssociatedDecl(EED);
-    if (HasTypeContext)
-      Builder.addFlair(CodeCompletionFlairBit::ExpressionSpecific);
 
     addLeadingDot(Builder);
     addValueBaseName(Builder, EED->getBaseIdentifier());
@@ -4372,6 +4370,23 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     addObjCPoundKeywordCompletions(/*needPound=*/true);
   }
 
+  /// Returns \c true if \p VD is an initializer on the \c Optional or \c
+  /// Id_OptionalNilComparisonType type from the Swift stdlib.
+  static bool isInitializerOnOptional(Type T, ValueDecl *VD) {
+    bool IsOptionalType = false;
+    IsOptionalType |= static_cast<bool>(T->getOptionalObjectType());
+    if (auto *NTD = T->getAnyNominal()) {
+      IsOptionalType |= NTD->getBaseIdentifier() ==
+                        VD->getASTContext().Id_OptionalNilComparisonType;
+    }
+    if (IsOptionalType && VD->getModuleContext()->isStdlibModule() &&
+        isa<ConstructorDecl>(VD)) {
+      return true;
+    } else {
+      return false;
+    }
+  }
+
   void getUnresolvedMemberCompletions(Type T) {
     if (!T->mayHaveMembers())
       return;
@@ -4389,16 +4404,11 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     // We can only say .foo where foo is a static member of the contextual
     // type and has the same type (or if the member is a function, then the
     // same result type) as the contextual type.
-    FilteredDeclConsumer consumer(*this, [=](ValueDecl *VD,
-                                             DeclVisibilityKind Reason) {
-      if (T->getOptionalObjectType() &&
-          VD->getModuleContext()->isStdlibModule()) {
-        // In optional context, ignore '.init(<some>)', 'init(nilLiteral:)',
-        if (isa<ConstructorDecl>(VD))
-          return false;
-      }
-      return true;
-    });
+    FilteredDeclConsumer consumer(
+        *this, [=](ValueDecl *VD, DeclVisibilityKind Reason) {
+          // In optional context, ignore '.init(<some>)', 'init(nilLiteral:)',
+          return !isInitializerOnOptional(T, VD);
+        });
 
     auto baseType = MetatypeType::get(T);
     llvm::SaveAndRestore<LookupKind> SaveLook(Kind, LookupKind::ValueExpr);
@@ -4410,6 +4420,21 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
                              /*includeProtocolExtensionMembers*/true);
   }
 
+  /// Complete all enum members declared on \p T.
+  void getEnumElementPatternCompletions(Type T) {
+    if (!isa_and_nonnull<EnumDecl>(T->getAnyNominal()))
+      return;
+
+    auto baseType = MetatypeType::get(T);
+    llvm::SaveAndRestore<LookupKind> SaveLook(Kind, LookupKind::EnumElement);
+    llvm::SaveAndRestore<Type> SaveType(ExprType, baseType);
+    llvm::SaveAndRestore<bool> SaveUnresolved(IsUnresolvedMember, true);
+    lookupVisibleMemberDecls(*this, baseType, CurrDeclContext,
+                             /*includeInstanceMembers=*/false,
+                             /*includeDerivedRequirements=*/false,
+                             /*includeProtocolExtensionMembers=*/true);
+  }
+
   void getUnresolvedMemberCompletions(ArrayRef<Type> Types) {
     NeedLeadingDot = !HaveDot;
 
@@ -6461,8 +6486,8 @@ static void deliverCompletionResults(CodeCompletionContext &CompletionContext,
 }
 
 void deliverUnresolvedMemberResults(
-    ArrayRef<UnresolvedMemberTypeCheckCompletionCallback::Result> Results,
-    DeclContext *DC, SourceLoc DotLoc,
+    ArrayRef<UnresolvedMemberTypeCheckCompletionCallback::ExprResult> Results,
+    ArrayRef<Type> EnumPatternTypes, DeclContext *DC, SourceLoc DotLoc,
     ide::CodeCompletionContext &CompletionCtx,
     CodeCompletionConsumer &Consumer) {
   ASTContext &Ctx = DC->getASTContext();
@@ -6471,7 +6496,7 @@ void deliverUnresolvedMemberResults(
 
   assert(DotLoc.isValid());
   Lookup.setHaveDot(DotLoc);
-  Lookup.shouldCheckForDuplicates(Results.size() > 1);
+  Lookup.shouldCheckForDuplicates(Results.size() + EnumPatternTypes.size() > 1);
 
   // Get the canonical versions of the top-level types
   SmallPtrSet<CanType, 4> originalTypes;
@@ -6496,6 +6521,22 @@ void deliverUnresolvedMemberResults(
     Lookup.getUnresolvedMemberCompletions(Result.ExpectedTy);
   }
 
+  // Offer completions when interpreting the pattern match as an
+  // EnumElementPattern.
+  for (auto &Ty : EnumPatternTypes) {
+    Lookup.setExpectedTypes({Ty}, /*IsImplicitSingleExpressionReturn=*/false,
+                            /*expectsNonVoid=*/true);
+    Lookup.setIdealExpectedType(Ty);
+
+    // We can pattern match MyEnum against Optional<MyEnum>
+    if (Ty->getOptionalObjectType()) {
+      Type Unwrapped = Ty->lookThroughAllOptionalTypes();
+      Lookup.getEnumElementPatternCompletions(Unwrapped);
+    }
+
+    Lookup.getEnumElementPatternCompletions(Ty);
+  }
+
   deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);
 }
 
@@ -6608,8 +6649,9 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
       Lookup.fallbackTypeCheck(CurDeclContext);
 
     addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
-    deliverUnresolvedMemberResults(Lookup.getResults(), CurDeclContext, DotLoc,
-                                   CompletionContext, Consumer);
+    deliverUnresolvedMemberResults(Lookup.getExprResults(),
+                                   Lookup.getEnumPatternTypes(), CurDeclContext,
+                                   DotLoc, CompletionContext, Consumer);
     return true;
   }
   case CompletionKind::KeyPathExprSwift: {

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -811,59 +811,11 @@ class CompletionContextFinder : public ASTWalker {
 
 } // end namespace
 
-// Determine if the target expression is the implicit BinaryExpr generated for
-// pattern-matching in a switch/if/guard case (<completion> ~= matchValue).
-static bool isForPatternMatch(SolutionApplicationTarget &target) {
-  if (target.getExprContextualTypePurpose() != CTP_Condition)
-    return false;
-  Expr *condition = target.getAsExpr();
-  if (!condition->isImplicit())
-    return false;
-  if (auto *BE = dyn_cast<BinaryExpr>(condition)) {
-    Identifier id;
-    if (auto *ODRE = dyn_cast<OverloadedDeclRefExpr>(BE->getFn())) {
-      id = ODRE->getDecls().front()->getBaseIdentifier();
-    } else if (auto *DRE = dyn_cast<DeclRefExpr>(BE->getFn())) {
-      id = DRE->getDecl()->getBaseIdentifier();
-    }
-    if (id != target.getDeclContext()->getASTContext().Id_MatchOperator)
-      return false;
-    return isa<CodeCompletionExpr>(BE->getLHS());
-  }
-  return false;
-}
-
-/// Remove any solutions from the provided vector that both require fixes and have a
-/// score worse than the best.
+/// Remove any solutions from the provided vector that both require fixes and
+/// have a score worse than the best.
 static void filterSolutions(SolutionApplicationTarget &target,
                             SmallVectorImpl<Solution> &solutions,
                             CodeCompletionExpr *completionExpr) {
-  // FIXME: this is only needed because in pattern matching position, the
-  // code completion expression always becomes an expression pattern, which
-  // requires the ~= operator to be defined on the type being matched against.
-  // Pattern matching against an enum doesn't require that however, so valid
-  // solutions always end up having fixes. This is a problem because there will
-  // always be a valid solution as well. Optional defines ~= between Optional
-  // and _OptionalNilComparisonType (which defines a nilLiteral initializer),
-  // and the matched-against value can implicitly be made Optional if it isn't
-  // already, so _OptionalNilComparisonType is always a valid solution for the
-  // completion. That only generates the 'nil' completion, which is rarely what
-  // the user intends to write in this position and shouldn't be preferred over
-  // the other formed solutions (which require fixes). We should generate enum
-  // pattern completions separately, but for now ignore the
-  // _OptionalNilComparisonType solution.
-  if (isForPatternMatch(target) && completionExpr) {
-    solutions.erase(llvm::remove_if(solutions, [&](const Solution &S) {
-      ASTContext &ctx = S.getConstraintSystem().getASTContext();
-      if (!S.hasType(completionExpr))
-        return false;
-      if (auto ty = S.getResolvedType(completionExpr))
-        if (auto *NTD = ty->getAnyNominal())
-          return NTD->getBaseIdentifier() == ctx.Id_OptionalNilComparisonType;
-      return false;
-    }), solutions.end());
-  }
-
   if (solutions.size() <= 1)
     return;
 
@@ -1286,6 +1238,69 @@ sawSolution(const constraints::Solution &S) {
   }
 }
 
+/// If the code completion variable occurs in a pattern matching position, we
+/// have an AST that looks like this.
+/// \code
+/// (binary_expr implicit type='$T3'
+///   (overloaded_decl_ref_expr function_ref=compound decls=[
+///     Swift.(file).~=,
+///     Swift.(file).Optional extension.~=])
+///   (tuple_expr implicit type='($T1, (OtherEnum))'
+///     (code_completion_expr implicit type='$T1')
+///     (declref_expr implicit decl=swift_ide_test.(file).foo(x:).$match)))
+/// \endcode
+/// If the code completion expression occurs in such an AST, return the
+/// declaration of the \c $match variable, otherwise return \c nullptr.
+VarDecl *getMatchVarIfInPatternMatch(CodeCompletionExpr *CompletionExpr,
+                                     ConstraintSystem &CS) {
+  auto &Context = CS.getASTContext();
+
+  TupleExpr *ArgTuple =
+      dyn_cast_or_null<TupleExpr>(CS.getParentExpr(CompletionExpr));
+  if (!ArgTuple || !ArgTuple->isImplicit() || ArgTuple->getNumElements() != 2) {
+    return nullptr;
+  }
+
+  auto Binary = dyn_cast_or_null<BinaryExpr>(CS.getParentExpr(ArgTuple));
+  if (!Binary || !Binary->isImplicit()) {
+    return nullptr;
+  }
+
+  auto CalledOperator = Binary->getFn();
+  if (!CalledOperator || !CalledOperator->isImplicit()) {
+    return nullptr;
+  }
+  // The reference to the ~= operator might be an OverloadedDeclRefExpr or a
+  // DeclRefExpr, depending on how many ~= operators are viable.
+  if (auto Overloaded =
+          dyn_cast_or_null<OverloadedDeclRefExpr>(CalledOperator)) {
+    if (!llvm::all_of(Overloaded->getDecls(), [&Context](ValueDecl *D) {
+          return D->getBaseName() == Context.Id_MatchOperator;
+        })) {
+      return nullptr;
+    }
+  } else if (auto Ref = dyn_cast_or_null<DeclRefExpr>(CalledOperator)) {
+    if (Ref->getDecl()->getBaseName() != Context.Id_MatchOperator) {
+      return nullptr;
+    }
+  } else {
+    return nullptr;
+  }
+
+  auto MatchArg = dyn_cast_or_null<DeclRefExpr>(ArgTuple->getElement(1));
+  if (!MatchArg || !MatchArg->isImplicit()) {
+    return nullptr;
+  }
+
+  auto MatchVar = MatchArg->getDecl();
+  if (MatchVar && MatchVar->isImplicit() &&
+      MatchVar->getBaseName() == Context.Id_PatternMatchVar) {
+    return dyn_cast<VarDecl>(MatchVar);
+  } else {
+    return nullptr;
+  }
+}
+
 void UnresolvedMemberTypeCheckCompletionCallback::
 sawSolution(const constraints::Solution &S) {
   GotCallback = true;
@@ -1295,18 +1310,34 @@ sawSolution(const constraints::Solution &S) {
   // If the type couldn't be determined (e.g. because there isn't any context
   // to derive it from), let's not attempt to do a lookup since it wouldn't
   // produce any useful results anyway.
-  if (!ExpectedTy || ExpectedTy->is<UnresolvedType>())
-    return;
-
-  // If ExpectedTy is a duplicate of any other result, ignore this solution.
-  if (llvm::any_of(Results, [&](const Result &R) {
-    return R.ExpectedTy->isEqual(ExpectedTy);
-  })) {
-    return;
+  if (ExpectedTy && !ExpectedTy->is<UnresolvedType>()) {
+    // If ExpectedTy is a duplicate of any other result, ignore this solution.
+    if (!llvm::any_of(ExprResults, [&](const ExprResult &R) {
+          return R.ExpectedTy->isEqual(ExpectedTy);
+        })) {
+      bool SingleExprBody =
+          isImplicitSingleExpressionReturn(CS, CompletionExpr);
+      ExprResults.push_back({ExpectedTy, SingleExprBody});
+    }
   }
 
-  bool SingleExprBody = isImplicitSingleExpressionReturn(CS, CompletionExpr);
-  Results.push_back({ExpectedTy, SingleExprBody});
+  if (auto MatchVar = getMatchVarIfInPatternMatch(CompletionExpr, CS)) {
+    Type MatchVarType;
+    // If the MatchVar has an explicit type, it's not part of the solution. But
+    // we can look it up in the constraint system directly.
+    if (auto T = S.getConstraintSystem().getVarType(MatchVar)) {
+      MatchVarType = T;
+    } else {
+      MatchVarType = S.getResolvedType(MatchVar);
+    }
+    if (MatchVarType && !MatchVarType->is<UnresolvedType>()) {
+      if (!llvm::any_of(EnumPatternTypes, [&](const Type &R) {
+            return R->isEqual(MatchVarType);
+          })) {
+        EnumPatternTypes.push_back(MatchVarType);
+      }
+    }
+  }
 }
 
 void KeyPathTypeCheckCompletionCallback::sawSolution(

lib/Sema/TypeCheckConstraints.cpp

@@ -633,11 +633,9 @@ bool TypeChecker::typeCheckExprPattern(ExprPattern *EP, DeclContext *DC,
   PrettyStackTracePattern stackTrace(Context, "type-checking", EP);
 
   // Create a 'let' binding to stand in for the RHS value.
-  auto *matchVar = new (Context) VarDecl(/*IsStatic*/false,
-                                         VarDecl::Introducer::Let,
-                                         EP->getLoc(),
-                                         Context.getIdentifier("$match"),
-                                         DC);
+  auto *matchVar =
+      new (Context) VarDecl(/*IsStatic*/ false, VarDecl::Introducer::Let,
+                            EP->getLoc(), Context.Id_PatternMatchVar, DC);
   matchVar->setInterfaceType(rhsType->mapTypeOutOfContext());
 
   matchVar->setImplicit();