[ConstraintSystem] Support solving expression patterns via injecting call to ~= operator

Augment the constraint solver to fallback to implicit `~=` application
when member couldn't be found for `EnumElement` patterns because
`case` statement should be able to match enum member directly, as well
as through an implicit `~=` operator application.

Author: xedin

lib/Sema/CSApply.cpp

@@ -9007,6 +9007,31 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     Type patternType = solution.simplifyType(solution.getType(info.pattern));
     patternType = patternType->reconstituteSugar(/*recursive=*/false);
 
+    // Check whether this enum element is resolved via ~= application.
+    if (auto *enumElement = dyn_cast<EnumElementPattern>(info.pattern)) {
+      if (auto target = cs.getSolutionApplicationTarget(enumElement)) {
+        auto *EP = target->getExprPattern();
+        auto enumType = solution.simplifyType(EP->getType());
+
+        auto *matchCall = target->getAsExpr();
+
+        auto *result = matchCall->walk(*this);
+        if (!result)
+          return None;
+
+        {
+          auto *matchVar = EP->getMatchVar();
+          matchVar->setInterfaceType(enumType->mapTypeOutOfContext());
+        }
+
+        EP->setMatchExpr(result);
+        EP->setType(enumType);
+
+        (*caseLabelItem)->setPattern(EP, /*resolved=*/true);
+        return target;
+      }
+    }
+
     // Coerce the pattern to its appropriate type.
     TypeResolutionOptions patternOptions(TypeResolverContext::InExpression);
     patternOptions |= TypeResolutionFlags::OverrideType;

lib/Sema/CSSimplify.cpp

@@ -8629,6 +8629,69 @@ fixMemberRef(ConstraintSystem &cs, Type baseTy,
   return nullptr;
 }
 
+/// Convert the given enum element pattern into an expression pattern
+/// and synthesize ~= operator application to find the type of the
+/// element.
+static bool inferEnumMemberThroughTildeEqualsOperator(
+    ConstraintSystem &cs, EnumElementPattern *pattern, Type enumTy,
+    Type elementTy, ConstraintLocator *locator) {
+  if (!pattern->hasUnresolvedOriginalExpr())
+    return true;
+
+  auto &DC = cs.DC;
+  auto &ctx = cs.getASTContext();
+
+  // Slots for expression and variable are going to be filled via
+  // synthesizing ~= operator application.
+  auto *EP = new (ctx) ExprPattern(pattern->getUnresolvedOriginalExpr(),
+                                   /*matchExpr=*/nullptr, /*matchVar=*/nullptr);
+
+  auto tildeEqualsApplication =
+      TypeChecker::synthesizeTildeEqualsOperatorApplication(EP, DC, enumTy);
+
+  if (!tildeEqualsApplication)
+    return true;
+
+  VarDecl *matchVar;
+  Expr *matchCall;
+
+  std::tie(matchVar, matchCall) = *tildeEqualsApplication;
+
+  // result of ~= operator is always a `Bool`.
+  auto target = SolutionApplicationTarget::forExprPattern(
+      matchCall, DC, EP, ctx.getBoolDecl()->getDeclaredInterfaceType());
+
+  DiagnosticTransaction diagnostics(ctx.Diags);
+  {
+    if (cs.preCheckTarget(target, /*replaceInvalidRefWithErrors=*/true,
+                          /*leaveClosureBodyUnchecked=*/false)) {
+      // Skip diagnostics if they are disabled, otherwise it would result in
+      // duplicate diagnostics, since this operation is going to be repeated
+      // in diagnostic mode.
+      if (!cs.shouldAttemptFixes())
+        diagnostics.abort();
+
+      return true;
+    }
+  }
+
+  cs.generateConstraints(target, FreeTypeVariableBinding::Disallow);
+
+  // Sub-expression associated with expression pattern is the enum element
+  // access which needs to be connected to the provided element type.
+  cs.addConstraint(ConstraintKind::Conversion, cs.getType(EP->getSubExpr()),
+                   elementTy, cs.getConstraintLocator(EP));
+
+  // Store the $match variable and binary expression for solution application.
+  EP->setMatchVar(matchVar);
+  EP->setMatchExpr(matchCall);
+  EP->setType(enumTy);
+
+  cs.setSolutionApplicationTarget(pattern, target);
+
+  return false;
+}
+
 ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
     ConstraintKind kind, Type baseTy, DeclNameRef member, Type memberTy,
     DeclContext *useDC, FunctionRefKind functionRefKind,
@@ -8834,6 +8897,30 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
     (void)recordFix(fix);
   }
 
