[CS] Allow ExprPatterns to be type-checked in the solver

Previously we would wait until CSApply, which
would trigger their type-checking in
`coercePatternToType`. This caused a number of
bugs, and hampered solver-based completion, which
does not run CSApply. Instead, form a single element
conjunction, which will generate constraints for the
expression when solved. The reason we need a
conjunction is to preserve the current behavior
where the expression pattern is type-checked
independently from the rest of the pattern.

We can then modify `coercePatternToType` to accept
a closure, which allows the solver to take over
rewriting the ExprPatterns it has already solved.

This then sets the stage for the complete removal
of `coercePatternToType`, and doing all pattern
type-checking in the solver.

Author: hamishknight

include/swift/Sema/ConstraintLocator.h

@@ -320,6 +320,14 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     return false;
   }
 
+  /// Determine whether this locator points directly to a given pattern.
+  template <typename P>
+  bool directlyAtPattern() const {
+    if (auto *pattern = getAnchor().dyn_cast<Pattern *>())
+      return isa<P>(pattern) && getPath().empty();
+    return false;
+  }
+
   /// Check whether the first element in the path of this locator (if any)
   /// is a given \c LocatorPathElt subclass.
   template <class T>

include/swift/Sema/ConstraintSystem.h

@@ -699,13 +699,27 @@ T *getAsStmt(ASTNode node) {
   return nullptr;
 }
 
+template <typename T>
+bool isPattern(ASTNode node) {
+  if (node.isNull() || !node.is<Pattern *>())
+    return false;
+
+  auto *P = node.get<Pattern *>();
+  return isa<T>(P);
+}
+
 template <typename T = Pattern>
 T *getAsPattern(ASTNode node) {
   if (auto *P = node.dyn_cast<Pattern *>())
     return dyn_cast_or_null<T>(P);
   return nullptr;
 }
 
+template <typename T = Pattern>
+T *castToPattern(ASTNode node) {
+  return cast<T>(node.get<Pattern *>());
+}
+
 template <typename T = Stmt> T *castToStmt(ASTNode node) {
   return cast<T>(node.get<Stmt *>());
 }
