Merge pull request swiftlang#75010 from hamishknight/elemental

hamishknight · web-flow · commit 156aef5a914d · 2024-07-07T10:20:33.000+01:00
[AST] Simplify construction of EnumElementPattern
diff --git a/include/swift/AST/Pattern.h b/include/swift/AST/Pattern.h
@@ -548,27 +548,49 @@ class EnumElementPattern : public Pattern {
   Pattern /*nullable*/ *SubPattern;
   DeclContext *DC;
 
-public:
   EnumElementPattern(TypeExpr *ParentType, SourceLoc DotLoc,
                      DeclNameLoc NameLoc, DeclNameRef Name,
-                     EnumElementDecl *Element, Pattern *SubPattern,
-                     DeclContext *DC)
+                     PointerUnion<EnumElementDecl *, Expr *> ElementOrOriginal,
+                     Pattern *SubPattern, DeclContext *DC)
       : Pattern(PatternKind::EnumElement), ParentType(ParentType),
         DotLoc(DotLoc), NameLoc(NameLoc), Name(Name),
-        ElementDeclOrUnresolvedOriginalExpr(Element), SubPattern(SubPattern),
-        DC(DC) {
-    assert(ParentType && "Missing parent type?");
+        ElementDeclOrUnresolvedOriginalExpr(ElementOrOriginal),
+        SubPattern(SubPattern), DC(DC) {}
+
+public:
+  /// Create an EnumElementPattern with a parent expression, e.g `E.foo`.
+  static EnumElementPattern *create(TypeExpr *parentExpr, SourceLoc dotLoc,
+                                    DeclNameLoc nameLoc, DeclNameRef name,
+                                    EnumElementDecl *decl, Pattern *subPattern,
+                                    DeclContext *DC) {
+    auto &ctx = DC->getASTContext();
+    return new (ctx) EnumElementPattern(parentExpr, dotLoc, nameLoc, name, decl,
+                                        subPattern, DC);
   }
 
   /// Create an unresolved EnumElementPattern for a `.foo` pattern relying on
   /// contextual type.
-  EnumElementPattern(SourceLoc DotLoc, DeclNameLoc NameLoc, DeclNameRef Name,
-                     Pattern *SubPattern, Expr *UnresolvedOriginalExpr,
-                     DeclContext *DC)
-      : Pattern(PatternKind::EnumElement), ParentType(nullptr), DotLoc(DotLoc),
-        NameLoc(NameLoc), Name(Name),
-        ElementDeclOrUnresolvedOriginalExpr(UnresolvedOriginalExpr),
-        SubPattern(SubPattern), DC(DC) {}
+  static EnumElementPattern *create(SourceLoc dotLoc, DeclNameLoc nameLoc,
+                                    DeclNameRef name,
+                                    Expr *unresolvedOriginalExpr,
+                                    Pattern *subPattern, DeclContext *DC) {
+    auto &ctx = DC->getASTContext();
+    return new (ctx)
+        EnumElementPattern(/*parent*/ nullptr, dotLoc, nameLoc, name,
+                           unresolvedOriginalExpr, subPattern, DC);
+  }
+
+  static EnumElementPattern *
+  createImplicit(Type parentTy, SourceLoc dotLoc, DeclNameLoc nameLoc,
+                 EnumElementDecl *decl, Pattern *subPattern, DeclContext *DC);
+
+  static EnumElementPattern *createImplicit(Type parentTy,
+                                            EnumElementDecl *decl,
+                                            Pattern *subPattern,
+                                            DeclContext *DC) {
+    return createImplicit(parentTy, SourceLoc(), DeclNameLoc(), decl,
+                          subPattern, DC);
+  }
 
   bool hasSubPattern() const { return SubPattern; }
 
diff --git a/lib/AST/Pattern.cpp b/lib/AST/Pattern.cpp
@@ -620,7 +620,20 @@ void ExprPattern::updateMatchExpr(Expr *e) const {
 
   MatchExprAndOperandOwnership = {e, walker.Ownership};
 }
-  
+
+EnumElementPattern *
+EnumElementPattern::createImplicit(Type parentTy, SourceLoc dotLoc,
+                                   DeclNameLoc nameLoc, EnumElementDecl *decl,
+                                   Pattern *subPattern, DeclContext *DC) {
+  auto &ctx = DC->getASTContext();
+  auto *parentExpr = TypeExpr::createImplicit(parentTy, ctx);
+  auto *P = new (ctx) EnumElementPattern(
+      parentExpr, dotLoc, nameLoc, decl->createNameRef(), decl, subPattern, DC);
+  P->setImplicit();
+  P->setType(parentTy);
+  return P;
+}
+
 SourceLoc EnumElementPattern::getStartLoc() const {
   return (ParentType && !ParentType->isImplicit())
              ? ParentType->getSourceRange().Start
diff --git a/lib/Sema/DerivedConformanceCodable.cpp b/lib/Sema/DerivedConformanceCodable.cpp
@@ -969,13 +969,10 @@ createEnumSwitch(ASTContext &C, DeclContext *DC, Expr *expr, EnumDecl *enumDecl,
 
     if (caseBody) {
       // generate: case .<Case>:
-      auto pat = new (C) EnumElementPattern(
-          TypeExpr::createImplicit(
-              DC->mapTypeIntoContext(
-                  targetElt->getParentEnum()->getDeclaredInterfaceType()),
-              C),
-          SourceLoc(), DeclNameLoc(), DeclNameRef(), targetElt, subpattern, DC);
-      pat->setImplicit();
+      auto parentTy = DC->mapTypeIntoContext(
+          targetElt->getParentEnum()->getDeclaredInterfaceType());
+      auto *pat = EnumElementPattern::createImplicit(parentTy, targetElt,
+                                                     subpattern, DC);
 
       auto labelItem = CaseLabelItem(pat);
       auto stmt =
diff --git a/lib/Sema/DerivedConformanceCodingKey.cpp b/lib/Sema/DerivedConformanceCodingKey.cpp
@@ -214,12 +214,9 @@ deriveBodyCodingKey_enum_stringValue(AbstractFunctionDecl *strValDecl, void *) {
   } else {
     SmallVector<ASTNode, 4> cases;
     for (auto *elt : elements) {
-      auto *baseTE = TypeExpr::createImplicit(enumType, C);
-      auto *pat = new (C) EnumElementPattern(baseTE, SourceLoc(), DeclNameLoc(),
-                                             DeclNameRef(), elt, nullptr,
-                                             /*DC*/ strValDecl);
-      pat->setImplicit();
-
+      auto *pat = EnumElementPattern::createImplicit(enumType, elt,
+                                                     /*subPattern*/ nullptr,
+                                                     /*DC*/ strValDecl);
       auto labelItem = CaseLabelItem(pat);
 
       auto *caseValue = new (C) StringLiteralExpr(elt->getNameStr(),
diff --git a/lib/Sema/DerivedConformanceComparable.cpp b/lib/Sema/DerivedConformanceComparable.cpp
@@ -114,23 +114,17 @@ deriveBodyComparable_enum_hasAssociatedValues_lt(AbstractFunctionDecl *ltDecl, v
 
     // .<elt>(let l0, let l1, ...)
     SmallVector<VarDecl*, 4> lhsPayloadVars;
-    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, /*DC*/ ltDecl);
-    lhsElemPat->setImplicit();
+    auto *lhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(
+        elt, 'l', ltDecl, lhsPayloadVars);
+    auto *lhsElemPat = EnumElementPattern::createImplicit(
+        enumType, elt, lhsSubpattern, /*DC*/ ltDecl);
 
     // .<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, /*DC*/ ltDecl);
-    rhsElemPat->setImplicit();
+    auto *rhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(
+        elt, 'r', ltDecl, rhsPayloadVars);
+    auto *rhsElemPat = EnumElementPattern::createImplicit(
+        enumType, elt, rhsSubpattern, /*DC*/ ltDecl);
 
     auto hasBoundDecls = !lhsPayloadVars.empty();
     std::optional<MutableArrayRef<VarDecl *>> caseBodyVarDecls;
diff --git a/lib/Sema/DerivedConformanceEquatableHashable.cpp b/lib/Sema/DerivedConformanceEquatableHashable.cpp
@@ -181,23 +181,17 @@ deriveBodyEquatable_enum_hasAssociatedValues_eq(AbstractFunctionDecl *eqDecl,
 
     // .<elt>(let l0, let l1, ...)
     SmallVector<VarDecl*, 3> lhsPayloadVars;
-    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, /*DC*/ eqDecl);
-    lhsElemPat->setImplicit();
+    auto *lhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(
+        elt, 'l', eqDecl, lhsPayloadVars);
+    auto *lhsElemPat = EnumElementPattern::createImplicit(
+        enumType, elt, lhsSubpattern, /*DC*/ eqDecl);
 
     // .<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, /*DC*/ eqDecl);
-    rhsElemPat->setImplicit();
+    auto *rhsSubpattern = DerivedConformance::enumElementPayloadSubpattern(
+        elt, 'r', eqDecl, rhsPayloadVars);
+    auto *rhsElemPat = EnumElementPattern::createImplicit(
+        enumType, elt, rhsSubpattern, /*DC*/ eqDecl);
 
     auto hasBoundDecls = !lhsPayloadVars.empty();
     std::optional<MutableArrayRef<VarDecl *>> caseBodyVarDecls;
@@ -709,11 +703,8 @@ 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, /*DC*/ hashIntoDecl);
-    pat->setImplicit();
+    auto *pat = EnumElementPattern::createImplicit(
+        enumType, elt, payloadPattern, /*DC*/ hashIntoDecl);
 
     auto labelItem = CaseLabelItem(pat);
 
diff --git a/lib/Sema/DerivedConformanceRawRepresentable.cpp b/lib/Sema/DerivedConformanceRawRepresentable.cpp
@@ -113,11 +113,8 @@ deriveBodyRawRepresentable_raw(AbstractFunctionDecl *toRawDecl, void *) {
 
   SmallVector<ASTNode, 4> cases;
   for (auto elt : enumDecl->getAllElements()) {
-    auto pat = new (C)
-        EnumElementPattern(TypeExpr::createImplicit(enumType, C), SourceLoc(),
-                           DeclNameLoc(), DeclNameRef(), elt, nullptr,
-                           /*DC*/ toRawDecl);
-    pat->setImplicit();
+    auto *pat = EnumElementPattern::createImplicit(
+        enumType, elt, /*subPattern*/ nullptr, /*DC*/ toRawDecl);
 
     auto labelItem = CaseLabelItem(pat);
 
diff --git a/lib/Sema/DerivedConformances.cpp b/lib/Sema/DerivedConformances.cpp
@@ -785,12 +785,8 @@ DeclRefExpr *DerivedConformance::convertEnumToIndex(SmallVectorImpl<ASTNode> &st
     }
 
     // generate: case .<Case>:
-    auto pat = new (C)
-        EnumElementPattern(TypeExpr::createImplicit(enumType, C), SourceLoc(),
-                           DeclNameLoc(), DeclNameRef(), elt, nullptr,
-                           /*DC*/ funcDecl);
-    pat->setImplicit();
-    pat->setType(enumType);
+    auto *pat = EnumElementPattern::createImplicit(
+        enumType, elt, /*subPattern*/ nullptr, /*DC*/ funcDecl);
 
     auto labelItem = CaseLabelItem(pat);
 
@@ -947,12 +943,8 @@ CaseStmt *DerivedConformance::unavailableEnumElementCaseStmt(
 
   auto createElementPattern = [&]() -> EnumElementPattern * {
     // .<elt>
-    EnumElementPattern *eltPattern = new (C) EnumElementPattern(
-        TypeExpr::createImplicit(enumType, C), SourceLoc(), DeclNameLoc(),
-        DeclNameRef(elt->getBaseIdentifier()), elt, nullptr, /*DC*/ parentDC);
-    eltPattern->setImplicit();
-    eltPattern->setType(enumType);
-    return eltPattern;
+    return EnumElementPattern::createImplicit(
+        enumType, elt, /*subPattern*/ nullptr, /*DC*/ parentDC);
   };
 
   Pattern *labelItemPattern;
diff --git a/lib/Sema/TypeCheckPattern.cpp b/lib/Sema/TypeCheckPattern.cpp
@@ -450,8 +450,9 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
     if (ume->getName().getBaseName().isSpecial())
       return nullptr;
 
-    return new (Context) EnumElementPattern(ume->getDotLoc(), ume->getNameLoc(),
-                                            ume->getName(), nullptr, ume, DC);
+    return EnumElementPattern::create(ume->getDotLoc(), ume->getNameLoc(),
+                                      ume->getName(), ume,
+                                      /*subPattern*/ nullptr, DC);
   }
   
   // Member syntax 'T.Element' forms a pattern if 'T' is an enum and the
@@ -491,9 +492,9 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
     auto *base =
         TypeExpr::createForMemberDecl(repr, ude->getNameLoc(), enumDecl);
     base->setType(MetatypeType::get(ty));
-    return new (Context)
-        EnumElementPattern(base, ude->getDotLoc(), ude->getNameLoc(),
-                           ude->getName(), referencedElement, nullptr, DC);
+    return EnumElementPattern::create(base, ude->getDotLoc(), ude->getNameLoc(),
+                                      ude->getName(), referencedElement,
+                                      /*subPattern*/ nullptr, DC);
   }
   
   // A DeclRef 'E' that refers to an enum element forms an EnumElementPattern.
@@ -506,9 +507,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, DC);
+    return EnumElementPattern::create(base, SourceLoc(), de->getNameLoc(),
+                                      elt->createNameRef(), elt,
+                                      /*subPattern*/ nullptr, DC);
   }
   Pattern *visitUnresolvedDeclRefExpr(UnresolvedDeclRefExpr *ude) {
     // FIXME: This shouldn't be needed.  It is only necessary because of the
@@ -523,9 +524,9 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
       auto enumTy = enumDecl->getDeclaredTypeInContext();
       auto *base = TypeExpr::createImplicit(enumTy, Context);
 
-      return new (Context)
-          EnumElementPattern(base, SourceLoc(), ude->getNameLoc(),
-                             ude->getName(), referencedElement, nullptr, DC);
+      return EnumElementPattern::create(base, SourceLoc(), ude->getNameLoc(),
+                                        ude->getName(), referencedElement,
+                                        /*subPattern*/ nullptr, DC);
     }
       
     
@@ -617,9 +618,9 @@ class ResolvePattern : public ASTVisitor<ResolvePattern,
     assert(!repr->hasGenericArgList() && "should be handled above");
 
     auto *subPattern = composeTupleOrParenPattern(ce->getArgs());
-    return new (Context) EnumElementPattern(
-        baseTE, SourceLoc(), repr->getNameLoc(), repr->getNameRef(),
-        referencedElement, subPattern, DC);
+    return EnumElementPattern::create(baseTE, SourceLoc(), repr->getNameLoc(),
+                                      repr->getNameRef(), referencedElement,
+                                      subPattern, DC);
   }
 };
 
@@ -1061,13 +1062,9 @@ NullablePtr<Pattern> TypeChecker::trySimplifyExprPattern(ExprPattern *EP,
   if (auto *NLE = dyn_cast<NilLiteralExpr>(EP->getSubExpr())) {
     if (patternTy->getOptionalObjectType()) {
       auto *NoneEnumElement = ctx.getOptionalNoneDecl();
-      auto *BaseTE = TypeExpr::createImplicit(patternTy, ctx);
-      auto *EEP = new (ctx)
-          EnumElementPattern(BaseTE, NLE->getLoc(), DeclNameLoc(NLE->getLoc()),
-                             NoneEnumElement->createNameRef(), NoneEnumElement,
-                             nullptr, EP->getDeclContext());
-      EEP->setType(patternTy);
-      return EEP;
+      return EnumElementPattern::createImplicit(
+          patternTy, NLE->getLoc(), DeclNameLoc(NLE->getLoc()), NoneEnumElement,
+          /*subPattern*/ nullptr, EP->getDeclContext());
     } else {
       // ...but for non-optional types it can never match! Diagnose it.
       ctx.Diags
@@ -1388,14 +1385,12 @@ Pattern *TypeChecker::coercePatternToType(
     if (numExtraOptionals > 0) {
       Pattern *sub = IP;
       auto extraOpts =
-          llvm::drop_begin(inputTypeOptionals, castTypeOptionals.size());
+          llvm::drop_end(inputTypeOptionals, castTypeOptionals.size());
       for (auto extraOptTy : llvm::reverse(extraOpts)) {
         auto some = Context.getOptionalSomeDecl();
-        auto *base = TypeExpr::createImplicit(extraOptTy, Context);
-        sub = new (Context) EnumElementPattern(
-            base, IP->getStartLoc(), DeclNameLoc(IP->getEndLoc()),
-            some->createNameRef(), nullptr, sub, dc);
-        sub->setImplicit();
+        sub = EnumElementPattern::createImplicit(extraOptTy, IP->getStartLoc(),
+                                                 DeclNameLoc(IP->getEndLoc()),
+                                                 some, sub, dc);
       }
 
       P = sub;
