Clean up WhereClauseOwner

Author: AnthonyLatsis

include/swift/AST/TypeCheckRequests.h

@@ -39,6 +39,7 @@ struct PropertyWrapperBackingPropertyInfo;
 struct PropertyWrapperMutability;
 class RequirementRepr;
 class SpecializeAttr;
+class TrailingWhereClause;
 class TypeAliasDecl;
 struct TypeLoc;
 class Witness;
@@ -369,9 +370,10 @@ struct WhereClauseOwner {
 
   /// The source of the where clause, which can be a generic parameter list
   /// or a declaration that can have a where clause.
-  llvm::PointerUnion<GenericParamList *, Decl *, SpecializeAttr *> source;
+  llvm::PointerUnion<GenericParamList *, TrailingWhereClause *, SpecializeAttr *> source;
 
-  WhereClauseOwner(Decl *decl);
+  WhereClauseOwner(GenericContext *genCtx);
+  WhereClauseOwner(AssociatedTypeDecl *atd);
 
   WhereClauseOwner(DeclContext *dc, GenericParamList *genericParams)
       : dc(dc), source(genericParams) {}
@@ -382,13 +384,12 @@ struct WhereClauseOwner {
   SourceLoc getLoc() const;
 
   friend hash_code hash_value(const WhereClauseOwner &owner) {
-    return llvm::hash_combine(owner.dc, owner.source.getOpaqueValue());
+    return llvm::hash_value(owner.source.getOpaqueValue());
   }
 
   friend bool operator==(const WhereClauseOwner &lhs,
                          const WhereClauseOwner &rhs) {
-    return lhs.dc == rhs.dc &&
-           lhs.source.getOpaqueValue() == rhs.source.getOpaqueValue();
+    return lhs.source.getOpaqueValue() == rhs.source.getOpaqueValue();
   }
 
   friend bool operator!=(const WhereClauseOwner &lhs,

lib/AST/TypeCheckRequests.cpp

@@ -337,12 +337,20 @@ void RequirementSignatureRequest::cacheResult(ArrayRef<Requirement> value) const
 // Requirement computation.
 //----------------------------------------------------------------------------//
 
-WhereClauseOwner::WhereClauseOwner(Decl *decl)
-  : dc(decl->getInnermostDeclContext()), source(decl) { }
+WhereClauseOwner::WhereClauseOwner(GenericContext *genCtx): dc(genCtx) {
+  if (const auto whereClause = genCtx->getTrailingWhereClause())
+    source = whereClause;
+  else
+    source = genCtx->getGenericParams();
+}
+
+WhereClauseOwner::WhereClauseOwner(AssociatedTypeDecl *atd)
+    : dc(atd->getInnermostDeclContext()),
+      source(atd->getTrailingWhereClause()) {}
 
 SourceLoc WhereClauseOwner::getLoc() const {
-  if (auto decl = source.dyn_cast<Decl *>())
-    return decl->getLoc();
+  if (auto where = source.dyn_cast<TrailingWhereClause *>())
+    return where->getWhereLoc();
 
   if (auto attr = source.dyn_cast<SpecializeAttr *>())
     return attr->getLocation();
@@ -352,8 +360,8 @@ SourceLoc WhereClauseOwner::getLoc() const {
 
 void swift::simple_display(llvm::raw_ostream &out,
                            const WhereClauseOwner &owner) {
-  if (auto decl = owner.source.dyn_cast<Decl *>()) {
-    simple_display(out, decl);
+  if (auto where = owner.source.dyn_cast<TrailingWhereClause *>()) {
+    simple_display(out, owner.dc->getAsDecl());
   } else if (owner.source.is<SpecializeAttr *>()) {
     out << "@_specialize";
   } else {
@@ -375,36 +383,13 @@ void RequirementRequest::noteCycleStep(DiagnosticEngine &diags) const {
 }
 
 MutableArrayRef<RequirementRepr> WhereClauseOwner::getRequirements() const {
-  if (auto genericParams = source.dyn_cast<GenericParamList *>()) {
+  if (const auto genericParams = source.dyn_cast<GenericParamList *>()) {
     return genericParams->getRequirements();
-  }
-
-  if (auto attr = source.dyn_cast<SpecializeAttr *>()) {
+  } else if (const auto attr = source.dyn_cast<SpecializeAttr *>()) {
     if (auto whereClause = attr->getTrailingWhereClause())
       return whereClause->getRequirements();
-    
-    return { };
-  }
-
-  auto decl = source.dyn_cast<Decl *>();
-  if (!decl)
-    return { };
-
-  if (auto proto = dyn_cast<ProtocolDecl>(decl)) {
-    if (auto whereClause = proto->getTrailingWhereClause())
-      return whereClause->getRequirements();
-
-    return { };
-  }
-
-  if (auto assocType = dyn_cast<AssociatedTypeDecl>(decl)) {
-    if (auto whereClause = assocType->getTrailingWhereClause())
-      return whereClause->getRequirements();
-  }
-
-  if (auto genericContext = decl->getAsGenericContext()) {
-    if (auto genericParams = genericContext->getGenericParams())
-      return genericParams->getRequirements();
+  } else if (const auto whereClause = source.get<TrailingWhereClause *>()) {
+    return whereClause->getRequirements();
   }
 
   return { };

lib/Sema/TypeCheckDecl.cpp

@@ -2218,13 +2218,13 @@ SelfAccessKindRequest::evaluate(Evaluator &evaluator, FuncDecl *FD) const {
   return SelfAccessKind::NonMutating;
 }
 
-/// Check the requirements in the where clause of the given \c source
+/// Check the requirements in the where clause of the given \c atd
 /// to ensure that they don't introduce additional 'Self' requirements.
 static void checkProtocolSelfRequirements(ProtocolDecl *proto,
-                                          TypeDecl *source) {
-  WhereClauseOwner(source).visitRequirements(
+                                          AssociatedTypeDecl *atd) {
+  WhereClauseOwner(atd).visitRequirements(
       TypeResolutionStage::Interface,
-      [&](const Requirement &req, RequirementRepr *reqRepr) {
+      [proto](const Requirement &req, RequirementRepr *reqRepr) {
         switch (req.getKind()) {
         case RequirementKind::Conformance:
         case RequirementKind::Layout: