AST: Generalize ASTContext::getOverrideGenericSignature() to work with protocols

We want to avoid feeding invalid type parameters into the Requirement Machine
when checking if a protocol requirement overrides another protocol requirement
in an inherited protocol.

In order to do that we need to make sure the potential override has a
compatible generic signature before we attempt substitution, just like we
already do for classes.

To do that, we need a way to 'rewrite' a generic signature for the base
requirement into a generic signature that can be compared with the derived
requirement, just as we do with class methods.

For protocols this is a little easier than with class methods, since protocols
only ever have a single 'Self' generic parameter. So just add the
'Self : DerivedProto' requirement to the base requirement's signature and
rebuild.

Author: slavapestov

lib/AST/ASTContext.cpp

@@ -101,14 +101,14 @@ using AssociativityCacheType =
 struct OverrideSignatureKey {
   GenericSignature baseMethodSig;
   GenericSignature derivedMethodSig;
-  Decl *subclassDecl;
+  NominalTypeDecl *derivedNominal;
 
   OverrideSignatureKey(GenericSignature baseMethodSignature,
                        GenericSignature derivedMethodSignature,
-                       Decl *subclassDecl)
+                       NominalTypeDecl *derivedNominal)
     : baseMethodSig(baseMethodSignature),
       derivedMethodSig(derivedMethodSignature),
-      subclassDecl(subclassDecl) {}
+      derivedNominal(derivedNominal) {}
 };
 
 namespace llvm {
@@ -120,27 +120,27 @@ template <> struct DenseMapInfo<OverrideSignatureKey> {
                       const OverrideSignatureKey rhs) {
     return lhs.baseMethodSig.getPointer() == rhs.baseMethodSig.getPointer() &&
            lhs.derivedMethodSig.getPointer() == rhs.derivedMethodSig.getPointer() &&
-           lhs.subclassDecl == rhs.subclassDecl;
+           lhs.derivedNominal == rhs.derivedNominal;
   }
 
   static inline OverrideSignatureKey getEmptyKey() {
     return OverrideSignatureKey(DenseMapInfo<GenericSignature>::getEmptyKey(),
                                 DenseMapInfo<GenericSignature>::getEmptyKey(),
-                                DenseMapInfo<Decl *>::getEmptyKey());
+                                DenseMapInfo<NominalTypeDecl *>::getEmptyKey());
   }
 
   static inline OverrideSignatureKey getTombstoneKey() {
     return OverrideSignatureKey(
         DenseMapInfo<GenericSignature>::getTombstoneKey(),
         DenseMapInfo<GenericSignature>::getTombstoneKey(),
-        DenseMapInfo<Decl *>::getTombstoneKey());
+        DenseMapInfo<NominalTypeDecl *>::getTombstoneKey());
   }
 
   static unsigned getHashValue(const OverrideSignatureKey &Val) {
     return hash_combine(
         DenseMapInfo<GenericSignature>::getHashValue(Val.baseMethodSig),
         DenseMapInfo<GenericSignature>::getHashValue(Val.derivedMethodSig),
-        DenseMapInfo<Decl *>::getHashValue(Val.subclassDecl));
+        DenseMapInfo<NominalTypeDecl *>::getHashValue(Val.derivedNominal));
   }
 };
 } // namespace llvm
