[AST] Lazily expand the nested types of an ArchetypeType.

An ArchetypeType has all of the information it needs to lazily compute
its set of nested types (hint: they're the associated types of all of
the protocols to which it conforms). Do so only when asked about the
nested types.

Not that there are two levels of laziness happening here: first, we
only create the (sorted) array of nested types when asked about the
nested types the first time. Second, even when we do create it, we
only record the names of these types---the actual archetypes (or
concrete types) behind these names are lazily created when requested.

Of course, all of this laziness is short-lived, because
ArchetypeBuilder::getGenericEnvironment() forces all of the archetypes
to be constructed.

Author: DougGregor

include/swift/AST/Types.h

@@ -287,10 +287,11 @@ class alignas(1 << TypeAlignInBits) TypeBase {
   struct ArchetypeTypeBitfields {
     unsigned : NumTypeBaseBits;
 
+    unsigned ExpandedNestedTypes : 1;
     unsigned HasSuperclass : 1;
     unsigned NumProtocols : 16;
   };
-  enum { NumArchetypeTypeBitfields = NumTypeBaseBits + 17 };
+  enum { NumArchetypeTypeBitfields = NumTypeBaseBits + 18 };
   static_assert(NumArchetypeTypeBitfields <= 32, "fits in an unsigned");
 
   struct TypeVariableTypeBitfields {
@@ -3554,6 +3555,16 @@ class ProtocolType : public NominalType, public llvm::FoldingSetNode {
   /// in the protocol list, then sorting them in some stable order.
   static void canonicalizeProtocols(SmallVectorImpl<ProtocolDecl *> &protocols);
 
+  /// Visit all of the protocols in the given list of protocols, along with their
+  ///
+  /// \param fn Visitor function called for each protocol (just once). If it
+  /// returns \c true, the visit operation will abort and return \c true.
+  ///
+  /// \returns \c true if any invocation of \c fn returns \c true, and \c false
+  /// otherwise.
+  static bool visitAllProtocols(ArrayRef<ProtocolDecl *> protocols,
+                                llvm::function_ref<bool(ProtocolDecl *)> fn);
+
   /// Compare two protocols to provide them with a stable ordering for
   /// use in sorting.
   static int compareProtocols(ProtocolDecl * const* PP1,
@@ -3801,6 +3812,7 @@ class ArchetypeType final : public SubstitutableType,
   llvm::PointerUnion<AssociatedTypeDecl *, Identifier> AssocTypeOrName;
   MutableArrayRef<std::pair<Identifier, NestedType>> NestedTypes;
 
+  void populateNestedTypes() const;
   void resolveNestedType(std::pair<Identifier, NestedType> &nested) const;
 
 public:
@@ -3936,9 +3948,9 @@ class ArchetypeType final : public SubstitutableType,
   getAllNestedTypes(bool resolveTypes = true) const;
 
   /// \brief Set the nested types to a copy of the given array of
-  /// archetypes, which will first be sorted in place.
+  /// archetypes.
   void setNestedTypes(ASTContext &Ctx,
-                      MutableArrayRef<std::pair<Identifier, NestedType>> Nested);
+                      ArrayRef<std::pair<Identifier, NestedType>> Nested);
 
   /// Register a nested type with the given name.
   void registerNestedType(Identifier name, NestedType nested);

lib/AST/ArchetypeBuilder.cpp

@@ -752,30 +752,6 @@ ArchetypeBuilder::PotentialArchetype::getTypeInContext(
     genericEnv->addMapping(getGenericParamKey(), arch);
   }
 
-  // Collect the set of nested types of this archetype, and put them into
-  // the archetype itself.
-  if (!representative->getNestedTypes().empty()) {
-    SmallVector<std::pair<Identifier, NestedType>, 4> FlatNestedTypes;
-    for (auto Nested : representative->getNestedTypes()) {
-      // Skip type aliases, which are just shortcuts.
-      if (Nested.second.front()->getTypeAliasDecl())
-        continue;
-      bool anyNotRenamed = false;
-      for (auto NestedPA : Nested.second) {
-        if (!NestedPA->wasRenamed()) {
-          anyNotRenamed = true;
-          break;
-        }
-      }
-
-      if (!anyNotRenamed)
-        continue;
-
-      FlatNestedTypes.push_back({ Nested.first, NestedType() });
-    }
-    arch->setNestedTypes(ctx, FlatNestedTypes);
-  }
-
   return NestedType::forArchetype(arch);
 }
 

lib/AST/Type.cpp

@@ -1044,6 +1044,37 @@ int ProtocolType::compareProtocols(ProtocolDecl * const* PP1,
   return P1->getName().str().compare(P2->getName().str());
 }
 
+bool ProtocolType::visitAllProtocols(
+                                 ArrayRef<ProtocolDecl *> protocols,
+                                 llvm::function_ref<bool(ProtocolDecl *)> fn) {
+  SmallVector<ProtocolDecl *, 4> stack;
+  SmallPtrSet<ProtocolDecl *, 4> knownProtocols;
+
+  // Prepopulate the stack.
+  for (auto proto : protocols) {
+    if (knownProtocols.insert(proto).second)
+      stack.push_back(proto);
+  }
+  std::reverse(stack.begin(), stack.end());
+
+  while (!stack.empty()) {
+    auto proto = stack.back();
+    stack.pop_back();
+
+    // Visit this protocol.
+    if (fn(proto))
+      return true;
+
+    // Add inherited protocols that we haven't seen already.
+    for (auto inherited : proto->getInheritedProtocols(nullptr)) {
+      if (knownProtocols.insert(inherited).second)
+        stack.push_back(inherited);
+    }
+  }
+
+  return false;
+}
+
 void ProtocolType::canonicalizeProtocols(
        SmallVectorImpl<ProtocolDecl *> &protocols) {
   llvm::SmallDenseMap<ProtocolDecl *, unsigned> known;
@@ -2494,6 +2525,7 @@ ArchetypeType::ArchetypeType(
   }
 
   // Set up the bits we need for trailing objects to work.
+  ArchetypeTypeBits.ExpandedNestedTypes = false;
   ArchetypeTypeBits.HasSuperclass = static_cast<bool>(Superclass);
   ArchetypeTypeBits.NumProtocols = ConformsTo.size();
 
@@ -2515,6 +2547,7 @@ ArchetypeType::ArchetypeType(const ASTContext &Ctx, Type Existential,
                         RecursiveTypeProperties::HasOpenedExistential)),
     ParentOrOpenedOrEnvironment(Existential.getPointer()) {
   // Set up the bits we need for trailing objects to work.
+  ArchetypeTypeBits.ExpandedNestedTypes = false;
   ArchetypeTypeBits.HasSuperclass = static_cast<bool>(Superclass);
   ArchetypeTypeBits.NumProtocols = ConformsTo.size();
 
@@ -2604,7 +2637,32 @@ namespace {
   };
 }
 
+void ArchetypeType::populateNestedTypes() const {
+  if (ArchetypeTypeBits.ExpandedNestedTypes) return;
+
+  // Collect the set of nested types of this archetype.
+  SmallVector<std::pair<Identifier, NestedType>, 4> nestedTypes;
+  llvm::SmallPtrSet<Identifier, 4> knownNestedTypes;
+  ProtocolType::visitAllProtocols(getConformsTo(),
+                                  [&](ProtocolDecl *proto) -> bool {
+    for (auto member : proto->getMembers()) {
+      if (auto assocType = dyn_cast<AssociatedTypeDecl>(member)) {
+        if (knownNestedTypes.insert(assocType->getName()).second)
+          nestedTypes.push_back({ assocType->getName(), NestedType() });
+      }
+    }
+
+    return false;
+  });
+
+  // Record the nested types.
+  auto mutableThis = const_cast<ArchetypeType *>(this);
+  mutableThis->setNestedTypes(mutableThis->getASTContext(), nestedTypes);
+}
+
 ArchetypeType::NestedType ArchetypeType::getNestedType(Identifier Name) const {
+  populateNestedTypes();
+
   auto Pos = std::lower_bound(NestedTypes.begin(), NestedTypes.end(), Name,
                               OrderArchetypeByName());
   if (Pos == NestedTypes.end() || Pos->first != Name) {
@@ -2623,6 +2681,8 @@ ArchetypeType::NestedType ArchetypeType::getNestedType(Identifier Name) const {
 
 Optional<ArchetypeType::NestedType> 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)
@@ -2632,13 +2692,17 @@ Optional<ArchetypeType::NestedType> ArchetypeType::getNestedTypeIfKnown(
 }
 
 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;
 }
 
 ArrayRef<std::pair<Identifier, ArchetypeType::NestedType>>
 ArchetypeType::getAllNestedTypes(bool resolveTypes) const {
+  populateNestedTypes();
+
   if (resolveTypes) {
     for (auto &nested : NestedTypes) {
       if (!nested.second)
@@ -2651,12 +2715,16 @@ ArchetypeType::getAllNestedTypes(bool resolveTypes) const {
 
 void ArchetypeType::setNestedTypes(
        ASTContext &Ctx,
-       MutableArrayRef<std::pair<Identifier, NestedType>> Nested) {
-  std::sort(Nested.begin(), Nested.end(), OrderArchetypeByName());
+       ArrayRef<std::pair<Identifier, NestedType>> Nested) {
+  assert(!ArchetypeTypeBits.ExpandedNestedTypes && "Already expanded");
   NestedTypes = Ctx.AllocateCopy(Nested);
+  std::sort(NestedTypes.begin(), NestedTypes.end(), OrderArchetypeByName());
+  ArchetypeTypeBits.ExpandedNestedTypes = true;
 }
 
 void ArchetypeType::registerNestedType(Identifier name, NestedType nested) {
+  populateNestedTypes();
+
   auto found = std::lower_bound(NestedTypes.begin(), NestedTypes.end(), name,
                                 OrderArchetypeByName());
   assert(found != NestedTypes.end() && found->first == name &&