Moving nested archetype storage into the generic environment.

Instead of storing nested archetypes hierarchically in `ArchetypeType`,
store them in a map (indexed by dependent member type) on the generic
environment itself. With this, we no longer need to create archetypes
for every type along the path, because archetypes are findable via
generic environment + interface type.

Author: DougGregor

include/swift/AST/GenericEnvironment.h

@@ -72,9 +72,12 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
     Opaque,
   };
 
+  class NestedTypeStorage;
+
 private:
   mutable llvm::PointerIntPair<GenericSignature, 2, Kind> SignatureAndKind{
       GenericSignature(), Kind::Normal};
+  NestedTypeStorage *nestedTypeStorage = nullptr;
 
   friend TrailingObjects;
   friend OpaqueTypeArchetypeType;
@@ -93,6 +96,9 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// generic signature.
   ArrayRef<Type> getContextTypes() const;
 
+  /// Get the nested type storage, allocating it if required.
+  NestedTypeStorage &getOrCreateNestedTypeStorage();
+
   explicit GenericEnvironment(GenericSignature signature, Kind kind);
   explicit GenericEnvironment(
       GenericSignature signature, OpaqueTypeDecl *opaque, SubstitutionMap subs);

include/swift/AST/Types.h

@@ -369,8 +369,7 @@ class alignas(1 << TypeAlignInBits) TypeBase
     NumParams : 16
   );
 
-  SWIFT_INLINE_BITFIELD_FULL(ArchetypeType, TypeBase, 1+1+1+16,
-    ExpandedNestedTypes : 1,
+  SWIFT_INLINE_BITFIELD_FULL(ArchetypeType, TypeBase, 1+1+16,
     HasSuperclass : 1,
     HasLayoutConstraint : 1,
     : NumPadBits,
@@ -5395,13 +5394,8 @@ class ArchetypeType : public SubstitutableType,
     return Bits.ArchetypeType.HasLayoutConstraint ? 1 : 0;
   }
   Type InterfaceType;
-  MutableArrayRef<std::pair<Identifier, Type>> NestedTypes;
   GenericEnvironment *Environment = nullptr;
 
-  void populateNestedTypes() const;
-  void resolveNestedType(std::pair<Identifier, Type> &nested) const;
-
-  
   // Helper to get the trailing objects of one of the subclasses.
   template<typename Type>
   const Type *getSubclassTrailingObjects() const;
@@ -5455,20 +5449,14 @@ class ArchetypeType : public SubstitutableType,
     return !getConformsTo().empty() || getSuperclass();
   }
 
-  /// Retrieve the nested type with the given name.
-  Type getNestedType(Identifier Name) const;
+  /// Retrieve the nested type with the given associated type.
+  Type getNestedType(AssociatedTypeDecl *assocType);
 
-  /// Retrieve the nested type with the given name, if it's already
-  /// known.
+  /// Retrieve the nested type with the given name.
   ///
-  /// This is an implementation detail used by the generic signature builder.
-  Optional<Type> getNestedTypeIfKnown(Identifier Name) const;
-
-  /// Check if the archetype contains a nested type with the given name.
-  bool hasNestedType(Identifier Name) const;
-
-  /// Register a nested type with the given name.
-  void registerNestedType(Identifier name, Type nested);
+  /// This is a slow operation because it must scan all of the protocols to
+  /// which the archetype conforms.
+  Type getNestedTypeByName(Identifier name);
 
   /// Retrieve the parent of this archetype, or null if this is a
   /// primary archetype.

lib/AST/ASTDemangler.cpp