@@ -5214,11 +5214,11 @@ ASTContext::getOverrideGenericSignature(const ValueDecl *base,
   assert(isa<AbstractFunctionDecl>(base) || isa<SubscriptDecl>(base));
   assert(isa<AbstractFunctionDecl>(derived) || isa<SubscriptDecl>(derived));
 
-  const auto baseClass = base->getDeclContext()->getSelfClassDecl();
-  const auto derivedClass = derived->getDeclContext()->getSelfClassDecl();
+  const auto baseNominal = base->getDeclContext()->getSelfNominalTypeDecl();
+  const auto derivedNominal = derived->getDeclContext()->getSelfNominalTypeDecl();
 
-  assert(baseClass != nullptr);
-  assert(derivedClass != nullptr);
+  assert(baseNominal != nullptr);
+  assert(derivedNominal != nullptr);
 
   const auto baseGenericSig =
       base->getAsGenericContext()->getGenericSignature();
@@ -5228,10 +5228,6 @@ ASTContext::getOverrideGenericSignature(const ValueDecl *base,
   if (base == derived)
     return derivedGenericSig;
 
-  const auto derivedSuperclass = derivedClass->getSuperclass();
-  if (derivedSuperclass.isNull())
-    return nullptr;
-
   if (derivedGenericSig.isNull())
     return nullptr;
 
@@ -5240,21 +5236,14 @@ ASTContext::getOverrideGenericSignature(const ValueDecl *base,
 
   auto key = OverrideSignatureKey(baseGenericSig,
                                   derivedGenericSig,
-                                  derivedClass);
+                                  derivedNominal);
 
   if (getImpl().overrideSigCache.find(key) !=
       getImpl().overrideSigCache.end()) {
     return getImpl().overrideSigCache.lookup(key);
   }
 
-  const auto derivedClassSig = derivedClass->getGenericSignature();
-
-  unsigned derivedDepth = 0;
-  unsigned baseDepth = 0;
-  if (derivedClassSig)
-    derivedDepth = derivedClassSig.getGenericParams().back()->getDepth() + 1;
-  if (const auto baseClassSig = baseClass->getGenericSignature())
-    baseDepth = baseClassSig.getGenericParams().back()->getDepth() + 1;
+  const auto derivedNominalSig = derivedNominal->getGenericSignature();
 
   SmallVector<GenericTypeParamType *, 2> addedGenericParams;
   if (const auto *gpList = derived->getAsGenericContext()->getGenericParams()) {
@@ -5264,38 +5253,59 @@ ASTContext::getOverrideGenericSignature(const ValueDecl *base,
     }
   }
 
-  const auto subMap = derivedSuperclass->getContextSubstitutionMap(
-      derivedClass->getModuleContext(), baseClass);
+  SmallVector<Requirement, 2> addedRequirements;
 
-  auto substFn = [&](SubstitutableType *type) -> Type {
-    auto *gp = cast<GenericTypeParamType>(type);
+  if (isa<ProtocolDecl>(baseNominal)) {
+    assert(isa<ProtocolDecl>(derivedNominal));
 
-    if (gp->getDepth() < baseDepth) {
-      return Type(gp).subst(subMap);
+    for (auto reqt : baseGenericSig.getRequirements()) {
+      addedRequirements.push_back(reqt);
     }
+  } else {
+    const auto derivedSuperclass = cast<ClassDecl>(derivedNominal)
+        ->getSuperclass();
+    if (derivedSuperclass.isNull())
+      return nullptr;
 
-    return CanGenericTypeParamType::get(
-        gp->isTypeSequence(), gp->getDepth() - baseDepth + derivedDepth,
-        gp->getIndex(), *this);
-  };
+    unsigned derivedDepth = 0;
+    unsigned baseDepth = 0;
+    if (derivedNominalSig)
+      derivedDepth = derivedNominalSig.getGenericParams().back()->getDepth() + 1;
+    if (const auto baseNominalSig = baseNominal->getGenericSignature())
+      baseDepth = baseNominalSig.getGenericParams().back()->getDepth() + 1;
 
-  auto lookupConformanceFn =
-      [&](CanType depTy, Type substTy,
-          ProtocolDecl *proto) -> ProtocolConformanceRef {
-    if (auto conf = subMap.lookupConformance(depTy, proto))
-      return conf;
+    const auto subMap = derivedSuperclass->getContextSubstitutionMap(
+        derivedNominal->getModuleContext(), baseNominal);
 
-    return ProtocolConformanceRef(proto);
-  };
+    auto substFn = [&](SubstitutableType *type) -> Type {
+      auto *gp = cast<GenericTypeParamType>(type);
 
-  SmallVector<Requirement, 2> addedRequirements;
-  for (auto reqt : baseGenericSig.getRequirements()) {
-    if (auto substReqt = reqt.subst(substFn, lookupConformanceFn)) {
-      addedRequirements.push_back(*substReqt);
+      if (gp->getDepth() < baseDepth) {
+        return Type(gp).subst(subMap);
+      }
+
+      return CanGenericTypeParamType::get(
+          gp->isTypeSequence(), gp->getDepth() - baseDepth + derivedDepth,
+          gp->getIndex(), *this);
+    };
+
+    auto lookupConformanceFn =
+        [&](CanType depTy, Type substTy,
+            ProtocolDecl *proto) -> ProtocolConformanceRef {
+      if (auto conf = subMap.lookupConformance(depTy, proto))
+        return conf;
+
+      return ProtocolConformanceRef(proto);
+    };
+
+    for (auto reqt : baseGenericSig.getRequirements()) {
+      if (auto substReqt = reqt.subst(substFn, lookupConformanceFn)) {
+        addedRequirements.push_back(*substReqt);
+      }
     }
   }
 
-  auto genericSig = buildGenericSignature(*this, derivedClassSig,
+  auto genericSig = buildGenericSignature(*this, derivedNominalSig,
                                           std::move(addedGenericParams),
                                           std::move(addedRequirements));
   getImpl().overrideSigCache.insert(std::make_pair(key, genericSig));