[AST] Stash an archetype builder in an "incomplete" generic environment.

This eliminates the really gross registration of archetype builders
within the ASTContext, and is another little step toward lazily
constructing archetypes.

include/swift/AST/ASTContext.h

@@ -762,8 +762,13 @@ public:
 
   /// Retrieve or create the stored archetype builder for the given
   /// canonical generic signature and module.
-  std::pair<ArchetypeBuilder *, GenericEnvironment *>
+  ArchetypeBuilder *getOrCreateArchetypeBuilder(CanGenericSignature sig,
-  getOrCreateArchetypeBuilder(CanGenericSignature sig, ModuleDecl *mod);
+                                                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 @@ private:
                         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

@@ -33,12 +33,20 @@ class SILType;
 /// generic parameters of a DeclContext.
 class alignas(1 << DeclAlignInBits) GenericEnvironment final {
   GenericSignature *Signature;
+  ArchetypeBuilder *Builder;
   TypeSubstitutionMap ArchetypeToInterfaceMap;
   TypeSubstitutionMap InterfaceToArchetypeMap;
 
   GenericEnvironment(GenericSignature *signature,
+                     ArchetypeBuilder *builder,
                      TypeSubstitutionMap interfaceToArchetypeMap);
 
+  friend class ArchetypeType;
+  friend class ArchetypeBuilder;
+  
+  ArchetypeBuilder *getArchetypeBuilder() const { return Builder; }
+  void clearArchetypeBuilder() { Builder = nullptr; }
+  
 public:
   GenericSignature *getGenericSignature() const {
     return Signature;
@@ -53,15 +61,14 @@ public:
   bool containsPrimaryArchetype(ArchetypeType *archetype) const;
 
   static
-  GenericEnvironment *get(ASTContext &ctx,
+  GenericEnvironment *get(GenericSignature *signature,
-                          GenericSignature *signature,
                           TypeSubstitutionMap interfaceToArchetypeMap);
 
   /// Create a new, "incomplete" generic environment that will be populated
   /// by calls to \c addMapping().
   static
-  GenericEnvironment *getIncomplete(ASTContext &ctx,
+  GenericEnvironment *getIncomplete(GenericSignature *signature,
-                                    GenericSignature *signature);
+                                    ArchetypeBuilder *builder);
 
   /// Add a mapping of a generic parameter to a specific type (which may be
   /// an archetype)

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.

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
+  /// Stored archetype builders for canonical generic signatures.
-  /// 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.
   llvm::DenseMap<std::pair<GenericSignature *, ModuleDecl *>,
-  std::pair<std::unique_ptr<ArchetypeBuilder>,
+                 std::unique_ptr<ArchetypeBuilder>>
-            GenericEnvironment *>> ArchetypeBuilders;
+    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 *>
+ArchetypeBuilder *ASTContext::getOrCreateArchetypeBuilder(
-ASTContext::getOrCreateArchetypeBuilder(CanGenericSignature sig,
+                                                      CanGenericSignature sig,
-                                        ModuleDecl *mod) {
+                                                      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,24 @@ GenericSignature *GenericSignature::get(ArrayRef<GenericTypeParamType *> params,
 }
 
 GenericEnvironment *
-GenericEnvironment::get(ASTContext &ctx,
+GenericEnvironment::get(GenericSignature *signature,
-                        GenericSignature *signature,
                         TypeSubstitutionMap interfaceToArchetypeMap) {
   assert(!interfaceToArchetypeMap.empty());
   assert(interfaceToArchetypeMap.size() == signature->getGenericParams().size()
          && "incorrect number of parameters");
 
 
-  return new (ctx) GenericEnvironment(signature, interfaceToArchetypeMap);
+  ASTContext &ctx = signature->getASTContext();
+  return new (ctx) 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();
+  return new (ctx) GenericEnvironment(signature, builder, empty);
 }
 
 void DeclName::CompoundDeclName::Profile(llvm::FoldingSetNodeID &id,
@@ -3886,27 +3896,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

lib/AST/ArchetypeBuilder.cpp

@@ -767,7 +767,6 @@ ArchetypeBuilder::PotentialArchetype::getTypeInContext(
   // Collect the set of nested types of this archetype, and put them into
   // the archetype itself.
   if (!representative->getNestedTypes().empty()) {
-    ctx.registerLazyArchetype(arch, builder, this);
     SmallVector<std::pair<Identifier, NestedType>, 4> FlatNestedTypes;
     for (auto Nested : representative->getNestedTypes()) {
       // Skip type aliases, which are just shortcuts.
@@ -790,8 +789,6 @@ ArchetypeBuilder::PotentialArchetype::getTypeInContext(
 
     // Force the resolution of the nested types.
     (void)arch->getAllNestedTypes();
-
-    ctx.unregisterLazyArchetype(arch);
   }
 
   return NestedType::forArchetype(arch);
@@ -799,15 +796,15 @@ ArchetypeBuilder::PotentialArchetype::getTypeInContext(
 
 void ArchetypeType::resolveNestedType(
        std::pair<Identifier, NestedType> &nested) const {
-  auto &ctx = const_cast<ArchetypeType *>(this)->getASTContext();
-  auto lazyArchetype = ctx.getLazyArchetype(this);
-
-  ArchetypeBuilder &builder = *lazyArchetype.first;
   auto genericEnv = getGenericEnvironment();
-  auto potentialArchetype =
+  auto &builder = *genericEnv->getArchetypeBuilder();
-    lazyArchetype.second->getNestedType(nested.first, builder);
 
-  auto result = potentialArchetype->getTypeInContext(builder, genericEnv);
+  Type interfaceType =
+    genericEnv->mapTypeOutOfContext(&builder.getModule(),
+                                    const_cast<ArchetypeType *>(this));
+  auto parentPA = builder.resolveArchetype(interfaceType);
+  auto memberPA = parentPA->getNestedType(nested.first, builder);
+  auto result = memberPA->getTypeInContext(builder, genericEnv);
   assert(!nested.second ||
          nested.second.getValue()->isEqual(result.getValue()) ||
          (nested.second.getValue()->hasError() &&
@@ -2235,11 +2232,10 @@ GenericSignature *ArchetypeBuilder::getGenericSignature() {
   return sig;
 }
 
-GenericEnvironment *ArchetypeBuilder::getGenericEnvironment(GenericSignature *signature) {
+GenericEnvironment *ArchetypeBuilder::getGenericEnvironment(
-  TypeSubstitutionMap interfaceToArchetypeMap;
+                                                  GenericSignature *signature) {
-
   // Compute the archetypes for the generic parameters.
-  auto genericEnv = GenericEnvironment::getIncomplete(Context, signature);
+  auto genericEnv = GenericEnvironment::getIncomplete(signature, this);
   for (auto pa : Impl->PotentialArchetypes) {
     Type contextType = pa->getTypeInContext(*this, genericEnv).getValue();
     if (!genericEnv->getMappingIfPresent(pa->getGenericParamKey()))
@@ -2268,6 +2264,7 @@ GenericEnvironment *ArchetypeBuilder::getGenericEnvironment(GenericSignature *si
   }
 #endif
 
+  genericEnv->clearArchetypeBuilder();
   return genericEnv;
 }
 

lib/AST/Builtins.cpp

@@ -198,8 +198,7 @@ getBuiltinGenericFunction(Identifier Id,
   GenericSignature *Sig =
     GenericSignature::get(GenericParamTypes, { });
   GenericEnvironment *Env =
-      GenericEnvironment::get(Context, Sig,
+    GenericEnvironment::get(Sig, InterfaceToArchetypeMap);
-                              InterfaceToArchetypeMap);
 
   Type InterfaceType = GenericFunctionType::get(Sig, ArgParamType, ResType,
                                                 AnyFunctionType::ExtInfo());

lib/AST/GenericEnvironment.cpp

@@ -22,8 +22,9 @@ using namespace swift;
 
 GenericEnvironment::GenericEnvironment(
     GenericSignature *signature,
+    ArchetypeBuilder *builder,
     TypeSubstitutionMap interfaceToArchetypeMap)
-  : Signature(signature)
+  : Signature(signature), Builder(builder)
 {
   // Build a mapping in both directions, making sure to canonicalize the
   // interface type where it is used as a key, so that substitution can

lib/AST/GenericSignature.cpp

@@ -75,7 +75,7 @@ ArchetypeBuilder *GenericSignature::getArchetypeBuilder(ModuleDecl &mod) {
 
   // Archetype builders are stored on the ASTContext.
   return getASTContext().getOrCreateArchetypeBuilder(CanGenericSignature(this),
-                                                     &mod).first;
+                                                     &mod);
 }
 
 bool GenericSignature::isCanonical() const {
@@ -522,8 +522,8 @@ CanType GenericSignature::getCanonicalTypeInContext(Type type, ModuleDecl &mod)
 GenericEnvironment *CanGenericSignature::getGenericEnvironment(
                                                      ModuleDecl &module) const {
   // Archetype builders are stored on the ASTContext.
-  return module.getASTContext().getOrCreateArchetypeBuilder(*this, &module)
+  return module.getASTContext().getOrCreateCanonicalGenericEnvironment(
-           .second;
+           module.getASTContext().getOrCreateArchetypeBuilder(*this, &module));
 }
 
 unsigned GenericParamKey::findIndexIn(

lib/Serialization/Deserialization.cpp

@@ -952,8 +952,7 @@ GenericEnvironment *ModuleFile::readGenericEnvironment(
   assert(!interfaceToArchetypeMap.empty() &&
          "no archetypes in generic function?");
 
-  return GenericEnvironment::get(getContext(), signature,
+  return GenericEnvironment::get(signature, interfaceToArchetypeMap);
-                                 interfaceToArchetypeMap);
 }
 
 GenericEnvironment *ModuleFile::maybeReadGenericEnvironment() {