[AST] Lazily compute the set of nested types of an ArchetypeType

include/swift/AST/Types.h

@@ -284,6 +284,16 @@ class alignas(1 << TypeAlignInBits) TypeBase {
   enum { NumAnyFunctionTypeBits = NumTypeBaseBits + 16 };
   static_assert(NumAnyFunctionTypeBits <= 32, "fits in an unsigned");
 
+  struct ArchetypeTypeBitfields {
+    unsigned : NumTypeBaseBits;
+
+    unsigned ExpandedNestedTypes : 1;
+    unsigned HasSuperclass : 1;
+    unsigned NumProtocols : 16;
+  };
+  enum { NumArchetypeTypeBitfields = NumTypeBaseBits + 18 };
+  static_assert(NumArchetypeTypeBitfields <= 32, "fits in an unsigned");
+
   struct TypeVariableTypeBitfields {
     unsigned : NumTypeBaseBits;
 
@@ -319,6 +329,7 @@ class alignas(1 << TypeAlignInBits) TypeBase {
     ErrorTypeBitfields ErrorTypeBits;
     AnyFunctionTypeBitfields AnyFunctionTypeBits;
     TypeVariableTypeBitfields TypeVariableTypeBits;
+    ArchetypeTypeBitfields ArchetypeTypeBits;
     SILFunctionTypeBitfields SILFunctionTypeBits;
     AnyMetatypeTypeBitfields AnyMetatypeTypeBits;
   };
@@ -3544,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,
@@ -3728,9 +3749,21 @@ DEFINE_EMPTY_CAN_TYPE_WRAPPER(SubstitutableType, Type)
 /// associated types, as well as the runtime type stored within an
 /// existential container.
 class ArchetypeType final : public SubstitutableType,
-    private llvm::TrailingObjects<ArchetypeType, UUID> {
+    private llvm::TrailingObjects<ArchetypeType, ProtocolDecl *, Type, UUID> {
   friend TrailingObjects;
 
+  size_t numTrailingObjects(OverloadToken<ProtocolDecl *>) const {
+    return ArchetypeTypeBits.NumProtocols;
+  }
+
+  size_t numTrailingObjects(OverloadToken<Type>) const {
+    return ArchetypeTypeBits.HasSuperclass ? 1 : 0;
+  }
+
+  size_t numTrailingObjects(OverloadToken<UUID>) const {
+    return getOpenedExistentialType() ? 1 : 0;
+  }
+
 public:
   /// A nested type. Either a dependent associated archetype, or a concrete
   /// type (which may be a bound archetype from an outer context).
@@ -3774,21 +3807,12 @@ class ArchetypeType final : public SubstitutableType,
   };
 
 private:
-  ArrayRef<ProtocolDecl *> ConformsTo;
-  Type Superclass;
-
   llvm::PointerUnion3<ArchetypeType *, TypeBase *,
                       GenericEnvironment *> ParentOrOpenedOrEnvironment;
   llvm::PointerUnion<AssociatedTypeDecl *, Identifier> AssocTypeOrName;
   MutableArrayRef<std::pair<Identifier, NestedType>> NestedTypes;
 
-  /// Set the ID number of this opened existential.
-  void setOpenedExistentialID(UUID value) {
-    assert(getOpenedExistentialType() && "Not an opened existential archetype");
-    // The UUID is tail-allocated at the end of opened existential archetypes.
-    *getTrailingObjects<UUID>() = value;
-  }
-
+  void populateNestedTypes() const;
   void resolveNestedType(std::pair<Identifier, NestedType> &nested) const;
 
 public:
@@ -3864,15 +3888,22 @@ class ArchetypeType final : public SubstitutableType,
 
   /// getConformsTo - Retrieve the set of protocols to which this substitutable
   /// type shall conform.
-  ArrayRef<ProtocolDecl *> getConformsTo() const { return ConformsTo; }
+  ArrayRef<ProtocolDecl *> getConformsTo() const {
+    return { getTrailingObjects<ProtocolDecl *>(),
+             ArchetypeTypeBits.NumProtocols };
+  }
 
   /// requiresClass - True if the type can only be substituted with class types.
   /// This is true if the type conforms to one or more class protocols or has
   /// a superclass constraint.
   bool requiresClass() const;
 
   /// \brief Retrieve the superclass of this type, if such a requirement exists.
-  Type getSuperclass() const { return Superclass; }
+  Type getSuperclass() const {
+    if (!ArchetypeTypeBits.HasSuperclass) return Type();
+
+    return *getTrailingObjects<Type>();
+  }
 
   /// \brief Return true if the archetype has any requirements at all.
   bool hasRequirements() const {
@@ -3917,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);
@@ -3960,29 +3991,11 @@ class ArchetypeType final : public SubstitutableType,
             ParentOrGenericEnv,
           llvm::PointerUnion<AssociatedTypeDecl *, Identifier> AssocTypeOrName,
           ArrayRef<ProtocolDecl *> ConformsTo,
-          Type Superclass)
-    : SubstitutableType(TypeKind::Archetype, &Ctx,
-                        RecursiveTypeProperties::HasArchetype),
-      ConformsTo(ConformsTo), Superclass(Superclass),
-      AssocTypeOrName(AssocTypeOrName) {
-    if (auto parent = ParentOrGenericEnv.dyn_cast<ArchetypeType *>()) {
-      ParentOrOpenedOrEnvironment = parent;
-    } else {
-      ParentOrOpenedOrEnvironment =
-        ParentOrGenericEnv.get<GenericEnvironment *>();
-    }
-  }
+          Type Superclass);
 
-  ArchetypeType(const ASTContext &Ctx, 
-                Type Existential,
-                ArrayRef<ProtocolDecl *> ConformsTo,
-                Type Superclass)
-    : SubstitutableType(TypeKind::Archetype, &Ctx,
-                        RecursiveTypeProperties(
-                          RecursiveTypeProperties::HasArchetype |
-                          RecursiveTypeProperties::HasOpenedExistential)),
-      ConformsTo(ConformsTo), Superclass(Superclass),
-      ParentOrOpenedOrEnvironment(Existential.getPointer()) { }
+  ArchetypeType(const ASTContext &Ctx, Type Existential,
+                ArrayRef<ProtocolDecl *> ConformsTo, Type Superclass,
+                UUID uuid);
 };
 BEGIN_CAN_TYPE_WRAPPER(ArchetypeType, SubstitutableType)
 CanArchetypeType getParent() const {

lib/AST/ASTContext.cpp

@@ -3418,7 +3418,7 @@ DependentMemberType *DependentMemberType::get(Type base,
 }
 
 CanArchetypeType ArchetypeType::getOpened(Type existential,
-                                        Optional<UUID> knownID) {
+                                          Optional<UUID> knownID) {
   auto &ctx = existential->getASTContext();
   auto &openedExistentialArchetypes = ctx.Impl.OpenedExistentialArchetypes;
   // If we know the ID already...
@@ -3437,19 +3437,20 @@ CanArchetypeType ArchetypeType::getOpened(Type existential,
     knownID = UUID::fromTime();
   }
 
-  auto arena = AllocationArena::Permanent;
   llvm::SmallVector<ProtocolDecl *, 4> conformsTo;
   assert(existential->isExistentialType());
   existential->getAnyExistentialTypeProtocols(conformsTo);
+  Type superclass = existential->getSuperclass(nullptr);
 
-  // Tail-allocate space for the UUID.
-  void *archetypeBuf = ctx.Allocate(totalSizeToAlloc<UUID>(1),
-                                    alignof(ArchetypeType), arena);
-
-  auto result = ::new (archetypeBuf) ArchetypeType(ctx, existential,
-                                       ctx.AllocateCopy(conformsTo),
-                                       existential->getSuperclass(nullptr));
-  result->setOpenedExistentialID(*knownID);
+  auto arena = AllocationArena::Permanent;
+  void *mem = ctx.Allocate(
+                totalSizeToAlloc<ProtocolDecl *, Type, UUID>(conformsTo.size(),
+                                                             superclass ? 1 : 0,
+                                                             1),
+                alignof(ArchetypeType), arena);
+
+  auto result = ::new (mem) ArchetypeType(ctx, existential, conformsTo,
+                                          superclass, *knownID);
   openedExistentialArchetypes[*knownID] = result;
 
   return CanArchetypeType(result);

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);
 }