[Sema] Introduce a couple of ExprPattern requests

include/swift/AST/Pattern.h

@@ -508,25 +508,29 @@ class EnumElementPattern : public Pattern {
   DeclNameRef Name;
   PointerUnion<EnumElementDecl *, Expr*> ElementDeclOrUnresolvedOriginalExpr;
   Pattern /*nullable*/ *SubPattern;
+  DeclContext *DC;
 
 public:
   EnumElementPattern(TypeExpr *ParentType, SourceLoc DotLoc,
                      DeclNameLoc NameLoc, DeclNameRef Name,
-                     EnumElementDecl *Element, Pattern *SubPattern)
+                     EnumElementDecl *Element, Pattern *SubPattern,
+                     DeclContext *DC)
       : Pattern(PatternKind::EnumElement), ParentType(ParentType),
         DotLoc(DotLoc), NameLoc(NameLoc), Name(Name),
-        ElementDeclOrUnresolvedOriginalExpr(Element), SubPattern(SubPattern) {
+        ElementDeclOrUnresolvedOriginalExpr(Element), SubPattern(SubPattern),
+        DC(DC) {
     assert(ParentType && "Missing parent type?");
   }
 
   /// Create an unresolved EnumElementPattern for a `.foo` pattern relying on
   /// contextual type.
   EnumElementPattern(SourceLoc DotLoc, DeclNameLoc NameLoc, DeclNameRef Name,
-                     Pattern *SubPattern, Expr *UnresolvedOriginalExpr)
+                     Pattern *SubPattern, Expr *UnresolvedOriginalExpr,
+                     DeclContext *DC)
       : Pattern(PatternKind::EnumElement), ParentType(nullptr), DotLoc(DotLoc),
         NameLoc(NameLoc), Name(Name),
         ElementDeclOrUnresolvedOriginalExpr(UnresolvedOriginalExpr),
-        SubPattern(SubPattern) {}
+        SubPattern(SubPattern), DC(DC) {}
 
   bool hasSubPattern() const { return SubPattern; }
 
