[Sema] Introduce EnumElementExprPatternRequest

hamishknight · hamishknight · commit 98855963ac3b · 2023-03-07T15:16:36.000Z
This is needed to ensure we only ever synthesize
a single unique ExprPattern when solving an
EnumElementPattern that we failed to lookup a
member for.
diff --git a/include/swift/AST/Pattern.h b/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 {
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -2265,6 +2265,28 @@ class ExprPatternMatchRequest
   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 *),
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -231,6 +231,9 @@ SWIFT_REQUEST(TypeChecker, NamingPatternRequest,
 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)
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -9928,13 +9928,12 @@ static bool inferEnumMemberThroughTildeEqualsOperator(
   if (!pattern->hasUnresolvedOriginalExpr())
     return true;
 
-  auto &DC = cs.DC;
+  auto *DC = pattern->getDeclContext();
   auto &ctx = cs.getASTContext();
 
-  // Slots for expression and variable are going to be filled via
-  // synthesizing ~= operator application.
-  auto *EP = ExprPattern::createResolved(
-      ctx, pattern->getUnresolvedOriginalExpr(), DC);
+  // Retrieve a corresponding ExprPattern which we can solve with ~=.
+  auto *EP =
+      llvm::cantFail(ctx.evaluator(EnumElementExprPatternRequest{pattern}));
 
   // result of ~= operator is always a `Bool`.
   auto *matchCall = EP->getMatchExpr();
diff --git a/lib/Sema/DerivedConformanceCodable.cpp b/lib/Sema/DerivedConformanceCodable.cpp
@@ -976,7 +976,7 @@ createEnumSwitch(ASTContext &C, DeclContext *DC, Expr *expr, EnumDecl *enumDecl,
               DC->mapTypeIntoContext(
                   targetElt->getParentEnum()->getDeclaredInterfaceType()),
               C),
-          SourceLoc(), DeclNameLoc(), DeclNameRef(), targetElt, subpattern);
+          SourceLoc(), DeclNameLoc(), DeclNameRef(), targetElt, subpattern, DC);
       pat->setImplicit();
 
       auto labelItem = CaseLabelItem(pat);
