[WIP] Move ExprPattern type-checking into the solver

include/swift/Sema/CompletionContextFinder.h

@@ -20,6 +20,10 @@
 
 namespace swift {
 
+namespace constraints {
+class SyntacticElementTarget;
+}
+
 class CompletionContextFinder : public ASTWalker {
   enum class ContextKind {
     FallbackExpression,
@@ -53,12 +57,9 @@ class CompletionContextFinder : public ASTWalker {
     return MacroWalking::Arguments;
   }
 
-  /// Finder for completion contexts within the provided initial expression.
-  CompletionContextFinder(ASTNode initialNode, DeclContext *DC)
-      : InitialExpr(initialNode.dyn_cast<Expr *>()), InitialDC(DC) {
-    assert(DC);
-    initialNode.walk(*this);
-  };
+  /// Finder for completion contexts within the provided SyntacticElementTarget.
+  CompletionContextFinder(constraints::SyntacticElementTarget target,
+                          DeclContext *DC);
 
   /// Finder for completion contexts within the outermost non-closure context of
   /// the code completion expression's direct context.

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -222,6 +222,9 @@ CUSTOM_LOCATOR_PATH_ELT(TernaryBranch)
 /// Performing a pattern patch.
 CUSTOM_LOCATOR_PATH_ELT(PatternMatch)
 
+/// A conjunction for the ExprPatterns in a pattern.
+SIMPLE_LOCATOR_PATH_ELT(ExprPatternConjunction)
+
 /// Points to a particular attribute associated with one of
 /// the arguments e.g. `inout` or its type e.g. `@escaping`.
 ///

include/swift/Sema/ConstraintSystem.h

@@ -706,6 +706,11 @@ T *getAsPattern(ASTNode node) {
   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 +1520,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;
 
@@ -1700,6 +1709,16 @@ class Solution {
         : nullptr;
   }
 
+  /// Retrieve the solved ExprPattern that corresponds to provided
+  /// sub-expression.
+  NullablePtr<ExprPattern> getExprPatternFor(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;
@@ -2163,6 +2182,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;
 
@@ -2754,6 +2777,9 @@ class ConstraintSystem {
     /// The length of \c caseLabelItems.
     unsigned numCaseLabelItems;
 
+    /// The length of \c exprPatterns.
+    unsigned numExprPatterns;
+
     /// The length of \c isolatedParams.
     unsigned numIsolatedParams;
 
@@ -3175,6 +3201,15 @@ class ConstraintSystem {
     caseLabelItems[item] = info;
   }
 
+  /// Record a given ExprPattern as the parent of its sub-expression.
+  void setExprPatternFor(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);
@@ -4299,6 +4334,11 @@ class ConstraintSystem {
   /// \returns \c true if constraint generation failed, \c false otherwise
   bool generateConstraints(SingleValueStmtExpr *E);
 
+  /// Generate constraints for an array of ExprPatterns, forming a conjunction
+  /// that solves each expression in turn.
+  void generateConstraints(ArrayRef<ExprPattern *> exprPatterns,
+                           ConstraintLocatorBuilder locator);
+
   /// Generate constraints for the given (unchecked) expression.
   ///
   /// \returns a possibly-sanitized expression, or null if an error occurred.

include/swift/Sema/SyntacticElementTarget.h

@@ -190,13 +190,6 @@ class SyntacticElementTarget {
                          ConstraintLocator *convertTypeLocator,
                          bool isDiscarded);
 
-  SyntacticElementTarget(Expr *expr, DeclContext *dc, ExprPattern *pattern,
-                         Type patternType)
-      : SyntacticElementTarget(expr, dc, CTP_ExprPattern, patternType,
-                               /*isDiscarded=*/false) {
-    setPattern(pattern);
-  }
-
   SyntacticElementTarget(ClosureExpr *closure, Type convertType) {
     kind = Kind::closure;
     this->closure.closure = closure;
@@ -296,11 +289,8 @@ class SyntacticElementTarget {
   forPropertyWrapperInitializer(VarDecl *wrappedVar, DeclContext *dc,
                                 Expr *initializer);
 
-  static SyntacticElementTarget forExprPattern(Expr *expr, DeclContext *dc,
-                                               ExprPattern *pattern,
-                                               Type patternTy) {
-    return {expr, dc, pattern, patternTy};
-  }
+  /// Form a target for the match expression of an ExprPattern.
+  static SyntacticElementTarget forExprPattern(ExprPattern *pattern);
 
   /// This is useful for code completion.
   ASTNode getAsASTNode() const {
@@ -781,11 +771,13 @@ class SyntacticElementTarget {
 
       // For an initialization, include the pattern in the range too.
       if (isForInitialization()) {
-        if (auto patternRange = getInitializationPattern()->getSourceRange()) {
-          if (range.isInvalid()) {
-            range = patternRange;
-          } else {
-            range.widen(patternRange);
+        if (auto *pattern = getInitializationPattern()) {
+          if (auto patternRange = pattern->getSourceRange()) {
+            if (range.isInvalid()) {
+              range = patternRange;
+            } else {
+              range.widen(patternRange);
+            }
           }
         }
       }

lib/AST/ASTContext.cpp

@@ -4178,7 +4178,6 @@ GenericFunctionType *GenericFunctionType::get(GenericSignature sig,
                                               Type result,
                                               Optional<ExtInfo> info) {
   assert(sig && "no generic signature for generic function type?!");
-  assert(!result->hasTypeVariable());
 
   llvm::FoldingSetNodeID id;
   GenericFunctionType::Profile(id, sig, params, result, info);

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.getExprPatternFor(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/BuilderTransform.cpp

@@ -2344,7 +2344,8 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
     SmallVector<Solution, 4> solutions;
     cs.solveForCodeCompletion(solutions);
 
-    CompletionContextFinder analyzer(func, func->getDeclContext());
+    SyntacticElementTarget funcTarget(func);
+    CompletionContextFinder analyzer(funcTarget, func->getDeclContext());
     if (analyzer.hasCompletion()) {
       filterSolutionsForCodeCompletion(solutions, analyzer);
       for (const auto &solution : solutions) {