Sink the generic environment of an archetype down into ArchetypeType.

All non-nested archetype types already have a generic environment field,
and nested archetype types will need one soon anyway on their way to
irrelevance, so sinking the field doesn't isn't a net loss. While
here, isolate the code that lazily creates and caches the environment,
because only one subclass uses it (for opened archetypes) and I had
inadvertently broke caching there recently.

Author: DougGregor

include/swift/AST/Types.h

@@ -5396,6 +5396,7 @@ class ArchetypeType : public SubstitutableType,
   }
   Type InterfaceType;
   MutableArrayRef<std::pair<Identifier, Type>> NestedTypes;
+  mutable GenericEnvironment *Environment = nullptr;
 
   void populateNestedTypes() const;
   void resolveNestedType(std::pair<Identifier, Type> &nested) const;
@@ -5494,7 +5495,8 @@ class ArchetypeType : public SubstitutableType,
                 RecursiveTypeProperties properties,
                 Type InterfaceType,
                 ArrayRef<ProtocolDecl *> ConformsTo,
-                Type Superclass, LayoutConstraint Layout);
+                Type Superclass, LayoutConstraint Layout,
+                GenericEnvironment *Environment);
 };
 BEGIN_CAN_TYPE_WRAPPER(ArchetypeType, SubstitutableType)
 END_CAN_TYPE_WRAPPER(ArchetypeType, SubstitutableType)
@@ -5507,8 +5509,6 @@ class PrimaryArchetypeType final : public ArchetypeType,
   friend TrailingObjects;
   friend ArchetypeType;
 
-  GenericEnvironment *Environment;
-
 public:
   /// getNew - Create a new primary archetype with the given name.
   ///
@@ -5521,11 +5521,6 @@ class PrimaryArchetypeType final : public ArchetypeType,
                                SmallVectorImpl<ProtocolDecl *> &ConformsTo,
                                Type Superclass, LayoutConstraint Layout);
 
-  /// Retrieve the generic environment in which this archetype resides.
-  GenericEnvironment *getGenericEnvironment() const {
-    return Environment;
-  }
-
   GenericTypeParamType *getInterfaceType() const {
     return cast<GenericTypeParamType>(InterfaceType.getPointer());
   }
@@ -5551,10 +5546,6 @@ class OpaqueTypeArchetypeType final : public ArchetypeType,
   friend ArchetypeType;
   friend GenericSignatureBuilder;
 
-  /// A GenericEnvironment with this opaque archetype bound to the interface
-  /// type of the output type from the OpaqueDecl.
-  GenericEnvironment *Environment = nullptr;
-
   friend class GenericEnvironment;
 
   static OpaqueTypeArchetypeType *getNew(
@@ -5575,9 +5566,6 @@ class OpaqueTypeArchetypeType final : public ArchetypeType,
   /// Retrieve the set of substitutions applied to the opaque type.
   SubstitutionMap getSubstitutions() const;
 
-  /// Get a generic environment that has this opaque archetype bound within it.
-  GenericEnvironment *getGenericEnvironment() const;
-
   /// Compute the canonical interface type within the environment of this
   /// opaque type archetype.
   CanType getCanonicalInterfaceType(Type interfaceType);
@@ -5662,6 +5650,11 @@ class OpenedArchetypeType final : public ArchetypeType,
   mutable GenericEnvironment *Environment = nullptr;
   TypeBase *Opened;
   UUID ID;
+
+  /// Create a generic environment with this opened type bound to its generic
+  /// parameter.
+  GenericEnvironment *createGenericEnvironment() const;
+
 public:
   /// Create a new archetype that represents the opened type
   /// of an existential value.
@@ -5687,11 +5680,7 @@ class OpenedArchetypeType final : public ArchetypeType,
   Type getOpenedExistentialType() const {
     return Opened;
   }
-  
-  /// Get a generic environment with this opened type bound to its generic
-  /// parameter.
-  GenericEnvironment *getGenericEnvironment() const;
-  
+
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::OpenedArchetype;
   }
@@ -5765,8 +5754,6 @@ class SequenceArchetypeType final
   friend TrailingObjects;
   friend ArchetypeType;
 
-  GenericEnvironment *Environment;
-
 public:
   /// getNew - Create a new sequence archetype with the given name.
   ///
@@ -5778,9 +5765,6 @@ class SequenceArchetypeType final
       SmallVectorImpl<ProtocolDecl *> &ConformsTo, Type Superclass,
       LayoutConstraint Layout);
 
-  /// Retrieve the generic environment in which this archetype resides.
-  GenericEnvironment *getGenericEnvironment() const { return Environment; }
-
   GenericTypeParamType *getInterfaceType() const {
     return cast<GenericTypeParamType>(InterfaceType.getPointer());
   }

lib/AST/ASTContext.cpp

@@ -4323,10 +4323,7 @@ CanOpenedArchetypeType OpenedArchetypeType::get(Type existential,
   return CanOpenedArchetypeType(result);
 }
 