diff --git a/lib/Sema/DerivedConformanceCodingKey.cpp b/lib/Sema/DerivedConformanceCodingKey.cpp
@@ -214,7 +214,8 @@ deriveBodyCodingKey_enum_stringValue(AbstractFunctionDecl *strValDecl, void *) {
     for (auto *elt : elements) {
       auto *baseTE = TypeExpr::createImplicit(enumType, C);
       auto *pat = new (C) EnumElementPattern(baseTE, SourceLoc(), DeclNameLoc(),
-                                             DeclNameRef(), elt, nullptr);
+                                             DeclNameRef(), elt, nullptr,
+                                             /*DC*/ strValDecl);
       pat->setImplicit();
 
       auto labelItem = CaseLabelItem(pat);
diff --git a/lib/Sema/DerivedConformanceComparable.cpp b/lib/Sema/DerivedConformanceComparable.cpp
@@ -116,19 +116,19 @@ deriveBodyComparable_enum_hasAssociatedValues_lt(AbstractFunctionDecl *ltDecl, v
     auto lhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(elt, 'l', ltDecl,
                                                       lhsPayloadVars);
     auto *lhsBaseTE = TypeExpr::createImplicit(enumType, C);
-    auto lhsElemPat =
-        new (C) EnumElementPattern(lhsBaseTE, SourceLoc(), DeclNameLoc(),
-                                   DeclNameRef(), elt, lhsSubpattern);
+    auto lhsElemPat = new (C)
+        EnumElementPattern(lhsBaseTE, SourceLoc(), DeclNameLoc(), DeclNameRef(),
+                           elt, lhsSubpattern, /*DC*/ ltDecl);
     lhsElemPat->setImplicit();
 
     // .<elt>(let r0, let r1, ...)
     SmallVector<VarDecl*, 4> rhsPayloadVars;
     auto rhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(elt, 'r', ltDecl,
                                                       rhsPayloadVars);
     auto *rhsBaseTE = TypeExpr::createImplicit(enumType, C);
-    auto rhsElemPat =
-        new (C) EnumElementPattern(rhsBaseTE, SourceLoc(), DeclNameLoc(),
-                                   DeclNameRef(), elt, rhsSubpattern);
+    auto rhsElemPat = new (C)
+        EnumElementPattern(rhsBaseTE, SourceLoc(), DeclNameLoc(), DeclNameRef(),
+                           elt, rhsSubpattern, /*DC*/ ltDecl);
     rhsElemPat->setImplicit();
 
     auto hasBoundDecls = !lhsPayloadVars.empty();
diff --git a/lib/Sema/DerivedConformanceEquatableHashable.cpp b/lib/Sema/DerivedConformanceEquatableHashable.cpp
@@ -181,19 +181,19 @@ deriveBodyEquatable_enum_hasAssociatedValues_eq(AbstractFunctionDecl *eqDecl,
     auto lhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(elt, 'l', eqDecl,
                                                       lhsPayloadVars);
     auto *lhsBaseTE = TypeExpr::createImplicit(enumType, C);
-    auto lhsElemPat =
-        new (C) EnumElementPattern(lhsBaseTE, SourceLoc(), DeclNameLoc(),
-                                   DeclNameRef(), elt, lhsSubpattern);
+    auto lhsElemPat = new (C)
+        EnumElementPattern(lhsBaseTE, SourceLoc(), DeclNameLoc(), DeclNameRef(),
+                           elt, lhsSubpattern, /*DC*/ eqDecl);
     lhsElemPat->setImplicit();
 
     // .<elt>(let r0, let r1, ...)
     SmallVector<VarDecl*, 3> rhsPayloadVars;
     auto rhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(elt, 'r', eqDecl,
                                                       rhsPayloadVars);
     auto *rhsBaseTE = TypeExpr::createImplicit(enumType, C);
-    auto rhsElemPat =
-        new (C) EnumElementPattern(rhsBaseTE, SourceLoc(), DeclNameLoc(),
-                                   DeclNameRef(), elt, rhsSubpattern);
+    auto rhsElemPat = new (C)
+        EnumElementPattern(rhsBaseTE, SourceLoc(), DeclNameLoc(), DeclNameRef(),
+                           elt, rhsSubpattern, /*DC*/ eqDecl);
     rhsElemPat->setImplicit();
 
     auto hasBoundDecls = !lhsPayloadVars.empty();
@@ -702,9 +702,10 @@ deriveBodyHashable_enum_hasAssociatedValues_hashInto(
 
     auto payloadPattern = DerivedConformance::enumElementPayloadSubpattern(elt, 'a', hashIntoDecl,
                                                        payloadVars);
-    auto pat = new (C) EnumElementPattern(
-        TypeExpr::createImplicit(enumType, C), SourceLoc(), DeclNameLoc(),
-        DeclNameRef(elt->getBaseIdentifier()), elt, payloadPattern);
+    auto pat = new (C)
+        EnumElementPattern(TypeExpr::createImplicit(enumType, C), SourceLoc(),
+                           DeclNameLoc(), DeclNameRef(elt->getBaseIdentifier()),
+                           elt, payloadPattern, /*DC*/ hashIntoDecl);
     pat->setImplicit();
 
     auto labelItem = CaseLabelItem(pat);
diff --git a/lib/Sema/DerivedConformanceRawRepresentable.cpp b/lib/Sema/DerivedConformanceRawRepresentable.cpp
@@ -114,7 +114,8 @@ deriveBodyRawRepresentable_raw(AbstractFunctionDecl *toRawDecl, void *) {
   for (auto elt : enumDecl->getAllElements()) {
     auto pat = new (C)
         EnumElementPattern(TypeExpr::createImplicit(enumType, C), SourceLoc(),
-                           DeclNameLoc(), DeclNameRef(), elt, nullptr);
+                           DeclNameLoc(), DeclNameRef(), elt, nullptr,
+                           /*DC*/ toRawDecl);
     pat->setImplicit();
 
     auto labelItem = CaseLabelItem(pat);
diff --git a/lib/Sema/DerivedConformances.cpp b/lib/Sema/DerivedConformances.cpp
@@ -737,7 +737,8 @@ DeclRefExpr *DerivedConformance::convertEnumToIndex(SmallVectorImpl<ASTNode> &st
     // generate: case .<Case>:
     auto pat = new (C)
         EnumElementPattern(TypeExpr::createImplicit(enumType, C), SourceLoc(),
-                           DeclNameLoc(), DeclNameRef(), elt, nullptr);
+                           DeclNameLoc(), DeclNameRef(), elt, nullptr,
+                           /*DC*/ funcDecl);
     pat->setImplicit();
     pat->setType(enumType);
 
diff --git a/lib/Sema/TypeCheckPattern.cpp b/lib/Sema/TypeCheckPattern.cpp
@@ -439,9 +439,8 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
     if (ume->getName().getBaseName().isSpecial())
       return nullptr;
 
-    return new (Context)
-        EnumElementPattern(ume->getDotLoc(), ume->getNameLoc(), ume->getName(),
-                           nullptr, ume);
+    return new (Context) EnumElementPattern(ume->getDotLoc(), ume->getNameLoc(),
+                                            ume->getName(), nullptr, ume, DC);
   }
   
   // Member syntax 'T.Element' forms a pattern if 'T' is an enum and the
@@ -483,7 +482,7 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
     base->setType(MetatypeType::get(ty));
     return new (Context)
         EnumElementPattern(base, ude->getDotLoc(), ude->getNameLoc(),
-                           ude->getName(), referencedElement, nullptr);
+                           ude->getName(), referencedElement, nullptr, DC);
   }
   
   // A DeclRef 'E' that refers to an enum element forms an EnumElementPattern.
@@ -496,8 +495,9 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
     auto enumTy = elt->getParentEnum()->getDeclaredTypeInContext();
     auto *base = TypeExpr::createImplicit(enumTy, Context);
 
-    return new (Context) EnumElementPattern(base, SourceLoc(), de->getNameLoc(),
-                                            elt->createNameRef(), elt, nullptr);
+    return new (Context)
+        EnumElementPattern(base, SourceLoc(), de->getNameLoc(),
+                           elt->createNameRef(), elt, nullptr, DC);
   }
   Pattern *visitUnresolvedDeclRefExpr(UnresolvedDeclRefExpr *ude) {
     // FIXME: This shouldn't be needed.  It is only necessary because of the
@@ -514,7 +514,7 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
 
       return new (Context)
           EnumElementPattern(base, SourceLoc(), ude->getNameLoc(),
-                             ude->getName(), referencedElement, nullptr);
+                             ude->getName(), referencedElement, nullptr, DC);
     }
       
     