@@ -540,6 +544,8 @@ class EnumElementPattern : public Pattern {
 
   void setSubPattern(Pattern *p) { SubPattern = p; }
 
+  DeclContext *getDeclContext() const { return DC; }
+
   DeclNameRef getName() const { return Name; }
 
   EnumElementDecl *getElementDecl() const {
@@ -647,43 +653,59 @@ class OptionalSomePattern : public Pattern {
 class ExprPattern : public Pattern {
   llvm::PointerIntPair<Expr *, 1, bool> SubExprAndIsResolved;
 
-  /// An expression constructed during type-checking that produces a call to the
-  /// '~=' operator comparing the match expression on the left to the matched
-  /// value on the right.
-  Expr *MatchExpr = nullptr;
+  DeclContext *DC;
+
+  /// A synthesized call to the '~=' operator comparing the match expression
+  /// on the left to the matched value on the right.
+  mutable Expr *MatchExpr = nullptr;
+
+  /// An implicit variable used to represent the RHS value of the synthesized
+  /// match expression.
+  mutable VarDecl *MatchVar = nullptr;
 
-  /// An implicit variable used to represent the RHS value of the match.
-  VarDecl *MatchVar = nullptr;
+  ExprPattern(Expr *E, DeclContext *DC, bool isResolved)
+      : Pattern(PatternKind::Expr), SubExprAndIsResolved(E, isResolved),
+        DC(DC) {}
 
-  ExprPattern(Expr *E, bool isResolved)
-      : Pattern(PatternKind::Expr), SubExprAndIsResolved(E, isResolved) {}
+  friend class ExprPatternMatchRequest;
 
 public:
   /// Create a new parsed unresolved ExprPattern.
-  static ExprPattern *createParsed(ASTContext &ctx, Expr *E);
+  static ExprPattern *createParsed(ASTContext &ctx, Expr *E, DeclContext *DC);
 
   /// Create a new resolved ExprPattern. This should be used in cases
   /// where a user-written expression should be treated as an ExprPattern.
-  static ExprPattern *createResolved(ASTContext &ctx, Expr *E);
+  static ExprPattern *createResolved(ASTContext &ctx, Expr *E, DeclContext *DC);
 
   /// Create a new implicit resolved ExprPattern.
-  static ExprPattern *createImplicit(ASTContext &ctx, Expr *E);
+  static ExprPattern *createImplicit(ASTContext &ctx, Expr *E, DeclContext *DC);
 
   Expr *getSubExpr() const { return SubExprAndIsResolved.getPointer(); }
   void setSubExpr(Expr *e) { SubExprAndIsResolved.setPointer(e); }
 
-  Expr *getMatchExpr() const { return MatchExpr; }
+  DeclContext *getDeclContext() const { return DC; }
+
+  /// The match expression if it has been computed, \c nullptr otherwise.
+  /// Should only be used by the ASTDumper and ASTWalker.
+  Expr *getCachedMatchExpr() const { return MatchExpr; }
+
+  /// The match variable if it has been computed, \c nullptr otherwise.
+  /// Should only be used by the ASTDumper and ASTWalker.
+  VarDecl *getCachedMatchVar() const { return MatchVar; }
+
+  /// A synthesized call to the '~=' operator comparing the match expression
+  /// on the left to the matched value on the right.
+  Expr *getMatchExpr() const;
+
+  /// An implicit variable used to represent the RHS value of the synthesized
+  /// match expression.
+  VarDecl *getMatchVar() const;
+
   void setMatchExpr(Expr *e) {
-    assert(isResolved() && "cannot set match fn for unresolved expr patter");
+    assert(MatchExpr && "Should only update an existing MatchExpr");
     MatchExpr = e;
   }
 
-  VarDecl *getMatchVar() const { return MatchVar; }
-  void setMatchVar(VarDecl *v) {
-    assert(isResolved() && "cannot set match var for unresolved expr patter");
-    MatchVar = v;
-  }
-
   SourceLoc getLoc() const;
   SourceRange getSourceRange() const;
 
@@ -851,6 +873,9 @@ class ContextualPattern {
 };
 
 void simple_display(llvm::raw_ostream &out, const ContextualPattern &pattern);
+void simple_display(llvm::raw_ostream &out, const Pattern *pattern);
+
+SourceLoc extractNearestSourceLoc(const Pattern *pattern);
 
 } // end namespace swift
 

include/swift/AST/TypeCheckRequests.h

@@ -2219,6 +2219,74 @@ class NamingPatternRequest
   void cacheResult(NamedPattern *P) const;
 };
 
+class ExprPatternMatchResult {
+  VarDecl *MatchVar;
+  Expr *MatchExpr;
+
+  friend class ExprPattern;
+
+  // Should only be used as the default value for the request, as the caching
+  // logic assumes the request always produces a non-null result.
+  ExprPatternMatchResult(llvm::NoneType)
+      : MatchVar(nullptr), MatchExpr(nullptr) {}
+
+public:
+  ExprPatternMatchResult(VarDecl *matchVar, Expr *matchExpr)
+      : MatchVar(matchVar), MatchExpr(matchExpr) {
+    // Note the caching logic currently requires the request to produce a
+    // non-null result.
+    assert(matchVar && matchExpr);
+  }
+
+  VarDecl *getMatchVar() const { return MatchVar; }
+  Expr *getMatchExpr() const { return MatchExpr; }
+};
+
+/// Compute the match VarDecl and expression for an ExprPattern, which applies
+/// the \c ~= operator.
+class ExprPatternMatchRequest
+    : public SimpleRequest<ExprPatternMatchRequest,
+                           ExprPatternMatchResult(const ExprPattern *),
+                           RequestFlags::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  ExprPatternMatchResult evaluate(Evaluator &evaluator,
+                                  const ExprPattern *EP) const;
+
+public:
+  // Separate caching.
+  bool isCached() const { return true; }
+  Optional<ExprPatternMatchResult> getCachedResult() const;
+  void cacheResult(ExprPatternMatchResult result) const;
+};
+
+/// Creates a corresponding ExprPattern from the original Expr of an
+/// EnumElementPattern. This needs to be a cached request to ensure we don't
+/// generate multiple ExprPatterns along different constraint solver paths.
+class EnumElementExprPatternRequest
+    : public SimpleRequest<EnumElementExprPatternRequest,
+                           ExprPattern *(const EnumElementPattern *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  ExprPattern *evaluate(Evaluator &evaluator,
+                        const EnumElementPattern *EEP) const;
+
+public:
+  // Cached.
+  bool isCached() const { return true; }
+};
+
 class InterfaceTypeRequest :
     public SimpleRequest<InterfaceTypeRequest,
                          Type (ValueDecl *),

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -228,6 +228,12 @@ SWIFT_REQUEST(TypeChecker, ModuleImplicitImportsRequest,
               ImplicitImportList(ModuleDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, NamingPatternRequest,
               NamedPattern *(VarDecl *), SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ExprPatternMatchRequest,
+              ExprPatternMatchResult(const ExprPattern *),
+              SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, EnumElementExprPatternRequest,
+              ExprPattern *(const EnumElementPattern *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, OpaqueReadOwnershipRequest,
               OpaqueReadOwnership(AbstractStorageDecl *), SeparatelyCached,
               NoLocationInfo)

include/swift/Sema/ConstraintSystem.h

@@ -1629,6 +1629,14 @@ class Solution {
     return None;
   }
 
+  Optional<SyntacticElementTarget>
+  getTargetFor(SyntacticElementTargetKey key) const {
+    auto known = targets.find(key);
+    if (known == targets.end())
+      return None;
+    return known->second;
+  }
+
   ConstraintLocator *getCalleeLocator(ConstraintLocator *locator,
                                       bool lookThroughApply = true) const;
 

lib/AST/ASTDumper.cpp

@@ -472,7 +472,7 @@ namespace {
     void visitExprPattern(ExprPattern *P) {
       printCommon(P, "pattern_expr");
       OS << '\n';
-      if (auto m = P->getMatchExpr())
+      if (auto m = P->getCachedMatchExpr())
         printRec(m);
       else
         printRec(P->getSubExpr());

lib/AST/ASTWalker.cpp

@@ -1980,8 +1980,8 @@ Pattern *Traversal::visitEnumElementPattern(EnumElementPattern *P) {
 Pattern *Traversal::visitExprPattern(ExprPattern *P) {
   // If the pattern has been type-checked, walk the match expression, which
   // includes the explicit subexpression.
-  if (P->getMatchExpr()) {
-    if (Expr *newMatch = doIt(P->getMatchExpr())) {
+  if (auto *match = P->getCachedMatchExpr()) {
+    if (Expr *newMatch = doIt(match)) {
       P->setMatchExpr(newMatch);
       return P;
     }

lib/AST/Pattern.cpp

@@ -19,6 +19,7 @@
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/TypeLoc.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/Basic/Statistic.h"
@@ -333,9 +334,8 @@ EnumElementDecl *OptionalSomePattern::getElementDecl() const {
 /// Return true if this is a non-resolved ExprPattern which is syntactically
 /// irrefutable.
 static bool isIrrefutableExprPattern(const ExprPattern *EP) {
-  // If the pattern has a registered match expression, it's
-  // a type-checked ExprPattern.
-  if (EP->getMatchExpr()) return false;
+  // If the pattern is resolved, it must be irrefutable.
+  if (EP->isResolved()) return false;
 
   auto expr = EP->getSubExpr();
   while (true) {
@@ -501,20 +501,35 @@ void IsPattern::setCastType(Type type) {
 
 TypeRepr *IsPattern::getCastTypeRepr() const { return CastType->getTypeRepr(); }
 
-ExprPattern *ExprPattern::createParsed(ASTContext &ctx, Expr *E) {
-  return new (ctx) ExprPattern(E, /*isResolved*/ false);
+ExprPattern *ExprPattern::createParsed(ASTContext &ctx, Expr *E,
+                                       DeclContext *DC) {
+  return new (ctx) ExprPattern(E, DC, /*isResolved*/ false);
 }
 
-ExprPattern *ExprPattern::createResolved(ASTContext &ctx, Expr *E) {
-  return new (ctx) ExprPattern(E, /*isResolved*/ true);
+ExprPattern *ExprPattern::createResolved(ASTContext &ctx, Expr *E,
+                                         DeclContext *DC) {
+  return new (ctx) ExprPattern(E, DC, /*isResolved*/ true);
 }
 
-ExprPattern *ExprPattern::createImplicit(ASTContext &ctx, Expr *E) {
-  auto *EP = ExprPattern::createResolved(ctx, E);
+ExprPattern *ExprPattern::createImplicit(ASTContext &ctx, Expr *E,
+                                         DeclContext *DC) {
+  auto *EP = ExprPattern::createResolved(ctx, E, DC);
   EP->setImplicit();
   return EP;
 }
 
+Expr *ExprPattern::getMatchExpr() const {
+  auto &eval = DC->getASTContext().evaluator;
+  return evaluateOrDefault(eval, ExprPatternMatchRequest{this}, None)
+      .getMatchExpr();
+}
+
+VarDecl *ExprPattern::getMatchVar() const {
+  auto &eval = DC->getASTContext().evaluator;
+  return evaluateOrDefault(eval, ExprPatternMatchRequest{this}, None)
+      .getMatchVar();
+}
+
 SourceLoc EnumElementPattern::getStartLoc() const {
   return (ParentType && !ParentType->isImplicit())
              ? ParentType->getSourceRange().Start
@@ -616,5 +631,13 @@ bool ContextualPattern::allowsInference() const {
 
 void swift::simple_display(llvm::raw_ostream &out,
                            const ContextualPattern &pattern) {
-  out << "(pattern @ " << pattern.getPattern() << ")";
+  simple_display(out, pattern.getPattern());
+}
+
+void swift::simple_display(llvm::raw_ostream &out, const Pattern *pattern) {
+  out << "(pattern @ " << pattern << ")";
+}
+
+SourceLoc swift::extractNearestSourceLoc(const Pattern *pattern) {
+  return pattern->getLoc();
 }

lib/AST/TypeCheckRequests.cpp

@@ -1011,6 +1011,25 @@ void NamingPatternRequest::cacheResult(NamedPattern *value) const {
   VD->NamingPattern = value;
 }
 
+//----------------------------------------------------------------------------//
+// ExprPatternMatchRequest caching.
+//----------------------------------------------------------------------------//
+
+Optional<ExprPatternMatchResult>
+ExprPatternMatchRequest::getCachedResult() const {
+  auto *EP = std::get<0>(getStorage());
+  if (!EP->MatchVar)
+    return None;
+
+  return ExprPatternMatchResult(EP->MatchVar, EP->MatchExpr);
+}
+
+void ExprPatternMatchRequest::cacheResult(ExprPatternMatchResult result) const {
+  auto *EP = std::get<0>(getStorage());
+  EP->MatchVar = result.getMatchVar();
+  EP->MatchExpr = result.getMatchExpr();
+}
+
 //----------------------------------------------------------------------------//
 // InterfaceTypeRequest computation.
 //----------------------------------------------------------------------------//

lib/IDE/ArgumentCompletion.cpp

@@ -95,19 +95,18 @@ static bool isExpressionResultTypeUnconstrained(const Solution &S, Expr *E) {
       return true;
     }
   }
-  auto targetIt = S.targets.find(E);
-  if (targetIt == S.targets.end()) {
+  auto target = S.getTargetFor(E);
+  if (!target)
     return false;
-  }
-  auto target = targetIt->second;
-  assert(target.kind == SyntacticElementTarget::Kind::expression);
-  switch (target.getExprContextualTypePurpose()) {
+
+  assert(target->kind == SyntacticElementTarget::Kind::expression);
+  switch (target->getExprContextualTypePurpose()) {
   case CTP_Unused:
     // If we aren't using the contextual type, its unconstrained by definition.
     return true;
   case CTP_Initialization: {
     // let x = <expr> is unconstrained
-    auto contextualType = target.getExprContextualType();
+    auto contextualType = target->getExprContextualType();
     return !contextualType || contextualType->is<UnresolvedType>();
   }
   default: