Always store a generic environment in a nested archetype.

Author: DougGregor

include/swift/AST/Types.h

@@ -5396,7 +5396,7 @@ class ArchetypeType : public SubstitutableType,
   }
   Type InterfaceType;
   MutableArrayRef<std::pair<Identifier, Type>> NestedTypes;
-  mutable GenericEnvironment *Environment = nullptr;
+  GenericEnvironment *Environment = nullptr;
 
   void populateNestedTypes() const;
   void resolveNestedType(std::pair<Identifier, Type> &nested) const;
@@ -5474,7 +5474,7 @@ class ArchetypeType : public SubstitutableType,
   ArchetypeType *getRoot() const;
 
   /// Get the generic environment this archetype lives in.
-  GenericEnvironment *getGenericEnvironment() const;
+  GenericEnvironment *getGenericEnvironment() const { return Environment; }
 
   /// Get the protocol/class existential type that most closely represents the
   /// set of constraints on this archetype.
@@ -5706,7 +5706,8 @@ class NestedArchetypeType final : public ArchetypeType,
   getNew(const ASTContext &Ctx, ArchetypeType *Parent,
          DependentMemberType *InterfaceType,
          SmallVectorImpl<ProtocolDecl *> &ConformsTo,
-         Type Superclass, LayoutConstraint Layout);
+         Type Superclass, LayoutConstraint Layout,
+         GenericEnvironment *Environment);
 
   /// Retrieve the parent of this archetype, or null if this is a
   /// primary archetype.
@@ -5729,7 +5730,8 @@ class NestedArchetypeType final : public ArchetypeType,
                      ArchetypeType *Parent,
                      Type InterfaceType,
                      ArrayRef<ProtocolDecl *> ConformsTo,
-                     Type Superclass, LayoutConstraint Layout);
+                     Type Superclass, LayoutConstraint Layout,
+                      GenericEnvironment *Environment);
 };
 BEGIN_CAN_TYPE_WRAPPER(NestedArchetypeType, ArchetypeType)
 CanArchetypeType getParent() const {

lib/AST/GenericEnvironment.cpp

@@ -281,7 +281,7 @@ GenericEnvironment::getOrCreateArchetypeFromInterfaceType(Type depType) {
     auto *depMemTy = requirements.anchor->castTo<DependentMemberType>();
     result = NestedArchetypeType::getNew(ctx, parentArchetype, depMemTy,
                                          requirements.protos, superclass,
-                                         requirements.layout);
+                                         requirements.layout, this);
     parentArchetype->registerNestedType(depMemTy->getName(), result);
   } else if (genericParam->isTypeSequence()) {
     result = SequenceArchetypeType::get(ctx, this, genericParam,

lib/AST/Type.cpp

@@ -3210,19 +3210,6 @@ ArchetypeType::ArchetypeType(TypeKind Kind,
                           getSubclassTrailingObjects<ProtocolDecl *>());
 }
 
-GenericEnvironment *ArchetypeType::getGenericEnvironment() const {
-  if (Environment) {
-    return Environment;
-  }
-
-  // Nested archetypes get their environment from their root.
-  if (auto nested = dyn_cast<NestedArchetypeType>(this)) {
-    return Environment = nested->getRoot()->getGenericEnvironment();
-  }
-
-  llvm_unreachable("Archetype without a generic environment!");
-}
-
 ArchetypeType *ArchetypeType::getRoot() const {
   auto parent = this;
   while (auto nested = dyn_cast<NestedArchetypeType>(parent)) {
@@ -3283,11 +3270,11 @@ NestedArchetypeType::NestedArchetypeType(const ASTContext &Ctx,
                                        Type InterfaceType,
                                        ArrayRef<ProtocolDecl *> ConformsTo,
                                        Type Superclass,
-                                       LayoutConstraint Layout)
+                                       LayoutConstraint Layout,
+                                       GenericEnvironment *Environment)
   : ArchetypeType(TypeKind::NestedArchetype, Ctx,
                   Parent->getRecursiveProperties(),
-                  InterfaceType, ConformsTo, Superclass, Layout,
-                  /*Environment=*/nullptr),
+                  InterfaceType, ConformsTo, Superclass, Layout, Environment),
     Parent(Parent)
 {
 }
@@ -3608,7 +3595,8 @@ CanNestedArchetypeType NestedArchetypeType::getNew(
                                    DependentMemberType *InterfaceType,
                                    SmallVectorImpl<ProtocolDecl *> &ConformsTo,
                                    Type Superclass,
-                                   LayoutConstraint Layout) {
+                                   LayoutConstraint Layout,
+                                   GenericEnvironment *Environment) {
   assert(!Superclass || Superclass->getClassOrBoundGenericClass());
 
   // Gather the set of protocol declarations to which this archetype conforms.
@@ -3621,7 +3609,7 @@ CanNestedArchetypeType NestedArchetypeType::getNew(
       alignof(NestedArchetypeType), arena);
 
   return CanNestedArchetypeType(::new (mem) NestedArchetypeType(
-      Ctx, Parent, InterfaceType, ConformsTo, Superclass, Layout));
+      Ctx, Parent, InterfaceType, ConformsTo, Superclass, Layout, Environment));
 }
 
 CanPrimaryArchetypeType

lib/IRGen/GenArchetype.cpp

@@ -186,8 +186,7 @@ llvm::Value *irgen::emitArchetypeWitnessTableRef(IRGenFunction &IGF,
   auto wtable = IGF.tryGetLocalTypeData(archetype, localDataKind);
   if (wtable) return wtable;
 
-  auto origRoot = archetype->getRoot();
-  auto environment = origRoot->getGenericEnvironment();
+  auto environment = archetype->getGenericEnvironment();
 
   // Otherwise, ask the generic signature for the environment for the best
   // path to the conformance.