@@ -614,7 +614,7 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
     auto *subPattern = composeTupleOrParenPattern(ce->getArgs());
     return new (Context) EnumElementPattern(
         baseTE, SourceLoc(), tailComponent->getNameLoc(),
-        tailComponent->getNameRef(), referencedElement, subPattern);
+        tailComponent->getNameRef(), referencedElement, subPattern, DC);
   }
 };
 
@@ -791,6 +791,15 @@ ExprPatternMatchRequest::evaluate(Evaluator &evaluator,
   return {matchVar, matchCall};
 }
 
+ExprPattern *
+EnumElementExprPatternRequest::evaluate(Evaluator &evaluator,
+                                        const EnumElementPattern *EEP) const {
+  assert(EEP->hasUnresolvedOriginalExpr());
+  auto *DC = EEP->getDeclContext();
+  return ExprPattern::createResolved(DC->getASTContext(),
+                                     EEP->getUnresolvedOriginalExpr(), DC);
+}
+
 Type PatternTypeRequest::evaluate(Evaluator &evaluator,
                                   ContextualPattern pattern) const {
   Pattern *P = pattern.getPattern();
@@ -1300,7 +1309,7 @@ Pattern *TypeChecker::coercePatternToType(ContextualPattern pattern,
         auto *BaseTE = TypeExpr::createImplicit(type, Context);
         P = new (Context) EnumElementPattern(
             BaseTE, NLE->getLoc(), DeclNameLoc(NLE->getLoc()),
-            NoneEnumElement->createNameRef(), NoneEnumElement, nullptr);
+            NoneEnumElement->createNameRef(), NoneEnumElement, nullptr, dc);
         return TypeChecker::coercePatternToType(
             pattern.forSubPattern(P, /*retainTopLevel=*/true), type, options);
       } else {
@@ -1355,7 +1364,7 @@ Pattern *TypeChecker::coercePatternToType(ContextualPattern pattern,
         auto *base = TypeExpr::createImplicit(extraOptTy, Context);
         sub = new (Context) EnumElementPattern(
             base, IP->getStartLoc(), DeclNameLoc(IP->getEndLoc()),
-            some->createNameRef(), nullptr, sub);
+            some->createNameRef(), nullptr, sub, dc);
         sub->setImplicit();
       }
 