-GenericEnvironment *OpenedArchetypeType::getGenericEnvironment() const {
-  if (Environment)
-    return Environment;
-  
+GenericEnvironment *OpenedArchetypeType::createGenericEnvironment() const {
   auto thisType = const_cast<OpenedArchetypeType*>(this);
   auto &ctx = thisType->getASTContext();
   // Create a generic environment to represent the opened type.

lib/AST/Type.cpp

@@ -3185,9 +3185,11 @@ ArchetypeType::ArchetypeType(TypeKind Kind,
                              RecursiveTypeProperties properties,
                              Type InterfaceType,
                              ArrayRef<ProtocolDecl *> ConformsTo,
-                             Type Superclass, LayoutConstraint Layout)
+                             Type Superclass, LayoutConstraint Layout,
+                             GenericEnvironment *Environment)
   : SubstitutableType(Kind, &Ctx, properties),
-    InterfaceType(InterfaceType)
+    InterfaceType(InterfaceType),
+    Environment(Environment)
 {
   // Set up the bits we need for trailing objects to work.
   Bits.ArchetypeType.ExpandedNestedTypes = false;
@@ -3209,20 +3211,21 @@ ArchetypeType::ArchetypeType(TypeKind Kind,
 }
 
 GenericEnvironment *ArchetypeType::getGenericEnvironment() const {
-  auto root = getRoot();
-  if (auto primary = dyn_cast<PrimaryArchetypeType>(root)) {
-    return primary->getGenericEnvironment();
+  if (Environment) {
+    return Environment;
   }
-  if (auto opened = dyn_cast<OpenedArchetypeType>(root)) {
-    return opened->getGenericEnvironment();
-  }
-  if (auto opaque = dyn_cast<OpaqueTypeArchetypeType>(root)) {
-    return opaque->getGenericEnvironment();
+
+  // Opened archetypes lazily create their generic environment.
+  if (auto opened = dyn_cast<OpenedArchetypeType>(this)) {
+    return Environment = opened->createGenericEnvironment();
   }
-  if (auto opaque = dyn_cast<SequenceArchetypeType>(root)) {
-    return opaque->getGenericEnvironment();
+
+  // Nested archetypes get their environment from their root.
+  if (auto nested = dyn_cast<NestedArchetypeType>(this)) {
+    return Environment = nested->getRoot()->getGenericEnvironment();
   }
-  llvm_unreachable("unhandled root archetype kind?!");
+
+  llvm_unreachable("Archetype without a generic environment!");
 }
 
 ArchetypeType *ArchetypeType::getRoot() const {
@@ -3261,8 +3264,7 @@ PrimaryArchetypeType::PrimaryArchetypeType(const ASTContext &Ctx,
                                      Type Superclass, LayoutConstraint Layout)
   : ArchetypeType(TypeKind::PrimaryArchetype, Ctx,
                   RecursiveTypeProperties::HasArchetype,
-                  InterfaceType, ConformsTo, Superclass, Layout),
-    Environment(GenericEnv)
+                  InterfaceType, ConformsTo, Superclass, Layout, GenericEnv)
 {
 }
 
@@ -3274,7 +3276,8 @@ OpenedArchetypeType::OpenedArchetypeType(const ASTContext &Ctx,
   : ArchetypeType(TypeKind::OpenedArchetype, Ctx,
                   RecursiveTypeProperties::HasArchetype
                     | RecursiveTypeProperties::HasOpenedExistential,
-                  Type(), ConformsTo, Superclass, Layout),
+                  Type(), ConformsTo, Superclass, Layout,
+                  /*Environment=*/nullptr),
     Opened(Existential.getPointer()),
     ID(uuid)
 {
@@ -3288,7 +3291,8 @@ NestedArchetypeType::NestedArchetypeType(const ASTContext &Ctx,
                                        LayoutConstraint Layout)
   : ArchetypeType(TypeKind::NestedArchetype, Ctx,
                   Parent->getRecursiveProperties(),
-                  InterfaceType, ConformsTo, Superclass, Layout),
+                  InterfaceType, ConformsTo, Superclass, Layout,
+                  /*Environment=*/nullptr),
     Parent(Parent)
 {
 }
@@ -3300,8 +3304,8 @@ OpaqueTypeArchetypeType::OpaqueTypeArchetypeType(
     ArrayRef<ProtocolDecl*> conformsTo,
     Type superclass, LayoutConstraint layout)
   : ArchetypeType(TypeKind::OpaqueTypeArchetype, interfaceType->getASTContext(),
-                  properties, interfaceType, conformsTo, superclass, layout),
-    Environment(environment)
+                  properties, interfaceType, conformsTo, superclass, layout,
+                  environment)
 {
 }
 
@@ -3315,8 +3319,7 @@ SequenceArchetypeType::SequenceArchetypeType(
     LayoutConstraint Layout)
     : ArchetypeType(TypeKind::SequenceArchetype, Ctx,
                     RecursiveTypeProperties::HasArchetype, InterfaceType,
-                    ConformsTo, Superclass, Layout),
-      Environment(GenericEnv) {
+                    ConformsTo, Superclass, Layout, GenericEnv) {
   assert(cast<GenericTypeParamType>(InterfaceType.getPointer())->isTypeSequence());
 }
 
@@ -3328,10 +3331,6 @@ CanType OpaqueTypeArchetypeType::getCanonicalInterfaceType(Type interfaceType) {
       ->getCanonicalType();
 }
 
-GenericEnvironment *OpaqueTypeArchetypeType::getGenericEnvironment() const {
-  return Environment;
-}
-
 OpaqueTypeDecl *OpaqueTypeArchetypeType::getDecl() const {
   return Environment->getOpaqueTypeDecl();
 }