Reflection: Walk up the superclass chain when resolving associated types

Now that IRGen emits the right metadata, we can make use of it to fix
associated type lookup.

Fixes <rdar://problem/28331935>.

Author: slavapestov

include/swift/Reflection/TypeRefBuilder.h

@@ -138,6 +138,10 @@ class TypeRefBuilder {
   /// this TypeRefBuilder and are automatically released.
   std::vector<std::unique_ptr<const TypeRef>> TypeRefPool;
 
+  /// Cache for associated type lookups.
+  std::unordered_map<TypeRefID, const TypeRef *,
+                     TypeRefID::Hash, TypeRefID::Equal> AssociatedTypeCache;
+
   TypeConverter TC;
   MetadataSourceBuilder MSB;
 
@@ -295,14 +299,21 @@ class TypeRefBuilder {
 
   const AssociatedTypeDescriptor *
   lookupAssociatedTypes(const std::string &MangledTypeName,
-                        const DependentMemberTypeRef *DependentMember);
+                        const TypeRef *Protocol);
 
 public:
   TypeConverter &getTypeConverter() { return TC; }
 
   const TypeRef *
-  getDependentMemberTypeRef(const std::string &MangledTypeName,
-                            const DependentMemberTypeRef *DependentMember);
+  lookupTypeWitness(const std::string &MangledTypeName,
+                    const std::string &Member,
+                    const TypeRef *Protocol);
+
+  const TypeRef *
+  lookupSuperclass(const std::string &MangledTypeName);
+
+  const TypeRef *
+  lookupSuperclass(const TypeRef *TR);
 
   /// Load unsubstituted field types for a nominal type.
   const FieldDescriptor *getFieldTypeInfo(const TypeRef *TR);

stdlib/public/Reflection/TypeRef.cpp

@@ -22,12 +22,6 @@
 using namespace swift;
 using namespace reflection;
 
-[[noreturn]]
-static void unreachable(const char *Message) {
-  std::cerr << "fatal error: " << Message << "\n";
-  std::abort();
-}
-
 class PrintTypeRef : public TypeRefVisitor<PrintTypeRef, void> {
   std::ostream &OS;
   unsigned Indent;
@@ -354,14 +348,9 @@ bool TypeRef::isConcreteAfterSubstitutions(
 
 unsigned NominalTypeTrait::getDepth() const {
   if (auto P = Parent) {
-    switch (P->getKind()) {
-    case TypeRefKind::Nominal:
-      return 1 + cast<NominalTypeRef>(P)->getDepth();
-    case TypeRefKind::BoundGeneric:
-      return 1 + cast<BoundGenericTypeRef>(P)->getDepth();
-    default:
-      unreachable("Asked for depth on non-nominal typeref");
-    }
+    if (auto *Nominal = dyn_cast<NominalTypeRef>(P))
+      return 1 + Nominal->getDepth();
+    return 1 + cast<BoundGenericTypeRef>(P)->getDepth();
   }
 
   return 0;
@@ -657,20 +646,29 @@ class TypeRefSubstitution
 
     const TypeRef *TypeWitness = nullptr;
 
-    // Get the original type of the witness from the conformance.
-    switch (SubstBase->getKind()) {
-    case TypeRefKind::Nominal: {
-      auto Nominal = cast<NominalTypeRef>(SubstBase);
-      TypeWitness = Builder.getDependentMemberTypeRef(Nominal->getMangledName(), DM);
-      break;
-    }
-    case TypeRefKind::BoundGeneric: {
-      auto BG = cast<BoundGenericTypeRef>(SubstBase);
-      TypeWitness = Builder.getDependentMemberTypeRef(BG->getMangledName(), DM);
-      break;
-    }
-    default:
-      unreachable("Unknown base type");
+    while (TypeWitness == nullptr) {
+      auto &Member = DM->getMember();
+      auto *Protocol = DM->getProtocol();
+
+      // Get the original type of the witness from the conformance.
+      if (auto *Nominal = dyn_cast<NominalTypeRef>(SubstBase)) {
+        TypeWitness = Builder.lookupTypeWitness(Nominal->getMangledName(),
+                                                Member, Protocol);
+      } else {
+        auto BG = cast<BoundGenericTypeRef>(SubstBase);
+        TypeWitness = Builder.lookupTypeWitness(BG->getMangledName(),
+                                                Member, Protocol);
+      }
+
+      if (TypeWitness != nullptr)
+        break;
+
+      // If we didn't find the member type, check the superclass.
+      auto *Superclass = Builder.lookupSuperclass(SubstBase);
+      if (Superclass == nullptr)
+        break;
+
+      SubstBase = Superclass;
     }
 
     // We didn't find the member type, so return something to let the

stdlib/public/Reflection/TypeRefBuilder.cpp

@@ -28,46 +28,72 @@ using namespace reflection;
 
 TypeRefBuilder::TypeRefBuilder() : TC(*this) {}
 
-const AssociatedTypeDescriptor * TypeRefBuilder::
-lookupAssociatedTypes(const std::string &MangledTypeName,
-                      const DependentMemberTypeRef *DependentMember) {
+const TypeRef * TypeRefBuilder::
+lookupTypeWitness(const std::string &MangledTypeName,
+                  const std::string &Member,
+                  const TypeRef *Protocol) {
+  TypeRefID key;
+  key.addString(MangledTypeName);
+  key.addString(Member);
+  key.addPointer(Protocol);
+  auto found = AssociatedTypeCache.find(key);
+  if (found != AssociatedTypeCache.end())
+    return found->second;
+
   // Cache missed - we need to look through all of the assocty sections
   // for all images that we've been notified about.
   for (auto &Info : ReflectionInfos) {
     for (const auto &AssocTyDescriptor : Info.assocty) {
       std::string ConformingTypeName(AssocTyDescriptor.ConformingTypeName);
       if (ConformingTypeName.compare(MangledTypeName) != 0)
         continue;
+
       std::string ProtocolMangledName(AssocTyDescriptor.ProtocolTypeName);
       auto DemangledProto = Demangle::demangleTypeAsNode(ProtocolMangledName);
       auto TR = swift::remote::decodeMangledType(*this, DemangledProto);
 
-      auto &Conformance = *DependentMember->getProtocol();
-      if (auto Protocol = dyn_cast<ProtocolTypeRef>(TR)) {
-        if (*Protocol != Conformance)
+      if (Protocol != TR)
+        continue;
+
+      for (auto &AssocTy : AssocTyDescriptor) {
+        if (Member.compare(AssocTy.getName()) != 0)
           continue;
-        return &AssocTyDescriptor;
+
+        auto SubstitutedTypeName = AssocTy.getMangledSubstitutedTypeName();
+        auto Demangled = Demangle::demangleTypeAsNode(SubstitutedTypeName);
+        auto *TypeWitness = swift::remote::decodeMangledType(*this, Demangled);
+
+        AssociatedTypeCache.insert(std::make_pair(key, TypeWitness));
+        return TypeWitness;
       }
     }
   }
   return nullptr;
 }
 
 const TypeRef * TypeRefBuilder::
-getDependentMemberTypeRef(const std::string &MangledTypeName,
-                          const DependentMemberTypeRef *DependentMember) {
-
-  if (auto AssocTys = lookupAssociatedTypes(MangledTypeName, DependentMember)) {
-    for (auto &AssocTy : *AssocTys) {
-      if (DependentMember->getMember().compare(AssocTy.getName()) != 0)
-        continue;
+lookupSuperclass(const std::string &MangledTypeName) {
+  // Superclasses are recorded as a special associated type named 'super'
+  // on the 'AnyObject' protocol.
+  return lookupTypeWitness(MangledTypeName, "super",
+                           ProtocolTypeRef::create(*this, "Ps9AnyObject_"));
+}
 
-      auto SubstitutedTypeName = AssocTy.getMangledSubstitutedTypeName();
-      auto Demangled = Demangle::demangleTypeAsNode(SubstitutedTypeName);
-      return swift::remote::decodeMangledType(*this, Demangled);
-    }
+const TypeRef * TypeRefBuilder::
+lookupSuperclass(const TypeRef *TR) {
+  const TypeRef *Superclass = nullptr;
+
+  if (auto *Nominal = dyn_cast<NominalTypeRef>(TR)) {
+    Superclass = lookupSuperclass(Nominal->getMangledName());
+  } else {
+    auto BG = cast<BoundGenericTypeRef>(TR);
+    Superclass = lookupSuperclass(BG->getMangledName());
   }
-  return nullptr;
+
+  if (Superclass == nullptr)
+    return nullptr;
+
+  return Superclass->subst(*this, TR->getSubstMap());
 }
 
 const FieldDescriptor *

test/Reflection/Inputs/TypeLowering.swift

@@ -17,19 +17,39 @@ public protocol P {
   associatedtype B
 }
 
-public struct Foo<T, U> : P {
+public protocol Q : P {}
+
+public struct ConformsP<T, U> : P {
+  public typealias A = Box<U>
+  public typealias B = Box<T>
+}
+
+public struct ConformsQ<T, U> : Q {
   public typealias A = Box<U>
   public typealias B = Box<T>
 }
 
+public class Base<T, U> : P {
+  public typealias A = Box<T>
+  public typealias B = Box<U>
+}
+
+public class Derived : Base<Int8, Int16> {}
+
+public class GenericDerived<T> : Base<T, T> {}
+
 public struct Bar<T : P> {
   public let a: T.A
   public let b: T.B
   public let c: (T.A, T.B)
 }
 
 public struct AssocTypeStruct {
-  public let t: Bar<Foo<Int8, Int16>>
+  public let t1: Bar<ConformsP<Int8, Int16>>
+  public let t2: Bar<ConformsQ<Int8, Int16>>
+  public let t3: Bar<Base<Int8, Int16>>
+  public let t4: Bar<Derived>
+  public let t5: Bar<GenericDerived<Int8>>
 }
 
 public class C {}