[Archetype builder] Make archetype construction more lazy

include/swift/AST/ASTContext.h

@@ -762,8 +762,13 @@ class ASTContext {
 
   /// Retrieve or create the stored archetype builder for the given
   /// canonical generic signature and module.
-  std::pair<ArchetypeBuilder *, GenericEnvironment *>
-  getOrCreateArchetypeBuilder(CanGenericSignature sig, ModuleDecl *mod);
+  ArchetypeBuilder *getOrCreateArchetypeBuilder(CanGenericSignature sig,
+                                                ModuleDecl *mod);
+
+  /// Retrieve or create the canonical generic environment of a canonical
+  /// archetype builder.
+  GenericEnvironment *getOrCreateCanonicalGenericEnvironment(
+                                                     ArchetypeBuilder *builder);
 
   /// Retrieve the inherited name set for the given class.
   const InheritedNameSet *getAllPropertyNames(ClassDecl *classDecl,
@@ -796,26 +801,6 @@ class ASTContext {
                         DeclContext *gpContext,
                         ArrayRef<Substitution> Subs) const;
 
-  /// Retrieve the archetype builder and potential archetype
-  /// corresponding to the given archetype type.
-  ///
-  /// This facility is only used by the archetype builder when forming
-  /// archetypes.a
-  std::pair<ArchetypeBuilder *, ArchetypeBuilder::PotentialArchetype *>
-  getLazyArchetype(const ArchetypeType *archetype);
-
-  /// Register information for a lazily-constructed archetype.
-  void registerLazyArchetype(
-         const ArchetypeType *archetype,
-         ArchetypeBuilder &builder,
-         ArchetypeBuilder::PotentialArchetype *potentialArchetype);
-
-  /// Unregister information about the given lazily-constructed archetype.
-  void unregisterLazyArchetype(const ArchetypeType *archetype);
-
-  friend class ArchetypeType;
-  friend class ArchetypeBuilder::PotentialArchetype;
-
   /// Provide context-level uniquing for SIL lowered type layouts.
   friend SILLayout;
   llvm::FoldingSet<SILLayout> *&getSILLayouts();

include/swift/AST/GenericEnvironment.h

@@ -20,6 +20,8 @@
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/AST/GenericParamKey.h"
 #include "swift/AST/GenericSignature.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/TrailingObjects.h"
 
 namespace swift {
 
@@ -31,14 +33,57 @@ class SILType;
 
 /// Describes the mapping between archetypes and interface types for the
 /// generic parameters of a DeclContext.
-class alignas(1 << DeclAlignInBits) GenericEnvironment final {
+class alignas(1 << DeclAlignInBits) GenericEnvironment final
+                  : private llvm::TrailingObjects<GenericEnvironment, Type> {
   GenericSignature *Signature;
+  ArchetypeBuilder *Builder;
   TypeSubstitutionMap ArchetypeToInterfaceMap;
-  TypeSubstitutionMap InterfaceToArchetypeMap;
+
+  friend TrailingObjects;
+
+  size_t numTrailingObjects(OverloadToken<Type>) const {
+    return Signature->getGenericParams().size();
+  }
+
+  /// Retrieve the array containing the context types associated with the
+  /// generic parameters, stored in parallel with the generic parameters of the
+  /// generic signature.
+  MutableArrayRef<Type> getContextTypes() {
+    return MutableArrayRef<Type>(getTrailingObjects<Type>(),
+                                 Signature->getGenericParams().size());
+  }
+
+  /// Retrieve the array containing the context types associated with the
+  /// generic parameters, stored in parallel with the generic parameters of the
+  /// generic signature.
+  ArrayRef<Type> getContextTypes() const {
+    return ArrayRef<Type>(getTrailingObjects<Type>(),
+                          Signature->getGenericParams().size());
+  }
 
   GenericEnvironment(GenericSignature *signature,
+                     ArchetypeBuilder *builder,
                      TypeSubstitutionMap interfaceToArchetypeMap);
 
+  friend class ArchetypeType;
+  friend class ArchetypeBuilder;
+
+  ArchetypeBuilder *getArchetypeBuilder() const { return Builder; }
+  void clearArchetypeBuilder() { Builder = nullptr; }
+
+  /// Query function suitable for use as a \c TypeSubstitutionFn that queries
+  /// the mapping of interface types to archetypes.
+  class QueryInterfaceTypeSubstitutions {
+    const GenericEnvironment *self;
+
+  public:
+    QueryInterfaceTypeSubstitutions(const GenericEnvironment *self)
+      : self(self) { }
+
+    Type operator()(SubstitutableType *type) const;
+  };
+  friend class QueryInterfaceTypeSubstitutions;
+
 public:
   GenericSignature *getGenericSignature() const {
     return Signature;
@@ -53,15 +98,14 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final {
   bool containsPrimaryArchetype(ArchetypeType *archetype) const;
 
   static
-  GenericEnvironment *get(ASTContext &ctx,
-                          GenericSignature *signature,
+  GenericEnvironment *get(GenericSignature *signature,
                           TypeSubstitutionMap interfaceToArchetypeMap);
 
   /// Create a new, "incomplete" generic environment that will be populated
   /// by calls to \c addMapping().
   static
-  GenericEnvironment *getIncomplete(ASTContext &ctx,
-                                    GenericSignature *signature);
+  GenericEnvironment *getIncomplete(GenericSignature *signature,
+                                    ArchetypeBuilder *builder);
 
   /// Add a mapping of a generic parameter to a specific type (which may be
   /// an archetype)
@@ -76,9 +120,11 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final {
   void *operator new(size_t Bytes) = delete;
   void operator delete(void *Data) = delete;
 
-  /// Only allow allocation of GenericEnvironments using the allocator
-  /// in ASTContext.
-  void *operator new(size_t bytes, const ASTContext &ctx);
+  /// Only allow placement new.
+  void *operator new(size_t Bytes, void *Mem) {
+    assert(Mem); 
+    return Mem; 
+  }
 
   /// Map a contextual type to an interface type.
   Type mapTypeOutOfContext(ModuleDecl *M, Type type) const;

include/swift/AST/GenericSignature.h

@@ -77,6 +77,8 @@ class alignas(1 << TypeAlignInBits) GenericSignature final
   /// Retrieve the archetype builder for the given generic signature.
   ArchetypeBuilder *getArchetypeBuilder(ModuleDecl &mod);
 
+  friend class ArchetypeType;
+
 public:
   /// Create a new generic signature with the given type parameters and
   /// requirements.
@@ -129,6 +131,13 @@ class alignas(1 << TypeAlignInBits) GenericSignature final
   using LookupConformanceFn =
       llvm::function_ref<ProtocolConformanceRef(CanType, Type, ProtocolType *)>;
 
+  /// Build an array of substitutions from an interface type substitution map,
+  /// using the given function to look up conformances.
+  void getSubstitutions(ModuleDecl &mod,
+                        TypeSubstitutionFn substitution,
+                        LookupConformanceFn lookupConformance,
+                        SmallVectorImpl<Substitution> &result) const;
+
   /// Build an array of substitutions from an interface type substitution map,
   /// using the given function to look up conformances.
   void getSubstitutions(ModuleDecl &mod,

include/swift/AST/Type.h

@@ -19,6 +19,7 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/STLExtras.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/AST/PrintOptions.h"
 #include "swift/AST/TypeAlignments.h"
@@ -52,6 +53,20 @@ class TypeWalker;
 /// replacements.
 typedef llvm::DenseMap<SubstitutableType *, Type> TypeSubstitutionMap;
 
+/// Function used to provide substitutions.
+///
+/// \returns A null \c Type to indicate that there is no substitution for
+/// this substitutable type; otherwise, the replacement type.
+typedef llvm::function_ref<Type(SubstitutableType *)> TypeSubstitutionFn;
+
+/// A function object suitable for use as a \c TypeSubstitutionFn that
+/// queries an underlying \c TypeSubstitutionMap.
+struct QueryTypeSubstitutionMap {
+  const TypeSubstitutionMap &substitutions;
+
+  Type operator()(SubstitutableType *type) const;
+};
+
 /// Flags that can be passed when substituting into a type.
 enum class SubstFlags {
   /// If a type cannot be produced because some member type is
@@ -187,6 +202,21 @@ class Type {
   Type subst(const SubstitutionMap &substitutions,
              SubstOptions options = None) const;
 
+  /// Replace references to substitutable types with new, concrete types and
+  /// return the substituted result.
+  ///
+  /// \param module The module to use for conformance lookups.
+  ///
+  /// \param substitutions A function mapping from substitutable types to their
+  /// replacements.
+  ///
+  /// \param options Options that affect the substitutions.
+  ///
+  /// \returns the substituted type, or a null type if an error occurred.
+  Type subst(ModuleDecl *module,
+             TypeSubstitutionFn substitutions,
+             SubstOptions options = None) const;
+
   bool isPrivateStdlibType(bool whitelistProtocols=true) const;
 
   void dump() const;

include/swift/AST/Types.h

@@ -3874,12 +3874,6 @@ class ArchetypeType final : public SubstitutableType,
   /// \brief Retrieve the superclass of this type, if such a requirement exists.
   Type getSuperclass() const { return Superclass; }
 
-  /// \brief Set the superclass of this type.
-  ///
-  /// This can only be performed in very narrow cases where the archetype is
-  /// being lazily constructed.
-  void setSuperclass(Type superclass) { Superclass = superclass; }
-
   /// \brief Return true if the archetype has any requirements at all.
   bool hasRequirements() const {
     return !getConformsTo().empty() || getSuperclass();

lib/AST/ASTContext.cpp

@@ -207,19 +207,16 @@ struct ASTContext::Implementation {
   llvm::DenseMap<Decl *, std::pair<LazyMemberLoader *, uint64_t>>
     ConformanceLoaders;
 
-  /// Mapping from archetypes with lazily-resolved nested types to the
-  /// archetype builder and potential archetype corresponding to that
-  /// archetype.
-  llvm::DenseMap<const ArchetypeType *, 
-                 std::pair<ArchetypeBuilder *,
-                           ArchetypeBuilder::PotentialArchetype *>>
-    LazyArchetypes;
-
-  /// \brief Stored archetype builders and their corresponding (canonical)
-  /// generic environments.
+  /// Stored archetype builders for canonical generic signatures.
   llvm::DenseMap<std::pair<GenericSignature *, ModuleDecl *>,
-  std::pair<std::unique_ptr<ArchetypeBuilder>,
-            GenericEnvironment *>> ArchetypeBuilders;
+                 std::unique_ptr<ArchetypeBuilder>>
+    ArchetypeBuilders;
+
+  /// Canonical generic environments for canonical generic signatures.
+  ///
+  /// The keys are the archetype builders in \c ArchetypeBuilders.
+  llvm::DenseMap<ArchetypeBuilder *, GenericEnvironment *>
+    CanonicalGenericEnvironments;
 
   /// The set of property names that show up in the defining module of a
   /// class.
@@ -1254,25 +1251,36 @@ void ASTContext::getVisibleTopLevelClangModules(
     collectAllModules(Modules);
 }
 
-std::pair<ArchetypeBuilder *, GenericEnvironment *>
-ASTContext::getOrCreateArchetypeBuilder(CanGenericSignature sig,
-                                        ModuleDecl *mod) {
+ArchetypeBuilder *ASTContext::getOrCreateArchetypeBuilder(
+                                                      CanGenericSignature sig,
+                                                      ModuleDecl *mod) {
   // Check whether we already have an archetype builder for this
   // signature and module.
   auto known = Impl.ArchetypeBuilders.find({sig, mod});
   if (known != Impl.ArchetypeBuilders.end())
-    return { known->second.first.get(), known->second.second };
+    return known->second.get();
 
   // Create a new archetype builder with the given signature.
   auto builder = new ArchetypeBuilder(*mod);
   builder->addGenericSignature(sig, nullptr);
-
-  // Store this archetype builder and its generic environment.
-  auto genericEnv = builder->getGenericEnvironment(sig);
-  Impl.ArchetypeBuilders[{sig, mod}]
-    = { std::unique_ptr<ArchetypeBuilder>(builder), genericEnv };
 
-  return { builder, genericEnv };
+  // Store this archetype builder (no generic environment yet).
+  Impl.ArchetypeBuilders[{sig, mod}] =
+    std::unique_ptr<ArchetypeBuilder>(builder);
+
+  return builder;
+}
+
+GenericEnvironment *ASTContext::getOrCreateCanonicalGenericEnvironment(
+                                                    ArchetypeBuilder *builder) {
+  auto known = Impl.CanonicalGenericEnvironments.find(builder);
+  if (known != Impl.CanonicalGenericEnvironments.find(builder))
+    return known->second;
+
+  auto sig = builder->getGenericSignature();
+  auto env = builder->getGenericEnvironment(sig);
+  Impl.CanonicalGenericEnvironments[builder] = env;
+  return env;
 }
 
 Module *
@@ -3551,22 +3559,30 @@ GenericSignature *GenericSignature::get(ArrayRef<GenericTypeParamType *> params,
 }
 
 GenericEnvironment *
-GenericEnvironment::get(ASTContext &ctx,
-                        GenericSignature *signature,
+GenericEnvironment::get(GenericSignature *signature,
                         TypeSubstitutionMap interfaceToArchetypeMap) {
   assert(!interfaceToArchetypeMap.empty());
   assert(interfaceToArchetypeMap.size() == signature->getGenericParams().size()
          && "incorrect number of parameters");
 
+  ASTContext &ctx = signature->getASTContext();
 
-  return new (ctx) GenericEnvironment(signature, interfaceToArchetypeMap);
+  // Allocate and construct the new environment.
+  size_t bytes = totalSizeToAlloc<Type>(signature->getGenericParams().size());
+  void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
+  return new (mem) GenericEnvironment(signature, nullptr,
+                                      interfaceToArchetypeMap);
 }
 
 GenericEnvironment *GenericEnvironment::getIncomplete(
-                                                  ASTContext &ctx,
-                                                  GenericSignature *signature) {
+                                                  GenericSignature *signature,
+                                                  ArchetypeBuilder *builder) {
   TypeSubstitutionMap empty;
-  return new (ctx) GenericEnvironment(signature, empty);
+  auto &ctx = signature->getASTContext();
+  // Allocate and construct the new environment.
+  size_t bytes = totalSizeToAlloc<Type>(signature->getGenericParams().size());
+  void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
+  return new (mem) GenericEnvironment(signature, builder, empty);
 }
 
 void DeclName::CompoundDeclName::Profile(llvm::FoldingSetNodeID &id,
@@ -3886,27 +3902,6 @@ Type ASTContext::getBridgedToObjC(const DeclContext *dc, Type type,
   return Type();
 }
 
-std::pair<ArchetypeBuilder *, ArchetypeBuilder::PotentialArchetype *>
-ASTContext::getLazyArchetype(const ArchetypeType *archetype) {
-  auto known = Impl.LazyArchetypes.find(archetype);
-  assert(known != Impl.LazyArchetypes.end());
-  return known->second;
-}
-
-void ASTContext::registerLazyArchetype(
-       const ArchetypeType *archetype,
-       ArchetypeBuilder &builder,
-       ArchetypeBuilder::PotentialArchetype *potentialArchetype) {
-  assert(Impl.LazyArchetypes.count(archetype) == 0);
-  Impl.LazyArchetypes[archetype] = { &builder, potentialArchetype };
-}
-
-void ASTContext::unregisterLazyArchetype(const ArchetypeType *archetype) {
-  auto known = Impl.LazyArchetypes.find(archetype);
-  assert(known != Impl.LazyArchetypes.end());
-  Impl.LazyArchetypes.erase(known);
-}
-
 const InheritedNameSet *ASTContext::getAllPropertyNames(ClassDecl *classDecl,
                                                         bool forInstance) {
   // If this class was defined in Objective-C, perform the lookup based on