+  // If there were no results from a direct enum lookup, let's attempt
+  // to resolve this member via ~= operator application.
+  if (candidates.empty()) {
+    if (auto patternLoc =
+            locator->getLastElementAs<LocatorPathElt::PatternMatch>()) {
+      if (auto *enumElement =
+              dyn_cast<EnumElementPattern>(patternLoc->getPattern())) {
+        auto enumType = baseObjTy->getMetatypeInstanceType();
+
+        // If the synthesis of ~= resulted in errors (i.e. broken stdlib)
+        // that would be diagnosed inline, so let's just fall through and
+        // let this situation be diagnosed as a missing member.
+        auto hadErrors = inferEnumMemberThroughTildeEqualsOperator(
+            *this, enumElement, enumType, memberTy, locator);
+
+        // Let's consider current member constraint solved because it's
+        // replaced by a new set of constraints that would resolve member
+        // type.
+        if (!hadErrors)
+          return SolutionKind::Solved;
+      }
+    }
+  }
+
   if (!candidates.empty()) {
     addOverloadSet(candidates, locator);
     return SolutionKind::Solved;

test/Constraints/result_builder.swift

@@ -882,3 +882,40 @@ func test_weak_optionality_stays_the_same() {
     }
   }
 }
+
+enum WrapperEnum<Wrapped> where Wrapped: RawRepresentable {
+case known(Wrapped)
+
+  static func ~= (lhs: Wrapped, rhs: WrapperEnum<Wrapped>) -> Bool where Wrapped: Equatable {
+    switch rhs {
+    case .known(let wrapped):
+      return wrapped == lhs
+    }
+  }
+}
+
+func test_custom_tilde_equals_operator_matching() {
+  @resultBuilder
+  struct Builder {
+    static func buildBlock<T>(_ t: T) -> T { t }
+    static func buildEither<T>(first: T) -> T { first }
+    static func buildEither<T>(second: T) -> T { second }
+  }
+
+  enum TildeTest : String {
+  case test = "test"
+  }
+
+  struct S {}
+
+  struct MyView {
+    var entry: WrapperEnum<TildeTest>
+
+    @Builder var body: S {
+      switch entry {
+      case .test: S() // Ok although `.test` comes from `TildeTest` instead of `WrapperEnum`
+      case .known(_): S() // Ok - `.known` comes directly from `WrapperEnum`
+      }
+    }
+  }
+}

test/Constraints/result_builder_diags.swift

@@ -442,9 +442,8 @@ func getSomeEnumOverloaded(_: Int) -> E2 { return .b(0, nil) }
 
 func testOverloadedSwitch() {
   tuplify(true) { c in
-    // FIXME: Bad source location.
-    switch getSomeEnumOverloaded(17) { // expected-error{{type 'E2' has no member 'a'; did you mean 'b'?}}
-    case .a:
+    switch getSomeEnumOverloaded(17) {
+    case .a: // expected-error{{type 'E2' has no member 'a'; did you mean 'b'?}}
       "a"
     default:
       "default"
@@ -775,13 +774,14 @@ func test_rdar65667992() {
     @Builder var body: S {
       switch entry { // expected-error {{type 'E' has no member 'unset'}}
       case .set(_, _): S()
-      case .unset(_): S()
+      case .unset(_): S() // expected-error {{'_' can only appear in a pattern or on the left side of an assignment}}
       default: S()
       }
     }
   }
 }
 
+
 func test_weak_with_nonoptional_type() {
   class X {
     func test() -> Int { 0 }

test/expr/closure/multi_statement.swift

@@ -178,3 +178,33 @@ func test_workaround_for_optional_void_result() {
     let _ = $0
   }
 }
+
+enum WrapperEnum<Wrapped> where Wrapped: RawRepresentable {
+case known(Wrapped)
+
+  static func ~= (lhs: Wrapped, rhs: WrapperEnum<Wrapped>) -> Bool where Wrapped: Equatable {
+    switch rhs {
+    case .known(let wrapped):
+      return wrapped == lhs
+    }
+  }
+}
+
+func test_custom_tilde_equals_operator_matching() {
+  enum TildeTest : String {
+  case test = "test"
+  case otherTest = ""
+  }
+
+  func test(_: (WrapperEnum<TildeTest>) -> Void) {}
+
+  test { v in
+    print(v)
+
+    switch v {
+    case .test: break // Ok although `.test` comes from `TildeTest` instead of `WrapperEnum`
+    case .otherTest: break // Ok although `.otherTest` comes from `TildeTest` instead of `WrapperEnum`
+    case .known(_): break // Ok - `.known` comes from `WrapperEnum`
+    }
+  }
+}