[Serialization/AST] Lazily construct generic environments for generic types.

When deserializing the generic environment for a generic type, only
immediately deserialize the generic signature. The generic environment
will be deserialized later, when it's needed.

include/swift/AST/ASTContext.h

@@ -703,11 +703,20 @@ public:
   LazyContextData *getOrCreateLazyContextData(const Decl *decl,
                                               LazyMemberLoader *lazyLoader);
 
+  /// Get the lazy function data for the given generic type.
+  ///
+  /// \param lazyLoader If non-null, the lazy loader to use when creating the
+  /// generic type data. The pointer must either be null or be consistent
+  /// across all calls for the same \p type.
+  LazyGenericTypeData *getOrCreateLazyGenericTypeData(
+                                                  const GenericTypeDecl *type,
+                                                  LazyMemberLoader *lazyLoader);
+
   /// Get the lazy function data for the given abstract function.
   ///
   /// \param lazyLoader If non-null, the lazy loader to use when creating the
   /// function data. The pointer must either be null or be consistent
-  /// across all calls for the same \p idc.
+  /// across all calls for the same \p func.
   LazyAbstractFunctionData *getOrCreateLazyFunctionContextData(
                                               const AbstractFunctionDecl *func,
                                               LazyMemberLoader *lazyLoader);

include/swift/AST/Decl.h

@@ -2241,17 +2241,21 @@ public:
 class GenericTypeDecl : public TypeDecl, public DeclContext {
   GenericParamList *GenericParams = nullptr;
 
-  /// \brief The generic context of this type.
+  /// The generic signature or environment of this type.
   ///
-  /// This is the mapping between interface types and archetypes for the
+  /// When this function stores only a signature, the generic environment
-  /// generic parameters of this type.
+  /// will be lazily loaded.
-  GenericEnvironment *GenericEnv = nullptr;
+  mutable llvm::PointerUnion<GenericSignature *, GenericEnvironment *>
+    GenericSigOrEnv;
 
   /// \brief Whether or not the generic signature of the type declaration is
   /// currently being validated.
   // TODO: Merge into GenericSig bits.
   unsigned ValidatingGenericSignature = false;
 
+  /// Lazily populate the generic environment.
+  GenericEnvironment *getLazyGenericEnvironmentSlow() const;
+
 public:
   GenericTypeDecl(DeclKind K, DeclContext *DC,
                   Identifier name, SourceLoc nameLoc,
@@ -2280,11 +2284,30 @@ public:
       return { };
   }
 
-  /// Retrieve the generic signature.
+  /// Retrieve the generic signature for this type.
   GenericSignature *getGenericSignature() const {
-    return GenericEnv ? GenericEnv->getGenericSignature() : nullptr;
+    if (auto genericEnv = GenericSigOrEnv.dyn_cast<GenericEnvironment *>())
+      return genericEnv->getGenericSignature();
+
+    if (auto genericSig = GenericSigOrEnv.dyn_cast<GenericSignature *>())
+      return genericSig;
+
+    return nullptr;
   }
-  
+
+  /// Retrieve the generic context for this type.
+  GenericEnvironment *getGenericEnvironment() const {
+    // Fast case: we already have a generic environment.
+    if (auto genericEnv = GenericSigOrEnv.dyn_cast<GenericEnvironment *>())
+      return genericEnv;
+
+    // If we only have a generic signature, build the generic environment.
+    if (GenericSigOrEnv.dyn_cast<GenericSignature *>())
+      return getLazyGenericEnvironmentSlow();
+
+    return nullptr;
+  }
+
   void setIsValidatingGenericSignature(bool validating=true) {
     ValidatingGenericSignature = validating;
   }
@@ -2293,16 +2316,21 @@ public:
     return ValidatingGenericSignature;
   }
 
-  /// Retrieve the generic context for this type.
+  /// Set a lazy generic environment.
-  GenericEnvironment *getGenericEnvironment() const { return GenericEnv; }
+  void setLazyGenericEnvironment(LazyMemberLoader *lazyLoader,
+                                 GenericSignature *genericSig,
+                                 uint64_t genericEnvData);
 
-  /// Set the generic context of this type.
+  /// Set the generic context of this function.
-  void setGenericEnvironment(GenericEnvironment *env) {
+  void setGenericEnvironment(GenericEnvironment *genericEnv) {
-    assert(!this->GenericEnv && "already have generic context?");
+    assert((GenericSigOrEnv.isNull() ||
-    this->GenericEnv = env;
+            getGenericSignature()->getCanonicalSignature() ==
+              genericEnv->getGenericSignature()->getCanonicalSignature()) &&
+           "set a generic environment with a different generic signature");
+    this->GenericSigOrEnv = genericEnv;
 
-    if (GenericEnv)
+    if (genericEnv)
-      GenericEnv->setOwningDeclContext(this);
+      genericEnv->setOwningDeclContext(this);
   }
 
   // Resolve ambiguity due to multiple base classes.

include/swift/AST/LazyResolver.h

@@ -213,8 +213,15 @@ public:
   uint64_t genericEnvData = 0;
 };
 