@@ -638,9 +638,8 @@ Type ASTBuilder::createDependentMemberType(StringRef member,
   auto identifier = Ctx.getIdentifier(member);
 
   if (auto *archetype = base->getAs<ArchetypeType>()) {
-    if (archetype->hasNestedType(identifier))
-      return archetype->getNestedType(identifier);
-
+      if (Type memberType = archetype->getNestedTypeByName(identifier))
+        return memberType;
   }
 
   if (base->isTypeParameter()) {
@@ -656,8 +655,8 @@ Type ASTBuilder::createDependentMemberType(StringRef member,
   auto identifier = Ctx.getIdentifier(member);
 
   if (auto *archetype = base->getAs<ArchetypeType>()) {
-    if (archetype->hasNestedType(identifier))
-      return archetype->getNestedType(identifier);
+    if (auto assocType = protocol->getAssociatedType(identifier))
+      return archetype->getNestedType(assocType);
   }
 
   if (base->isTypeParameter()) {

lib/AST/GenericEnvironment.cpp

@@ -218,6 +218,31 @@ Type TypeBase::mapTypeOutOfContext() {
     SubstFlags::AllowLoweredTypes);
 }
 
+class GenericEnvironment::NestedTypeStorage
+    : public llvm::DenseMap<CanType, Type> { };
+
+auto GenericEnvironment::getOrCreateNestedTypeStorage() -> NestedTypeStorage & {
+  if (nestedTypeStorage)
+    return *nestedTypeStorage;
+
+  nestedTypeStorage = new NestedTypeStorage();
+  ASTContext &ctx = getGenericParams().front()->getASTContext();
+  ctx.addCleanup([nestedTypeStorage=this->nestedTypeStorage]() {
+    delete nestedTypeStorage;
+  });
+
+  return *nestedTypeStorage;
+}
+
+static Type stripBoundDependentMemberTypes(Type t) {
+  if (auto *depMemTy = t->getAs<DependentMemberType>()) {
+    return DependentMemberType::get(
+      stripBoundDependentMemberTypes(depMemTy->getBase()),
+      depMemTy->getName());
+  }
+
+  return t;
+}
 Type
 GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
   auto genericSig = getGenericSignature();
@@ -248,18 +273,16 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
   // First, write an ErrorType to the location where this type is cached,
   // to catch re-entrant lookups that might arise from an invalid generic
   // signature (eg, <X where X == Array<X>>).
-  ArchetypeType *parentArchetype = nullptr;
+  CanDependentMemberType nestedDependentMemberType;
   GenericTypeParamType *genericParam = nullptr;
   if (auto depMemTy = requirements.anchor->getAs<DependentMemberType>()) {
-    parentArchetype =
-      getOrCreateArchetypeFromInterfaceType(depMemTy->getBase())
-        ->castTo<ArchetypeType>();
-
-    auto name = depMemTy->getName();
-    if (auto type = parentArchetype->getNestedTypeIfKnown(name))
-      return *type;
+    nestedDependentMemberType = cast<DependentMemberType>(
+        stripBoundDependentMemberTypes(depMemTy)->getCanonicalType());
+    auto &entry = getOrCreateNestedTypeStorage()[nestedDependentMemberType];
+    if (entry)
+      return entry;
 
-    parentArchetype->registerNestedType(name, ErrorType::get(ctx));
+    entry = ErrorType::get(ctx);
   } else {
     genericParam = requirements.anchor->castTo<GenericTypeParamType>();
     if (auto type = getMappingIfPresent(genericParam))
@@ -277,12 +300,12 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
 
   Type result;
 
-  if (parentArchetype) {
+  if (nestedDependentMemberType) {
     auto *depMemTy = requirements.anchor->castTo<DependentMemberType>();
     result = NestedArchetypeType::getNew(ctx, depMemTy,
                                          requirements.protos, superclass,
                                          requirements.layout, this);
-    parentArchetype->registerNestedType(depMemTy->getName(), result);
+    getOrCreateNestedTypeStorage()[nestedDependentMemberType] = result;
   } else if (genericParam->isTypeSequence()) {
     result = SequenceArchetypeType::get(ctx, this, genericParam,
                                         requirements.protos, superclass,
@@ -310,21 +333,6 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
   return result;
 }
 
-void ArchetypeType::resolveNestedType(
-                                    std::pair<Identifier, Type> &nested) const {
-  Type interfaceType = getInterfaceType();
-  Type memberInterfaceType =
-      DependentMemberType::get(interfaceType, nested.first);
-
-  Type result = getGenericEnvironment()->getOrCreateArchetypeFromInterfaceType(
-      memberInterfaceType);
-
-  assert(!nested.second ||
-         nested.second->isEqual(result) ||
-         nested.second->is<ErrorType>());
-  nested.second = result;
-}
-
 Type QueryInterfaceTypeSubstitutions::operator()(SubstitutableType *type) const{
   if (auto gp = type->getAs<GenericTypeParamType>()) {
     // Find the index into the parallel arrays of generic parameters and

lib/AST/GenericSignatureBuilder.cpp

@@ -2583,7 +2583,7 @@ static void concretizeNestedTypeFromConcreteParent(
     // Otherwise we have an abstract conformance to an opaque result type.
     assert(conformance.isAbstract());
     auto archetype = concreteParent->castTo<ArchetypeType>();
-    witnessType = archetype->getNestedType(assocType->getName());
+    witnessType = archetype->getNestedType(assocType);
   }
 
   builder.addSameTypeRequirement(

lib/AST/Type.cpp

@@ -3192,7 +3192,6 @@ ArchetypeType::ArchetypeType(TypeKind Kind,
     Environment(Environment)
 {
   // Set up the bits we need for trailing objects to work.
-  Bits.ArchetypeType.ExpandedNestedTypes = false;
   Bits.ArchetypeType.HasSuperclass = static_cast<bool>(Superclass);
   Bits.ArchetypeType.HasLayoutConstraint = static_cast<bool>(Layout);
   Bits.ArchetypeType.NumProtocols = ConformsTo.size();
@@ -3695,78 +3694,25 @@ namespace {
   };
 } // end anonymous namespace
 
-void ArchetypeType::populateNestedTypes() const {
-  if (Bits.ArchetypeType.ExpandedNestedTypes) return;
-
-  // Collect the set of nested types of this archetype.
-  SmallVector<std::pair<Identifier, Type>, 4> nestedTypes;
-  llvm::SmallPtrSet<Identifier, 4> knownNestedTypes;
-  ProtocolType::visitAllProtocols(getConformsTo(),
-                                  [&](ProtocolDecl *proto) -> bool {
-    for (auto assocType : proto->getAssociatedTypeMembers()) {
-      if (knownNestedTypes.insert(assocType->getName()).second)
-        nestedTypes.push_back({ assocType->getName(), Type() });
-    }
-
-    return false;
-  });
-
-  // Record the nested types.
-  auto mutableThis = const_cast<ArchetypeType *>(this);
-
-  std::sort(nestedTypes.begin(), nestedTypes.end(), OrderArchetypeByName());
-  auto &Ctx = mutableThis->getASTContext();
-  mutableThis->NestedTypes = Ctx.AllocateCopy(nestedTypes);
-  mutableThis->Bits.ArchetypeType.ExpandedNestedTypes = true;
+Type ArchetypeType::getNestedType(AssociatedTypeDecl *assocType) {
+  Type interfaceType = getInterfaceType();
+  Type memberInterfaceType =
+      DependentMemberType::get(interfaceType, assocType->getName());
+  return getGenericEnvironment()->getOrCreateArchetypeFromInterfaceType(
+      memberInterfaceType);
 }
 
-Type ArchetypeType::getNestedType(Identifier Name) const {
-  populateNestedTypes();
-
-  auto Pos = std::lower_bound(NestedTypes.begin(), NestedTypes.end(), Name,
-                              OrderArchetypeByName());
-  if (Pos == NestedTypes.end() || Pos->first != Name) {
-    return ErrorType::get(const_cast<ArchetypeType *>(this)->getASTContext());
-  }
-
-  // If the type is null, lazily resolve it. 
-  if (!Pos->second) {
-    resolveNestedType(*Pos);
-  }
-
-  return Pos->second;
-}
-
-Optional<Type> ArchetypeType::getNestedTypeIfKnown(Identifier Name) const {
-  populateNestedTypes();
-
-  auto Pos = std::lower_bound(NestedTypes.begin(), NestedTypes.end(), Name,
-                              OrderArchetypeByName());
-  if (Pos == NestedTypes.end() || Pos->first != Name || !Pos->second)
-    return None;
-
-  return Pos->second;
-}
-
-bool ArchetypeType::hasNestedType(Identifier Name) const {
-  populateNestedTypes();
-
-  auto Pos = std::lower_bound(NestedTypes.begin(), NestedTypes.end(), Name,
-                              OrderArchetypeByName());
-  return Pos != NestedTypes.end() && Pos->first == Name;
-}
+Type ArchetypeType::getNestedTypeByName(Identifier name) {
+  AssociatedTypeDecl *assocType = nullptr;
+  ProtocolType::visitAllProtocols(getConformsTo(), [&](ProtocolDecl *proto) {
+    assocType = proto->getAssociatedType(name);
+    return assocType != nullptr;
+  });
 
-void ArchetypeType::registerNestedType(Identifier name, Type nested) {
-  populateNestedTypes();
+  if (assocType)
+    return getNestedType(assocType);
 
-  auto found = std::lower_bound(NestedTypes.begin(), NestedTypes.end(), name,
-                                OrderArchetypeByName());
-  assert(found != NestedTypes.end() && found->first == name &&
-         "Unable to find nested type?");
-  assert(!found->second ||
-         found->second->isEqual(nested) ||
-         found->second->is<ErrorType>());
-  found->second = nested;
+  return Type();
 }
 
 AssociatedTypeDecl *NestedArchetypeType::getAssocType() const {
@@ -3974,8 +3920,12 @@ static Type getMemberForBaseType(LookupConformanceFn lookupConformances,
   // If the parent is an archetype, extract the child archetype with the
   // given name.
   if (auto archetypeParent = substBase->getAs<ArchetypeType>()) {
-    if (archetypeParent->hasNestedType(name))
-      return archetypeParent->getNestedType(name);
+//    if (assocType && !archetypeParent->getSuperclass())
+//      return archetypeParent->getNestedType(assocType);
+
+    // FIXME: This is really slow for the case where we have an associated type.
+    if (Type memberArchetypeByName = archetypeParent->getNestedTypeByName(name))
+      return memberArchetypeByName;
 
     // If looking for an associated type and the archetype is constrained to a
     // class, continue to the default associated type lookup