AST: Factor out a new Requirement::getProtocolDecl() utility method

slavapestov · slavapestov · commit 53e06d69b5f2 · 2021-02-25T17:21:18.000-05:00
diff --git a/include/swift/AST/Requirement.h b/include/swift/AST/Requirement.h
@@ -74,6 +74,8 @@ class Requirement
     llvm_unreachable("Unhandled RequirementKind in switch.");
   }
 
+  ProtocolDecl *getProtocolDecl() const;
+
   SWIFT_DEBUG_DUMP;
   void dump(raw_ostream &out) const;
   void print(raw_ostream &os, const PrintOptions &opts) const;
diff --git a/include/swift/IRGen/Linking.h b/include/swift/IRGen/Linking.h
@@ -565,8 +565,7 @@ class LinkEntity {
     for (const auto &reqt : proto->getRequirementSignature()) {
       if (reqt.getKind() == RequirementKind::Conformance &&
           reqt.getFirstType()->getCanonicalType() == associatedType &&
-          reqt.getSecondType()->castTo<ProtocolType>()->getDecl() ==
-                                                                requirement) {
+          reqt.getProtocolDecl() == requirement) {
         return index;
       }
       ++index;
@@ -590,7 +589,7 @@ class LinkEntity {
     auto &reqt = proto->getRequirementSignature()[index];
     assert(reqt.getKind() == RequirementKind::Conformance);
     return { reqt.getFirstType()->getCanonicalType(),
-             reqt.getSecondType()->castTo<ProtocolType>()->getDecl() };
+             reqt.getProtocolDecl() };
   }
 
   static std::pair<CanType, ProtocolDecl*>
diff --git a/include/swift/SIL/SILWitnessVisitor.h b/include/swift/SIL/SILWitnessVisitor.h
@@ -64,9 +64,7 @@ template <class T> class SILWitnessVisitor : public ASTVisitor<T> {
       case RequirementKind::Conformance: {
         auto type = reqt.getFirstType()->getCanonicalType();
         assert(type->isTypeParameter());
-        auto requirement =
-          cast<ProtocolType>(reqt.getSecondType()->getCanonicalType())
-            ->getDecl();
+        auto requirement = reqt.getProtocolDecl();
 
         // ObjC protocols do not have witnesses.
         if (!Lowering::TypeConverter::protocolRequiresWitnessTable(requirement))
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -709,8 +709,7 @@ FuncDecl *ASTContext::getPlusFunctionOnRangeReplaceableCollection() const {
         continue;
       for (auto Req: FD->getGenericRequirements()) {
         if (Req.getKind() == RequirementKind::Conformance &&
-              Req.getSecondType()->getNominalOrBoundGenericNominal() ==
-            getRangeReplaceableCollectionDecl()) {
+              Req.getProtocolDecl() == getRangeReplaceableCollectionDecl()) {
           getImpl().PlusFunctionOnRangeReplaceableCollection = FD;
         }
       }
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -1584,8 +1584,7 @@ static bool containsRetroactiveConformance(
   for (auto requirement : rootConformance->getConditionalRequirements()) {
     if (requirement.getKind() != RequirementKind::Conformance)
       continue;
-    ProtocolDecl *proto =
-        requirement.getSecondType()->castTo<ProtocolType>()->getDecl();
+    ProtocolDecl *proto = requirement.getProtocolDecl();
     auto conformance = subMap.lookupConformance(
         requirement.getFirstType()->getCanonicalType(), proto);
     if (conformance.isInvalid()) {
@@ -2598,8 +2597,7 @@ void ASTMangler::appendRequirement(const Requirement &reqt) {
   case RequirementKind::Layout: {
   } break;
   case RequirementKind::Conformance: {
-    Type SecondTy = reqt.getSecondType();
-    appendProtocolName(SecondTy->castTo<ProtocolType>()->getDecl());
+    appendProtocolName(reqt.getProtocolDecl());
   } break;
   case RequirementKind::Superclass:
   case RequirementKind::SameType: {
@@ -3047,7 +3045,7 @@ static unsigned conformanceRequirementIndex(
       continue;
 
     if (req.getFirstType()->isEqual(entry.first) &&
-        req.getSecondType()->castTo<ProtocolType>()->getDecl() == entry.second)
+        req.getProtocolDecl() == entry.second)
       return result;
 
     ++result;
@@ -3175,8 +3173,7 @@ void ASTMangler::appendConcreteProtocolConformance(
       if (type->hasArchetype())
         type = type->mapTypeOutOfContext();
       CanType canType = type->getCanonicalType(CurGenericSignature);
-      auto proto =
-        conditionalReq.getSecondType()->castTo<ProtocolType>()->getDecl();
+      auto proto = conditionalReq.getProtocolDecl();
       
       ProtocolConformanceRef conformance;
       
diff --git a/lib/AST/ASTPrinter.cpp b/lib/AST/ASTPrinter.cpp
@@ -2533,8 +2533,7 @@ static bool usesFeatureRethrowsProtocol(
           ->getGenericSignatureOfContext()) {
     for (const auto &req : genericSig->getRequirements()) {
       if (req.getKind() == RequirementKind::Conformance &&
-          usesFeatureRethrowsProtocol(
-            req.getSecondType()->getAs<ProtocolType>()->getDecl(), checked))
+          usesFeatureRethrowsProtocol(req.getProtocolDecl(), checked))
         return true;
     }
   }
diff --git a/lib/AST/ASTVerifier.cpp b/lib/AST/ASTVerifier.cpp
@@ -2752,8 +2752,7 @@ class Verifier : public ASTWalker {
             abort();
           }
 
-          auto reqProto =
-            req.getSecondType()->castTo<ProtocolType>()->getDecl();
+          auto reqProto = req.getProtocolDecl();
           if (reqProto != conformances[idx].getRequirement()) {
             Out << "error: wrong protocol in signature conformances: have "
               << conformances[idx].getRequirement()->getName().str()
diff --git a/lib/AST/GenericSignature.cpp b/lib/AST/GenericSignature.cpp
@@ -308,9 +308,8 @@ CanGenericSignature::getCanonical(TypeArrayView<GenericTypeParamType> params,
     assert(reqt.getKind() == RequirementKind::Conformance &&
            "Only conformance requirements can have multiples");
 
-    auto prevProto =
-      prevReqt.getSecondType()->castTo<ProtocolType>()->getDecl();
-    auto proto = reqt.getSecondType()->castTo<ProtocolType>()->getDecl();
+    auto prevProto = prevReqt.getProtocolDecl();
+    auto proto = reqt.getProtocolDecl();
     assert(TypeDecl::compare(prevProto, proto) < 0 &&
            "Out-of-order conformance requirements");
   }
@@ -545,8 +544,7 @@ bool GenericSignatureImpl::isRequirementSatisfied(
 
   switch (requirement.getKind()) {
   case RequirementKind::Conformance: {
-    auto protocolType = requirement.getSecondType()->castTo<ProtocolType>();
-    auto protocol = protocolType->getDecl();
+    auto *protocol = requirement.getProtocolDecl();
 
     if (canFirstType->isTypeParameter())
       return requiresProtocol(canFirstType, protocol);
@@ -746,8 +744,7 @@ static bool hasConformanceInSignature(ArrayRef<Requirement> requirements,
   for (const auto &req: requirements) {
     if (req.getKind() == RequirementKind::Conformance &&
         req.getFirstType()->isEqual(subjectType) &&
-        req.getSecondType()->castTo<ProtocolType>()->getDecl()
-          == proto) {
+        req.getProtocolDecl() == proto) {
       return true;
     }
   }
@@ -1007,7 +1004,6 @@ bool Requirement::isCanonical() const {
   return true;
 }
 
-
 /// Get the canonical form of this requirement.
 Requirement Requirement::getCanonical() const {
   Type firstType = getFirstType();
@@ -1029,3 +1025,8 @@ Requirement Requirement::getCanonical() const {
   }
   llvm_unreachable("Unhandled RequirementKind in switch");
 }
+
+ProtocolDecl *Requirement::getProtocolDecl() const {
+  assert(getKind() == RequirementKind::Conformance);
+  return getSecondType()->castTo<ProtocolType>()->getDecl();
+}
diff --git a/lib/AST/NameLookup.cpp b/lib/AST/NameLookup.cpp
@@ -2700,8 +2700,7 @@ swift::getDirectlyInheritedNominalTypeDecls(
       if (!req.getFirstType()->isEqual(protoSelfTy))
         continue;
 
-      result.emplace_back(req.getSecondType()->castTo<ProtocolType>()->getDecl(),
-                          loc);
+      result.emplace_back(req.getProtocolDecl(), loc);
     }
     return result;
   }
diff --git a/lib/AST/ProtocolConformance.cpp b/lib/AST/ProtocolConformance.cpp
@@ -572,8 +572,7 @@ void NormalProtocolConformance::setSignatureConformances(
              "Should have interface types here");
       assert(idx < conformances.size());
       assert(conformances[idx].isInvalid() ||
-             conformances[idx].getRequirement() ==
-               req.getSecondType()->castTo<ProtocolType>()->getDecl());
+             conformances[idx].getRequirement() == req.getProtocolDecl());
       ++idx;
     }
   }
@@ -772,7 +771,7 @@ NormalProtocolConformance::getAssociatedConformance(Type assocType,
     if (reqt.getKind() == RequirementKind::Conformance) {
       // Is this the conformance we're looking for?
       if (reqt.getFirstType()->isEqual(assocType) &&
-          reqt.getSecondType()->castTo<ProtocolType>()->getDecl() == protocol)
+          reqt.getProtocolDecl() == protocol)
         return getSignatureConformances()[conformanceIndex];
 
       ++conformanceIndex;
@@ -840,7 +839,7 @@ void NormalProtocolConformance::finishSignatureConformances() {
       auto *depMemTy = origTy->castTo<DependentMemberType>();
       substTy = recursivelySubstituteBaseType(module, this, depMemTy);
     }
-    auto reqProto = req.getSecondType()->castTo<ProtocolType>()->getDecl();
+    auto reqProto = req.getProtocolDecl();
 
     // Looking up a conformance for a contextual type and mapping the
     // conformance context produces a more accurate result than looking
diff --git a/lib/AST/SubstitutionMap.cpp b/lib/AST/SubstitutionMap.cpp
@@ -234,8 +234,7 @@ SubstitutionMap SubstitutionMap::get(GenericSignature genericSig,
 
     CanType depTy = req.getFirstType()->getCanonicalType();
     auto replacement = depTy.subst(subs, lookupConformance);
-    auto protoType = req.getSecondType()->castTo<ProtocolType>();
-    auto proto = protoType->getDecl();
+    auto *proto = req.getProtocolDecl();
     auto conformance = lookupConformance(depTy, replacement, proto);
     conformances.push_back(conformance);
   }
@@ -334,38 +333,27 @@ SubstitutionMap::lookupConformance(CanType type, ProtocolDecl *proto) const {
 
   auto genericSig = getGenericSignature();
 
-  // Fast path
-  unsigned index = 0;
-  for (auto reqt : genericSig->getRequirements()) {
-    if (reqt.getKind() == RequirementKind::Conformance) {
-      if (reqt.getFirstType()->isEqual(type) &&
-          reqt.getSecondType()->isEqual(proto->getDeclaredInterfaceType()))
-        return getConformances()[index];
+  auto getSignatureConformance =
+      [&](Type type, ProtocolDecl *proto) -> Optional<ProtocolConformanceRef> {
+    unsigned index = 0;
+    for (auto reqt : genericSig->getRequirements()) {
+      if (reqt.getKind() == RequirementKind::Conformance) {
+        if (reqt.getFirstType()->isEqual(type) &&
+            reqt.getProtocolDecl() == proto)
+          return getConformances()[index];
 
-      ++index;
-    }
-  }
-
-  // Retrieve the starting conformance from the conformance map.
-  auto getInitialConformance =
-      [&](Type type, ProtocolDecl *proto) -> ProtocolConformanceRef {
-    unsigned conformanceIndex = 0;
-    for (const auto &req : getGenericSignature()->getRequirements()) {
-      if (req.getKind() != RequirementKind::Conformance)
-        continue;
-
-      // Is this the conformance we're looking for?
-      if (req.getFirstType()->isEqual(type) &&
-          req.getSecondType()->castTo<ProtocolType>()->getDecl() == proto) {
-        return getConformances()[conformanceIndex];
+        ++index;
       }
-
-      ++conformanceIndex;
     }
 
-    return ProtocolConformanceRef::forInvalid();
+    return None;
   };
 
+  // Fast path -- check if the generic signature directly states the
+  // conformance.
+  if (auto directConformance = getSignatureConformance(type, proto))
+    return *directConformance;
+
   // Check whether the superclass conforms.
   if (auto superclass = genericSig->getSuperclassBound(type)) {
     LookUpConformanceInSignature lookup(getGenericSignature().getPointer());
@@ -388,16 +376,16 @@ SubstitutionMap::lookupConformance(CanType type, ProtocolDecl *proto) const {
   for (const auto &step : accessPath) {
     // For the first step, grab the initial conformance.
     if (conformance.isInvalid()) {
-      conformance = getInitialConformance(step.first, step.second);
-      if (conformance.isInvalid())
-        return ProtocolConformanceRef::forInvalid();
+      if (auto initialConformance = getSignatureConformance(
+            step.first, step.second)) {
+        conformance = *initialConformance;
+        continue;
+      }
 
-      continue;
+      // We couldn't find the initial conformance, fail.
+      return ProtocolConformanceRef::forInvalid();
     }
 
-    if (conformance.isInvalid())
-      return conformance;
-
     // If we've hit an abstract conformance, everything from here on out is
     // abstract.
     // FIXME: This may not always be true, but it holds for now.
@@ -436,6 +424,8 @@ SubstitutionMap::lookupConformance(CanType type, ProtocolDecl *proto) const {
 
     // Get the associated conformance.
     conformance = concrete->getAssociatedConformance(step.first, step.second);
+    if (conformance.isInvalid())
+      return conformance;
   }
 
   return conformance;
diff --git a/lib/AST/Type.cpp b/lib/AST/Type.cpp
@@ -2012,7 +2012,7 @@ class IsBindableVisitor
       if (req.getKind() != RequirementKind::Conformance) continue;
 
       auto canTy = req.getFirstType()->getCanonicalType();
-      auto *proto = req.getSecondType()->castTo<ProtocolType>()->getDecl();
+      auto *proto = req.getProtocolDecl();
       auto origConf = origSubMap.lookupConformance(canTy, proto);
       auto substConf = substSubMap.lookupConformance(canTy, proto);
 
diff --git a/lib/IRGen/GenMeta.cpp b/lib/IRGen/GenMeta.cpp
@@ -5243,8 +5243,7 @@ GenericRequirementsMetadata irgen::addGenericRequirements(
       break;
 
     case RequirementKind::Conformance: {
-      auto protocol = requirement.getSecondType()->castTo<ProtocolType>()
-        ->getDecl();
+      auto protocol = requirement.getProtocolDecl();
 
       // Marker protocols do not record generic requirements at all.
       if (protocol->isMarkerProtocol()) {
diff --git a/lib/IRGen/GenProto.cpp b/lib/IRGen/GenProto.cpp
@@ -266,8 +266,7 @@ irgen::enumerateGenericSignatureRequirements(CanGenericSignature signature,
 
       case RequirementKind::Conformance: {
         auto type = CanType(reqt.getFirstType());
-        auto protocol =
-          cast<ProtocolType>(CanType(reqt.getSecondType()))->getDecl();
+        auto protocol = reqt.getProtocolDecl();
         if (Lowering::TypeConverter::protocolRequiresWitnessTable(protocol)) {
           callback({type, protocol});
         }
@@ -926,7 +925,7 @@ static bool isDependentConformance(
     if (req.getKind() != RequirementKind::Conformance)
       continue;
 
-    auto assocProtocol = req.getSecondType()->castTo<ProtocolType>()->getDecl();
+    auto assocProtocol = req.getProtocolDecl();
     if (assocProtocol->isObjC())
       continue;
 
@@ -2801,8 +2800,7 @@ void NecessaryBindings::addAbstractConditionalRequirements(
   for (auto req : condRequirements) {
     if (req.getKind() != RequirementKind::Conformance)
       continue;
-    auto *proto =
-        req.getSecondType()->castTo<ProtocolType>()->getDecl();
+    auto *proto = req.getProtocolDecl();
     auto ty = req.getFirstType()->getCanonicalType();
     auto archetype = dyn_cast<ArchetypeType>(ty);
     if (!archetype)
diff --git a/lib/SIL/IR/SILWitnessTable.cpp b/lib/SIL/IR/SILWitnessTable.cpp
@@ -187,7 +187,7 @@ bool SILWitnessTable::enumerateWitnessTableConditionalConformances(
     if (req.getKind() != RequirementKind::Conformance)
       continue;
 
-    auto proto = req.getSecondType()->castTo<ProtocolType>()->getDecl();
+    auto proto = req.getProtocolDecl();
 
     if (Lowering::TypeConverter::protocolRequiresWitnessTable(proto)) {
       if (fn(conformanceIndex, req.getFirstType()->getCanonicalType(), proto))
diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp
@@ -266,8 +266,7 @@ static bool diagnoseUnsatisfiedRequirements(ADContext &context,
     }
     // Check conformance requirements.
     case RequirementKind::Conformance: {
-      auto protocolType = req.getSecondType()->castTo<ProtocolType>();
-      auto protocol = protocolType->getDecl();
+      auto *protocol = req.getProtocolDecl();
       assert(protocol && "Expected protocol in generic signature requirement");
       // If the first type does not conform to the second type in the current
       // module, then record the unsatisfied requirement.
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -1381,14 +1381,13 @@ void ConstraintSystem::openGenericRequirements(
     auto kind = req.getKind();
     switch (kind) {
     case RequirementKind::Conformance: {
-      auto proto = req.getSecondType()->castTo<ProtocolType>();
-      auto protoDecl = proto->getDecl();
+      auto protoDecl = req.getProtocolDecl();
       // Determine whether this is the protocol 'Self' constraint we should
       // skip.
       if (skipProtocolSelfConstraint && protoDecl == outerDC &&
           protoDecl->getSelfInterfaceType()->isEqual(req.getFirstType()))
         continue;
-      openedReq = Requirement(kind, openedFirst, proto);
+      openedReq = Requirement(kind, openedFirst, req.getSecondType());
       break;
     }
     case RequirementKind::Superclass:
diff --git a/lib/Sema/DerivedConformanceDifferentiable.cpp b/lib/Sema/DerivedConformanceDifferentiable.cpp
diff --git a/lib/Sema/TypeCheckEffects.cpp b/lib/Sema/TypeCheckEffects.cpp
diff --git a/lib/Sema/TypeCheckGeneric.cpp b/lib/Sema/TypeCheckGeneric.cpp
diff --git a/lib/Sema/TypeCheckProtocol.cpp b/lib/Sema/TypeCheckProtocol.cpp
diff --git a/lib/Sema/TypeCheckProtocolInference.cpp b/lib/Sema/TypeCheckProtocolInference.cpp
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftDocSupport.cpp b/tools/SourceKit/lib/SwiftLang/SwiftDocSupport.cpp

Original file line number	Diff line number	Diff line change
`@@ -74,6 +74,8 @@ class Requirement`
`74`	`74`	`llvm_unreachable("Unhandled RequirementKind in switch.");`
`75`	`75`	`}`
`76`	`76`
	`77`	`+ ProtocolDecl *getProtocolDecl() const;`
	`78`	`+`
`77`	`79`	`SWIFT_DEBUG_DUMP;`
`78`	`80`	`void dump(raw_ostream &out) const;`
`79`	`81`	`void print(raw_ostream &os, const PrintOptions &opts) const;`
Original file line number	Diff line number	Diff line change
`@@ -709,8 +709,7 @@ FuncDecl *ASTContext::getPlusFunctionOnRangeReplaceableCollection() const {`
`709`	`709`	`continue;`
`710`	`710`	`for (auto Req: FD->getGenericRequirements()) {`
`711`	`711`	`if (Req.getKind() == RequirementKind::Conformance &&`
`712`		`- Req.getSecondType()->getNominalOrBoundGenericNominal() ==`
`713`		`- getRangeReplaceableCollectionDecl()) {`
	`712`	`+ Req.getProtocolDecl() == getRangeReplaceableCollectionDecl()) {`
`714`	`713`	`getImpl().PlusFunctionOnRangeReplaceableCollection = FD;`
`715`	`714`	`}`
`716`	`715`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2533,8 +2533,7 @@ static bool usesFeatureRethrowsProtocol(`
`2533`	`2533`	`->getGenericSignatureOfContext()) {`
`2534`	`2534`	`for (const auto &req : genericSig->getRequirements()) {`
`2535`	`2535`	`if (req.getKind() == RequirementKind::Conformance &&`
`2536`		`- usesFeatureRethrowsProtocol(`
`2537`		`- req.getSecondType()->getAs<ProtocolType>()->getDecl(), checked))`
	`2536`	`+ usesFeatureRethrowsProtocol(req.getProtocolDecl(), checked))`
`2538`	`2537`	`return true;`
`2539`	`2538`	`}`
`2540`	`2539`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2752,8 +2752,7 @@ class Verifier : public ASTWalker {`
`2752`	`2752`	`abort();`
`2753`	`2753`	`}`
`2754`	`2754`
`2755`		`- auto reqProto =`
`2756`		`- req.getSecondType()->castTo<ProtocolType>()->getDecl();`
	`2755`	`+ auto reqProto = req.getProtocolDecl();`
`2757`	`2756`	`if (reqProto != conformances[idx].getRequirement()) {`
`2758`	`2757`	`Out << "error: wrong protocol in signature conformances: have "`
`2759`	`2758`	`<< conformances[idx].getRequirement()->getName().str()`
Original file line number	Diff line number	Diff line change
`@@ -2700,8 +2700,7 @@ swift::getDirectlyInheritedNominalTypeDecls(`
`2700`	`2700`	`if (!req.getFirstType()->isEqual(protoSelfTy))`
`2701`	`2701`	`continue;`
`2702`	`2702`
`2703`		`- result.emplace_back(req.getSecondType()->castTo<ProtocolType>()->getDecl(),`
`2704`		`- loc);`
	`2703`	`+ result.emplace_back(req.getProtocolDecl(), loc);`
`2705`	`2704`	`}`
`2706`	`2705`	`return result;`
`2707`	`2706`	`}`