+/// Context data for generic type declarations.
+class LazyGenericTypeData : public LazyContextData {
+public:
+  /// The context data used for loading the generic environment.
+  uint64_t genericEnvData = 0;
+};
+
 /// Context data for iterable decl contexts.
-class LazyIterableDeclContextData : public LazyContextData {
+class LazyIterableDeclContextData : public LazyGenericTypeData {
 public:
   /// The context data used for loading all of the members of the iterable
   /// context.

include/swift/Serialization/ModuleFile.h

@@ -449,6 +449,13 @@ private:
   void readGenericRequirements(SmallVectorImpl<Requirement> &requirements,
                                llvm::BitstreamCursor &Cursor);
 
+  /// Allocate a lazy generic environment map for use with lazily deserialized
+  /// generic environments.
+  uint64_t allocateLazyGenericEnvironmentMap(TypeSubstitutionMap &&map);
+
+  /// Set up a lazy generic environment for the given type.
+  void readLazyGenericEnvironment(GenericTypeDecl *type);
+
   /// Read the generic signature and type substitution map for a
   /// generic environment from \c Cursor.
   std::pair<GenericSignature *, TypeSubstitutionMap>

lib/AST/ASTContext.cpp

@@ -1483,6 +1483,14 @@ LazyContextData *ASTContext::getOrCreateLazyContextData(
     return contextData;
   }
 
+  // Create new lazy generic type data with the given loader.
+  if (isa<GenericTypeDecl>(decl)) {
+    auto *contextData = Allocate<LazyGenericTypeData>();
+    contextData->loader = lazyLoader;
+    Impl.LazyContexts[decl] = contextData;
+    return contextData;
+  }
+
   // Create new lazy function context data with the given loader.
   if (isa<AbstractFunctionDecl>(decl)) {
     auto *contextData = Allocate<LazyAbstractFunctionData>();
@@ -1514,6 +1522,11 @@ LazyAbstractFunctionData *ASTContext::getOrCreateLazyFunctionContextData(
                                                                 lazyLoader);
 }
 
+LazyGenericTypeData *ASTContext::getOrCreateLazyGenericTypeData(
+                                               const GenericTypeDecl *type,
+                                               LazyMemberLoader *lazyLoader) {
+  return (LazyGenericTypeData *)getOrCreateLazyContextData(type, lazyLoader);
+}
 void ASTContext::addDelayedConformanceDiag(
        NormalProtocolConformance *conformance,
        DelayedConformanceDiag fn) {

lib/AST/Decl.cpp

@@ -2140,6 +2140,30 @@ void GenericTypeDecl::setGenericParams(GenericParamList *params) {
       Param->setDeclContext(this);
 }
 
+GenericEnvironment *
+GenericTypeDecl::getLazyGenericEnvironmentSlow() const {
+  assert(GenericSigOrEnv.is<GenericSignature *>() &&
+         "not a lazily deserialized generic environment");
+  auto contextData = getASTContext().getOrCreateLazyGenericTypeData(this,
+                                                                    nullptr);
+  auto genericEnv = contextData->loader->loadGenericEnvironment(
+                                            this, contextData->genericEnvData);
+  const_cast<GenericTypeDecl *>(this)->setGenericEnvironment(genericEnv);
+  ++NumLazyGenericEnvironmentsLoaded;
+  return genericEnv;
+}
+
+void GenericTypeDecl::setLazyGenericEnvironment(LazyMemberLoader *lazyLoader,
+                                                GenericSignature *genericSig,
+                                                uint64_t genericEnvData) {
+  assert(GenericSigOrEnv.isNull() && "already have a generic signature");
+  GenericSigOrEnv = genericSig;
+
+  auto contextData =
+    getASTContext().getOrCreateLazyGenericTypeData(this, lazyLoader);
+  contextData->genericEnvData = genericEnvData;
+  ++NumLazyGenericEnvironments;
+}
 
 TypeAliasDecl::TypeAliasDecl(SourceLoc TypeAliasLoc, Identifier Name,
                              SourceLoc NameLoc, TypeLoc UnderlyingTy,

lib/Serialization/Deserialization.cpp

@@ -912,6 +912,29 @@ void ModuleFile::readGenericRequirements(
   }
 }
 
+uint64_t ModuleFile::allocateLazyGenericEnvironmentMap(
+                                                   TypeSubstitutionMap &&map) {
+  auto storedGenericEnvMap = new TypeSubstitutionMap(std::move(map));
+  GenericEnvironmentMaps.push_back(
+                    std::unique_ptr<TypeSubstitutionMap>(storedGenericEnvMap));
+  return reinterpret_cast<uint64_t>(storedGenericEnvMap);
+}
+
+void ModuleFile::readLazyGenericEnvironment(GenericTypeDecl *type) {
+  // Read the generic environment.
+  GenericSignature *genericSig;
+  TypeSubstitutionMap genericEnvMap;
+  std::tie(genericSig, genericEnvMap) =
+    readGenericEnvironmentPieces(DeclTypeCursor);
+
+  // Set up the lazy generic environment.
+  if (genericSig) {
+    type->setLazyGenericEnvironment(
+                   this, genericSig,
+                   allocateLazyGenericEnvironmentMap(std::move(genericEnvMap)));
+  }
+}
+
 std::pair<GenericSignature *, TypeSubstitutionMap>
 ModuleFile::readGenericEnvironmentPieces(
                              llvm::BitstreamCursor &Cursor,
@@ -2298,9 +2321,7 @@ Decl *ModuleFile::getDecl(DeclID DID, Optional<DeclContext *> ForcedContext) {
     declOrOffset = alias;
 
     if (genericParams) {
-      auto *env = maybeReadGenericEnvironment();
+      readLazyGenericEnvironment(alias);
-      assert(env && "generic typealias without environment");
-      alias->setGenericEnvironment(env);
     }
 
     alias->setDeserializedUnderlyingType(getType(underlyingTypeID));
@@ -2428,8 +2449,8 @@ Decl *ModuleFile::getDecl(DeclID DID, Optional<DeclContext *> ForcedContext) {
     if (isImplicit)
       theStruct->setImplicit();
 
-    auto *env = maybeReadGenericEnvironment();
+    // Read the generic environment.
-    theStruct->setGenericEnvironment(env);
+    readLazyGenericEnvironment(theStruct);
 
     theStruct->computeType();
 
@@ -2727,13 +2748,9 @@ Decl *ModuleFile::getDecl(DeclID DID, Optional<DeclContext *> ForcedContext) {
 
     // If there is a generic environment, lazily wire it up.
     if (genericSig) {
-      auto storedGenericEnvMap =
-        new TypeSubstitutionMap(std::move(genericEnvMap));
       fn->setLazyGenericEnvironment(
-                              this, genericSig,
+                  this, genericSig,
-                              reinterpret_cast<uint64_t>(storedGenericEnvMap));
+                  allocateLazyGenericEnvironmentMap(std::move(genericEnvMap)));
-      GenericEnvironmentMaps.push_back(
-                     std::unique_ptr<TypeSubstitutionMap>(storedGenericEnvMap));
     }
 
     SmallVector<ParameterList*, 2> paramLists;
@@ -2854,8 +2871,8 @@ Decl *ModuleFile::getDecl(DeclID DID, Optional<DeclContext *> ForcedContext) {
     if (auto genericParams = maybeReadGenericParams(DC))
       proto->setGenericParams(genericParams);
 
-    auto *env = maybeReadGenericEnvironment();
+    // Read the generic environment.
-    proto->setGenericEnvironment(env);
+    readLazyGenericEnvironment(proto);
 
     if (isImplicit)
       proto->setImplicit();
@@ -3030,8 +3047,8 @@ Decl *ModuleFile::getDecl(DeclID DID, Optional<DeclContext *> ForcedContext) {
     if (requiresStoredPropertyInits)
       theClass->setRequiresStoredPropertyInits(true);
 
-    auto *env = maybeReadGenericEnvironment();
+    // Read the generic environment.
-    theClass->setGenericEnvironment(env);
+    readLazyGenericEnvironment(theClass);
 
     theClass->computeType();
 
@@ -3086,8 +3103,8 @@ Decl *ModuleFile::getDecl(DeclID DID, Optional<DeclContext *> ForcedContext) {
       theEnum->setImplicit();
     theEnum->setRawType(getType(rawTypeID));
 
-    auto *env = maybeReadGenericEnvironment();
+    // Read the generic environment.
-    theEnum->setGenericEnvironment(env);
+    readLazyGenericEnvironment(theEnum);
 
     theEnum->computeType();
 
@@ -4342,6 +4359,8 @@ GenericEnvironment *ModuleFile::loadGenericEnvironment(const Decl *decl,
   GenericSignature *genericSig;
   if (auto func = dyn_cast<AbstractFunctionDecl>(decl))
     genericSig = func->getGenericSignature();
+  else if (auto type = dyn_cast<GenericTypeDecl>(decl))
+    genericSig = type->getGenericSignature();
   else
     llvm_unreachable("Cannot lazily deserialize generic environment");