@@ -1515,6 +1529,10 @@ class Solution {
   llvm::SmallMapVector<const CaseLabelItem *, CaseLabelItemInfo, 4>
       caseLabelItems;
 
+  /// A map of expressions to the ExprPatterns that they are being solved as
+  /// a part of.
+  llvm::SmallMapVector<Expr *, ExprPattern *, 2> exprPatterns;
+
   /// The set of parameters that have been inferred to be 'isolated'.
   llvm::SmallVector<ParamDecl *, 2> isolatedParams;
 
@@ -1697,6 +1715,20 @@ class Solution {
         : nullptr;
   }
 
+  CaseLabelItemInfo getCaseLabelItemInfo(CaseLabelItem *labelItem) {
+    auto result = caseLabelItems.find(labelItem);
+    assert(result != caseLabelItems.end());
+    return result->second;
+  }
+
+  NullablePtr<ExprPattern> getExprPattern(Expr *E) const {
+    auto result = exprPatterns.find(E);
+    if (result == exprPatterns.end())
+      return nullptr;
+
+    return result->second;
+  }
+
   /// This method implements functionality of `Expr::isTypeReference`
   /// with data provided by a given solution.
   bool isTypeReference(Expr *E) const;
@@ -2161,6 +2193,10 @@ class ConstraintSystem {
   llvm::SmallMapVector<const CaseLabelItem *, CaseLabelItemInfo, 4>
       caseLabelItems;
 
+  /// A map of expressions to the ExprPatterns that they are being solved as
+  /// a part of.
+  llvm::SmallMapVector<Expr *, ExprPattern *, 2> exprPatterns;
+
   /// The set of parameters that have been inferred to be 'isolated'.
   llvm::SmallSetVector<ParamDecl *, 2> isolatedParams;
 
@@ -2751,6 +2787,9 @@ class ConstraintSystem {
     /// The length of \c caseLabelItems.
     unsigned numCaseLabelItems;
 
+    /// The length of \c exprPatterns.
+    unsigned numExprPatterns;
+
     /// The length of \c isolatedParams.
     unsigned numIsolatedParams;
 
@@ -3164,6 +3203,14 @@ class ConstraintSystem {
     caseLabelItems[item] = info;
   }
 
+  void setExprPattern(Expr *E, ExprPattern *EP) {
+    assert(E);
+    assert(EP);
+    auto inserted = exprPatterns.insert({E, EP}).second;
+    assert(inserted && "Mapping already defined?");
+    (void)inserted;
+  }
+
   Optional<CaseLabelItemInfo> getCaseLabelItemInfo(
       const CaseLabelItem *item) const {
     auto known = caseLabelItems.find(item);
@@ -4287,6 +4334,9 @@ class ConstraintSystem {
   /// \returns \c true if constraint generation failed, \c false otherwise
   bool generateConstraints(SingleValueStmtExpr *E);
 
+  LLVM_NODISCARD
+  Type generateConstraints(ExprPattern *EP);
+
   /// Generate constraints for the given (unchecked) expression.
   ///
   /// \returns a possibly-sanitized expression, or null if an error occurred.
@@ -4311,6 +4361,10 @@ class ConstraintSystem {
   LLVM_NODISCARD
   bool generateConstraints(StmtCondition condition, DeclContext *dc);
 
+  LLVM_NODISCARD
+  bool generateConstraints(CaseLabelItem *caseLabelItem, DeclContext *dc,
+                           Type convertTy, ConstraintLocator *locator);
+
   /// Generate constraints for a case statement.
   ///
   /// \param subjectType The type of the "subject" expression in the enclosing

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -81,7 +81,13 @@ Type swift::ide::getTypeForCompletion(const constraints::Solution &S,
 /// \endcode
 /// If the code completion expression occurs in such an AST, return the
 /// declaration of the \c $match variable, otherwise return \c nullptr.
-static VarDecl *getMatchVarIfInPatternMatch(Expr *E, ConstraintSystem &CS) {
+static VarDecl *getMatchVarIfInPatternMatch(Expr *E, const Solution &S) {
+  if (auto EP = S.getExprPattern(E))
+    return EP.get()->getMatchVar();
+
+  // TODO: Once ExprPattern type-checking is fully moved into the solver,
+  // the below can be deleted.
+  auto &CS = S.getConstraintSystem();
   auto &Context = CS.getASTContext();
 
   auto *Binary = dyn_cast_or_null<BinaryExpr>(CS.getParentExpr(E));
@@ -109,20 +115,21 @@ static VarDecl *getMatchVarIfInPatternMatch(Expr *E, ConstraintSystem &CS) {
 }
 
 Type swift::ide::getPatternMatchType(const constraints::Solution &S, Expr *E) {
-  if (auto MatchVar = getMatchVarIfInPatternMatch(E, S.getConstraintSystem())) {
-    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 = getTypeForCompletion(S, MatchVar);
-    }
-    if (MatchVarType) {
-      return MatchVarType;
-    }
-  }
-  return nullptr;
+  auto MatchVar = getMatchVarIfInPatternMatch(E, S);
+  if (!MatchVar)
+    return nullptr;
+
+  if (S.hasType(MatchVar))
+    return S.getResolvedType(MatchVar);
+
+  // If the ExprPattern wasn't solved as part of the constraint system, it's
+  // not part of the solution.
+  // TODO: This can be removed once ExprPattern type-checking is fully part
+  // of the constraint system.
+  if (auto T = S.getConstraintSystem().getVarType(MatchVar))
+    return T;
+
+  return getTypeForCompletion(S, MatchVar);
 }
 
 void swift::ide::getSolutionSpecificVarTypes(

lib/Sema/CSApply.cpp

@@ -8554,6 +8554,9 @@ namespace {
       return Action::SkipChildren();
     }
 
+    NullablePtr<Pattern>
+    rewritePattern(Pattern *pattern, DeclContext *DC);
+
     /// Rewrite the target, producing a new target.
     Optional<SolutionApplicationTarget>
     rewriteTarget(SolutionApplicationTarget target);
@@ -8800,12 +8803,68 @@ static Expr *wrapAsyncLetInitializer(
   return resultInit;
 }
 
+static Pattern *rewriteExprPattern(SolutionApplicationTarget matchTarget,
+                                   Type patternTy,
+                                   RewriteTargetFn rewriteTarget) {
+  auto *EP = matchTarget.getExprPattern();
+
+  // See if we can simplify to another kind of pattern.
+  if (auto simplified = TypeChecker::trySimplifyExprPattern(EP, patternTy))
+    return simplified.get();
+
+  auto resultTarget = rewriteTarget(matchTarget);
+  if (!resultTarget)
+    return nullptr;
+
+  matchTarget = *resultTarget;
+
+  EP->setMatchExpr(matchTarget.getAsExpr());
+  EP->getMatchVar()->setInterfaceType(patternTy->mapTypeOutOfContext());
+  EP->setType(patternTy);
+  return EP;
+}
+
+static Optional<Pattern *>
+tryRewriteExprPattern(Pattern *P, Solution &solution, Type patternTy,
+                      RewriteTargetFn rewriteTarget) {
+  // We may either have an ExprPattern, or a pattern that was mapped to an
+  // ExprPattern, such as an EnumElementPattern.
+  auto &cs = solution.getConstraintSystem();
+  auto matchTarget = cs.getSolutionApplicationTarget(P);
+  if (!matchTarget)
+    return None;
+
+  return rewriteExprPattern(*matchTarget, patternTy, rewriteTarget);
+}
+
+NullablePtr<Pattern> ExprWalker::rewritePattern(Pattern *pattern,
+                                                DeclContext *DC) {
+  auto &solution = Rewriter.solution;
+
+  // Figure out the pattern type.
+  auto patternTy = solution.getResolvedType(pattern);
+  patternTy = patternTy->reconstituteSugar(/*recursive=*/false);
+
+  // Coerce the pattern to its appropriate type.
+  TypeResolutionOptions patternOptions(TypeResolverContext::InExpression);
+  patternOptions |= TypeResolutionFlags::OverrideType;
+
+  auto tryRewritePattern = [&](Pattern *EP, Type ty) {
+    return ::tryRewriteExprPattern(
+        EP, solution, ty, [&](auto target) { return rewriteTarget(target); });
+  };
+
+  auto contextualPattern = ContextualPattern::forRawPattern(pattern, DC);
+  return TypeChecker::coercePatternToType(contextualPattern, patternTy,
+                                          patternOptions, tryRewritePattern);
+}
+
 /// Apply the given solution to the initialization target.
 ///
 /// \returns the resulting initialization expression.
 static Optional<SolutionApplicationTarget> applySolutionToInitialization(
-    Solution &solution, SolutionApplicationTarget target,
-    Expr *initializer) {
+    Solution &solution, SolutionApplicationTarget target, Expr *initializer,
+    RewriteTargetFn rewriteTarget) {
   auto wrappedVar = target.getInitializationWrappedVar();
   Type initType;
   if (wrappedVar) {
@@ -8870,10 +8929,15 @@ static Optional<SolutionApplicationTarget> applySolutionToInitialization(
 
   finalPatternType = finalPatternType->reconstituteSugar(/*recursive =*/false);
 
+  auto tryRewritePattern = [&](Pattern *EP, Type ty) {
+    return ::tryRewriteExprPattern(
+        EP, solution, ty, [&](auto target) { return rewriteTarget(target); });
+  };
+
   // Apply the solution to the pattern as well.
   auto contextualPattern = target.getContextualPattern();
   if (auto coercedPattern = TypeChecker::coercePatternToType(
-          contextualPattern, finalPatternType, options)) {
+          contextualPattern, finalPatternType, options, tryRewritePattern)) {
     resultTarget.setPattern(coercedPattern);
   } else {
     return None;
@@ -9023,10 +9087,16 @@ static Optional<SolutionApplicationTarget> applySolutionToForEachStmt(
     TypeResolutionOptions options(TypeResolverContext::ForEachStmt);
     options |= TypeResolutionFlags::OverrideType;
 
+    auto tryRewritePattern = [&](Pattern *EP, Type ty) {
+      return ::tryRewriteExprPattern(
+          EP, solution, ty, [&](auto target) { return rewriteTarget(target); });
+    };
+
     // Apply the solution to the pattern as well.
     auto contextualPattern = target.getContextualPattern();
     auto coercedPattern = TypeChecker::coercePatternToType(
-        contextualPattern, forEachStmtInfo.initType, options);
+        contextualPattern, forEachStmtInfo.initType, options,
+        tryRewritePattern);
     if (!coercedPattern)
       return None;
 
@@ -9114,7 +9184,8 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     switch (target.getExprContextualTypePurpose()) {
     case CTP_Initialization: {
       auto initResultTarget = applySolutionToInitialization(
-          solution, target, rewrittenExpr);
+          solution, target, rewrittenExpr,
+          [&](auto target) { return rewriteTarget(target); });
       if (!initResultTarget)
         return None;
 
@@ -9205,47 +9276,11 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     ConstraintSystem &cs = solution.getConstraintSystem();
     auto info = *cs.getCaseLabelItemInfo(*caseLabelItem);
 
-    // Figure out the pattern type.
-    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.getResolvedType(EP);
-
-        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;
-    auto contextualPattern =
-        ContextualPattern::forRawPattern(info.pattern,
-                                         target.getDeclContext());
-    if (auto coercedPattern = TypeChecker::coercePatternToType(
-            contextualPattern, patternType, patternOptions)) {
-      (*caseLabelItem)->setPattern(coercedPattern, /*resolved=*/true);
-    } else {
+    auto pattern = rewritePattern(info.pattern, target.getDeclContext());
+    if (!pattern)
       return None;
-    }
+
+    (*caseLabelItem)->setPattern(pattern.get(), /*resolved=*/true);
 
     // If there is a guard expression, coerce that.
     if (auto *guardExpr = info.guardExpr) {
@@ -9314,8 +9349,13 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
       options |= TypeResolutionFlags::OverrideType;
     }
 
+    auto tryRewritePattern = [&](Pattern *EP, Type ty) {
+      return ::tryRewriteExprPattern(
+          EP, solution, ty, [&](auto target) { return rewriteTarget(target); });
+    };
+
     if (auto coercedPattern = TypeChecker::coercePatternToType(
-            contextualPattern, patternType, options)) {
+            contextualPattern, patternType, options, tryRewritePattern)) {
       auto resultTarget = target;
       resultTarget.setPattern(coercedPattern);
       return resultTarget;