Merge pull request #41570 from xedin/add-expr-pattern-handling-to-solver

[ConstraintSystem] Add support for expression patterns

Author: xedin

include/swift/Sema/ConstraintLocator.h

@@ -85,6 +85,9 @@ enum ContextualTypePurpose : uint8_t {
   CTP_ComposedPropertyWrapper, ///< Composed wrapper type expected to match
                                ///< former 'wrappedValue' type
 
+  CTP_ExprPattern,      ///< `~=` operator application associated with expression
+                        /// pattern.
+
   CTP_CannotFail,       ///< Conversion can never fail. abort() if it does.
 };
 

include/swift/Sema/ConstraintSystem.h

@@ -968,6 +968,7 @@ class SolutionApplicationTargetsKey {
     stmtCondElement,
     expr,
     stmt,
+    pattern,
     patternBindingEntry,
     varDecl,
   };
@@ -982,6 +983,8 @@ class SolutionApplicationTargetsKey {
 
     const Stmt *stmt;
 
+    const Pattern *pattern;
+
     struct PatternBindingEntry {
       const PatternBindingDecl *patternBinding;
       unsigned index;
@@ -1011,6 +1014,11 @@ class SolutionApplicationTargetsKey {
     storage.stmt = stmt;
   }
 
+  SolutionApplicationTargetsKey(const Pattern *pattern) {
+    kind = Kind::pattern;
+    storage.pattern = pattern;
+  }
+
   SolutionApplicationTargetsKey(
       const PatternBindingDecl *patternBinding, unsigned index) {
     kind = Kind::patternBindingEntry;
@@ -1042,6 +1050,9 @@ class SolutionApplicationTargetsKey {
     case Kind::stmt:
       return lhs.storage.stmt == rhs.storage.stmt;
 
+    case Kind::pattern:
+      return lhs.storage.pattern == rhs.storage.pattern;
+
     case Kind::patternBindingEntry:
       return (lhs.storage.patternBindingEntry.patternBinding
                 == rhs.storage.patternBindingEntry.patternBinding) &&
@@ -1083,6 +1094,11 @@ class SolutionApplicationTargetsKey {
           DenseMapInfo<unsigned>::getHashValue(static_cast<unsigned>(kind)),
           DenseMapInfo<void *>::getHashValue(storage.stmt));
 
+    case Kind::pattern:
+      return hash_combine(
+          DenseMapInfo<unsigned>::getHashValue(static_cast<unsigned>(kind)),
+          DenseMapInfo<void *>::getHashValue(storage.pattern));
+
     case Kind::patternBindingEntry:
       return hash_combine(
           DenseMapInfo<unsigned>::getHashValue(static_cast<unsigned>(kind)),
@@ -1701,6 +1717,13 @@ class SolutionApplicationTarget {
                             ContextualTypePurpose contextualPurpose,
                             TypeLoc convertType, bool isDiscarded);
 
+  SolutionApplicationTarget(Expr *expr, DeclContext *dc, ExprPattern *pattern,
+                            Type patternType)
+      : SolutionApplicationTarget(expr, dc, CTP_ExprPattern, patternType,
+                                  /*isDiscarded=*/false) {
+    setPattern(pattern);
+  }
+
   SolutionApplicationTarget(AnyFunctionRef fn)
       : SolutionApplicationTarget(fn, fn.getBody()) { }
 
@@ -1786,6 +1809,12 @@ class SolutionApplicationTarget {
   static SolutionApplicationTarget forPropertyWrapperInitializer(
       VarDecl *wrappedVar, DeclContext *dc, Expr *initializer);
 
+  static SolutionApplicationTarget forExprPattern(Expr *expr, DeclContext *dc,
+                                                  ExprPattern *pattern,
+                                                  Type patternTy) {
+    return {expr, dc, pattern, patternTy};
+  }
+
   Expr *getAsExpr() const {
     switch (kind) {
     case Kind::expression:
@@ -1888,6 +1917,12 @@ class SolutionApplicationTarget {
     return expression.pattern;
   }
 
+  ExprPattern *getExprPattern() const {
+    assert(kind == Kind::expression);
+    assert(expression.contextualPurpose == CTP_ExprPattern);
+    return cast<ExprPattern>(expression.pattern);
+  }
+
   /// For a pattern initialization target, retrieve the contextual pattern.
   ContextualPattern getContextualPattern() const;
 
@@ -2008,7 +2043,8 @@ class SolutionApplicationTarget {
     assert(kind == Kind::expression);
     assert(expression.contextualPurpose == CTP_Initialization ||
            expression.contextualPurpose == CTP_ForEachStmt ||
-           expression.contextualPurpose == CTP_ForEachSequence);
+           expression.contextualPurpose == CTP_ForEachSequence ||
+           expression.contextualPurpose == CTP_ExprPattern);
     expression.pattern = pattern;
   }
 
@@ -5107,6 +5143,15 @@ class ConstraintSystem {
                              = FreeTypeVariableBinding::Disallow);
 
 public:
+  /// Pre-check the target, validating any types that occur in it
+  /// and folding sequence expressions.
+  ///
+  /// \param replaceInvalidRefsWithErrors Indicates whether it's allowed
+  /// to replace any discovered invalid member references with `ErrorExpr`.
+  static bool preCheckTarget(SolutionApplicationTarget &target,
+                             bool replaceInvalidRefsWithErrors,
+                             bool leaveClosureBodiesUnchecked);
+
   /// Pre-check the expression, validating any types that occur in the
   /// expression and folding sequence expressions.
   ///

lib/Sema/CSApply.cpp

@@ -8936,6 +8936,7 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     case CTP_Unused:
     case CTP_CaseStmt:
     case CTP_ReturnStmt:
+    case CTP_ExprPattern:
     case swift::CTP_ReturnSingleExpr:
     case swift::CTP_YieldByValue:
     case swift::CTP_YieldByReference:

lib/Sema/CSDiagnostics.cpp

@@ -710,6 +710,7 @@ Optional<Diag<Type, Type>> GenericArgumentsMismatchFailure::getDiagnosticFor(
   case CTP_YieldByReference:
   case CTP_CalleeResult:
   case CTP_EnumCaseRawValue:
+  case CTP_ExprPattern:
     break;
   }
   return None;
@@ -2538,6 +2539,7 @@ getContextualNilDiagnostic(ContextualTypePurpose CTP) {
   case CTP_YieldByReference:
   case CTP_WrappedProperty:
   case CTP_ComposedPropertyWrapper:
+  case CTP_ExprPattern:
     return None;
 
   case CTP_EnumCaseRawValue:
@@ -3302,6 +3304,7 @@ ContextualFailure::getDiagnosticFor(ContextualTypePurpose context,
   case CTP_CannotFail:
   case CTP_YieldByReference:
   case CTP_CalleeResult:
+  case CTP_ExprPattern:
     break;
   }
   return None;

lib/Sema/CSSimplify.cpp

@@ -12996,6 +12996,7 @@ void ConstraintSystem::addContextualConversionConstraint(
   case CTP_ForEachStmt:
   case CTP_WrappedProperty:
   case CTP_ComposedPropertyWrapper:
+  case CTP_ExprPattern:
     break;
   }
 

lib/Sema/ConstraintSystem.cpp

@@ -5887,6 +5887,7 @@ bool SolutionApplicationTarget::contextualTypeIsOnlyAHint() const {
   case CTP_WrappedProperty:
   case CTP_ComposedPropertyWrapper:
   case CTP_CannotFail:
+  case CTP_ExprPattern:
     return false;
   }
   llvm_unreachable("invalid contextual type");

lib/Sema/PreCheckExpr.cpp

@@ -2110,6 +2110,40 @@ Expr *PreCheckExpression::simplifyTypeConstructionWithLiteralArg(Expr *E) {
              : nullptr;
 }
 
+bool ConstraintSystem::preCheckTarget(SolutionApplicationTarget &target,
+                                      bool replaceInvalidRefsWithErrors,
+                                      bool leaveClosureBodiesUnchecked) {
+  auto *DC = target.getDeclContext();
+
+  bool hadErrors = false;
+
+  if (auto *expr = target.getAsExpr()) {
+    hadErrors |= preCheckExpression(expr, DC, replaceInvalidRefsWithErrors,
+                                    leaveClosureBodiesUnchecked);
+    // Even if the pre-check fails, expression still has to be re-set.
+    target.setExpr(expr);
+  }
+
+  if (target.isForEachStmt()) {
+    auto &info = target.getForEachStmtInfo();
+
+    if (info.whereExpr)
+      hadErrors |= preCheckExpression(info.whereExpr, DC,
+                                      /*replaceInvalidRefsWithErrors=*/true,
+                                      /*leaveClosureBodiesUnchecked=*/false);
+
+    // Update sequence and where expressions to pre-checked versions.
+    if (!hadErrors) {
+      info.stmt->setSequence(target.getAsExpr());
+
+      if (info.whereExpr)
+        info.stmt->setWhere(info.whereExpr);
+    }
+  }
+
+  return hadErrors;
+}
+
 /// Pre-check the expression, validating any types that occur in the
 /// expression and folding sequence expressions.
 bool ConstraintSystem::preCheckExpression(Expr *&expr, DeclContext *dc,

lib/Sema/TypeCheckCodeCompletion.cpp

@@ -579,12 +579,10 @@ bool TypeChecker::typeCheckForCodeCompletion(
   if (needsPrecheck) {
     // First, pre-check the expression, validating any types that occur in the
     // expression and folding sequence expressions.
-    auto failedPreCheck = ConstraintSystem::preCheckExpression(
-        expr, DC,
-        /*replaceInvalidRefsWithErrors=*/true,
-        /*leaveClosureBodiesUnchecked=*/true);
-
-    target.setExpr(expr);
+    auto failedPreCheck =
+        ConstraintSystem::preCheckTarget(target,
+                                         /*replaceInvalidRefsWithErrors=*/true,
+                                         /*leaveClosureBodiesUnchecked=*/true);
 
     if (failedPreCheck)
       return false;

lib/Sema/TypeCheckConstraints.cpp

@@ -337,25 +337,23 @@ Type TypeChecker::typeCheckExpression(Expr *&expr, DeclContext *dc,
 }
 
 Optional<SolutionApplicationTarget>
-TypeChecker::typeCheckExpression(
-    SolutionApplicationTarget &target,
-    TypeCheckExprOptions options) {
-  Expr *expr = target.getAsExpr();
+TypeChecker::typeCheckExpression(SolutionApplicationTarget &target,
+                                 TypeCheckExprOptions options) {
   DeclContext *dc = target.getDeclContext();
   auto &Context = dc->getASTContext();
-  FrontendStatsTracer StatsTracer(Context.Stats,
-                                  "typecheck-expr", expr);
-  PrettyStackTraceExpr stackTrace(Context, "type-checking", expr);
+  FrontendStatsTracer StatsTracer(Context.Stats, "typecheck-expr",
+                                  target.getAsExpr());
+  PrettyStackTraceExpr stackTrace(Context, "type-checking", target.getAsExpr());
 
   // First, pre-check the expression, validating any types that occur in the
   // expression and folding sequence expressions.
-  if (ConstraintSystem::preCheckExpression(
-        expr, dc, /*replaceInvalidRefsWithErrors=*/true,
-        options.contains(TypeCheckExprFlags::LeaveClosureBodyUnchecked))) {
-    target.setExpr(expr);
+  if (ConstraintSystem::preCheckTarget(
+          target, /*replaceInvalidRefsWithErrors=*/true,
+          options.contains(TypeCheckExprFlags::LeaveClosureBodyUnchecked))) {
     return None;
   }
-  target.setExpr(expr);
+
+  auto *expr = target.getAsExpr();
 
   // Check whether given expression has a code completion token which requires
   // special handling.
@@ -810,25 +808,6 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
   if (!sequenceProto)
     return failed();
 
-  // Precheck the sequence.
-  Expr *sequence = stmt->getSequence();
-  if (ConstraintSystem::preCheckExpression(
-          sequence, dc, /*replaceInvalidRefsWithErrors=*/true,
-          /*leaveClosureBodiesUnchecked=*/false))
-    return failed();
-  stmt->setSequence(sequence);
-
-  // Precheck the filtering condition.
-  if (Expr *whereExpr = stmt->getWhere()) {
-    if (ConstraintSystem::preCheckExpression(
-            whereExpr, dc,
-            /*replaceInvalidRefsWithErrors=*/true,
-            /*leaveClosureBodiesUnchecked=*/false))
-      return failed();
-
-    stmt->setWhere(whereExpr);
-  }
-
   auto target = SolutionApplicationTarget::forForEachStmt(
       stmt, sequenceProto, dc, /*bindPatternVarsOneWay=*/false);
   if (!typeCheckExpression(target))
@@ -913,13 +892,22 @@ bool TypeChecker::typeCheckExprPattern(ExprPattern *EP, DeclContext *DC,
                                                 /*Implicit=*/true);
   Expr *matchCall = BinaryExpr::create(Context, EP->getSubExpr(), matchOp,
                                        matchVarRef, /*implicit*/ true);
+
+  // Result of `~=` should always be a boolean.
+  auto contextualTy = Context.getBoolDecl()->getDeclaredInterfaceType();
+  auto target = SolutionApplicationTarget::forExprPattern(matchCall, DC, EP,
+                                                          contextualTy);
+
   // Check the expression as a condition.
-  bool hadError = typeCheckCondition(matchCall, DC);
+  auto result = typeCheckExpression(target);
+  if (!result)
+    return true;
+
   // Save the type-checked expression in the pattern.
-  EP->setMatchExpr(matchCall);
+  EP->setMatchExpr(result->getAsExpr());
   // Set the type on the pattern.
   EP->setType(rhsType);
-  return hadError;
+  return false;
 }
 
 static Type replaceArchetypesWithTypeVariables(ConstraintSystem &cs,