Merge pull request #22758 from jckarter/opened-archetype-generic-environment

Generic environments for opened archetypes

Author: jckarter

include/swift/AST/Types.h

@@ -4695,8 +4695,11 @@ class ArchetypeType : public SubstitutableType,
   /// Register a nested type with the given name.
   void registerNestedType(Identifier name, Type nested);
 
-  /// getPrimary - Return the primary archetype parent of this archetype.
-  PrimaryArchetypeType *getPrimary() const;
+  /// Return the root archetype parent of this archetype.
+  ArchetypeType *getRoot() const;
+  
+  /// Get the generic environment this archetype lives in.
+  GenericEnvironment *getGenericEnvironment() const;
 
   // Implement isa/cast/dyncast/etc.
   static bool classof(const TypeBase *T) {
@@ -4760,7 +4763,8 @@ class OpenedArchetypeType final : public ArchetypeType,
 {
   friend TrailingObjects;
   friend ArchetypeType;
-      
+  
+  mutable GenericEnvironment *Environment = nullptr;
   TypeBase *Opened;
   UUID ID;
 public:
@@ -4789,12 +4793,17 @@ class OpenedArchetypeType final : public ArchetypeType,
     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;
   }
 
 private:
-  OpenedArchetypeType(const ASTContext &Ctx, Type Existential,
+  OpenedArchetypeType(const ASTContext &Ctx,
+                      Type Existential,
                       ArrayRef<ProtocolDecl *> ConformsTo, Type Superclass,
                       LayoutConstraint Layout, UUID uuid);
 };
@@ -4831,6 +4840,10 @@ class NestedArchetypeType final : public ArchetypeType,
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::NestedArchetype;
   }
+  
+  DependentMemberType *getInterfaceType() const {
+    return cast<DependentMemberType>(InterfaceType.getPointer());
+  }
 
 private:
   NestedArchetypeType(const ASTContext &Ctx,

include/swift/Serialization/DeclTypeRecordNodes.def

@@ -88,7 +88,9 @@ TRAILING_INFO(TUPLE_TYPE_ELT)
 TYPE(FUNCTION)
 TRAILING_INFO(FUNCTION_PARAM)
 TYPE(METATYPE)
-TYPE(ARCHETYPE)
+TYPE(PRIMARY_ARCHETYPE)
+TYPE(OPENED_ARCHETYPE)
+TYPE(NESTED_ARCHETYPE)
 TYPE(PROTOCOL_COMPOSITION)
 TYPE(BOUND_GENERIC)
 

include/swift/Serialization/ModuleFormat.h

@@ -52,7 +52,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 474; // Last change: @_alwaysEmitIntoClient
+const uint16_t SWIFTMODULE_VERSION_MINOR = 475; // Last change: Generalize nested archetype serialization
 
 using DeclIDField = BCFixed<31>;
 
@@ -766,17 +766,24 @@ namespace decls_block {
     MetatypeRepresentationField        // representation
   >;
 
-  using ArchetypeTypeLayout = BCRecordLayout<
-    ARCHETYPE_TYPE,
+  using PrimaryArchetypeTypeLayout = BCRecordLayout<
+    PRIMARY_ARCHETYPE_TYPE,
     GenericEnvironmentIDField, // generic environment
-    TypeIDField                // interface type
+    BCVBR<4>, // generic type parameter depth
+    BCVBR<4>  // index + 1, or zero if we have a generic type parameter decl
   >;
 
-  using OpenedExistentialTypeLayout = BCRecordLayout<
-    OPENED_EXISTENTIAL_TYPE,
+  using OpenedArchetypeTypeLayout = BCRecordLayout<
+    OPENED_ARCHETYPE_TYPE,
     TypeIDField         // the existential type
   >;
-
+  
+  using NestedArchetypeTypeLayout = BCRecordLayout<
+    NESTED_ARCHETYPE_TYPE,
+    TypeIDField, // root archetype
+    TypeIDField // interface type relative to root
+  >;
+  
   using DynamicSelfTypeLayout = BCRecordLayout<
     DYNAMIC_SELF_TYPE,
     TypeIDField          // self type
@@ -1058,7 +1065,7 @@ namespace decls_block {
     // - the foreign error convention, if any
     // - inlinable body text, if any
   >;
-
+  
   // TODO: remove the unnecessary FuncDecl components here
   using AccessorLayout = BCRecordLayout<
     ACCESSOR_DECL,

lib/AST/ASTContext.cpp

@@ -4316,11 +4316,28 @@ CanOpenedArchetypeType OpenedArchetypeType::get(Type existential,
       ::new (mem) OpenedArchetypeType(ctx, existential,
                                 protos, layoutSuperclass,
                                 layoutConstraint, *knownID);
+  result->InterfaceType = GenericTypeParamType::get(0, 0, ctx);
+  
   openedExistentialArchetypes[*knownID] = result;
-
   return CanOpenedArchetypeType(result);
 }
 
+GenericEnvironment *OpenedArchetypeType::getGenericEnvironment() const {
+  if (Environment)
+    return Environment;
+  
+  auto thisType = Type(const_cast<OpenedArchetypeType*>(this));
+  auto &ctx = thisType->getASTContext();
+  // Create a generic environment to represent the opened type.
+  auto signature = ctx.getExistentialSignature(Opened->getCanonicalType(),
+                                               nullptr);
+  auto env = signature->createGenericEnvironment();
+  env->addMapping(signature->getGenericParams()[0], thisType);
+  Environment = env;
+  
+  return env;
+}
+
 CanType OpenedArchetypeType::getAny(Type existential) {
   if (auto metatypeTy = existential->getAs<ExistentialMetatypeType>()) {
     auto instanceTy = metatypeTy->getInstanceType();

lib/AST/ASTVerifier.cpp

@@ -632,13 +632,15 @@ class Verifier : public ASTWalker {
           // Only visit each archetype once.
           if (!visitedArchetypes.insert(archetype).second)
             return false;
+          
+          auto root = archetype->getRoot();
 
           // We should know about archetypes corresponding to opened
           // existential archetypes.
-          if (auto opened = dyn_cast<OpenedArchetypeType>(archetype)) {
+          if (auto opened = dyn_cast<OpenedArchetypeType>(root)) {
             if (OpenedExistentialArchetypes.count(opened) == 0) {
               Out << "Found opened existential archetype "
-                  << archetype->getString()
+                  << root->getString()
                   << " outside enclosing OpenExistentialExpr\n";
               return true;
             }
@@ -649,18 +651,18 @@ class Verifier : public ASTWalker {
           // Otherwise, the archetype needs to be from this scope.
           if (GenericEnv.empty() || !GenericEnv.back()) {
             Out << "AST verification error: archetype outside of generic "
-                   "context: " << archetype->getString() << "\n";
+                   "context: " << root->getString() << "\n";
             return true;
           }
 
           // Get the primary archetype.
-          auto *parent = archetype->getPrimary();
+          auto rootPrimary = cast<PrimaryArchetypeType>(root);
 
-          if (!GenericEnv.back().containsPrimaryArchetype(parent)) {
+          if (!GenericEnv.back().containsPrimaryArchetype(rootPrimary)) {
             Out << "AST verification error: archetype "
-                << archetype->getString() << " not allowed in this context\n";
+                << root->getString() << " not allowed in this context\n";
 
-            if (auto env = parent->getGenericEnvironment()) {
+            if (auto env = rootPrimary->getGenericEnvironment()) {
               if (auto owningDC = env->getOwningDeclContext()) {
                 llvm::errs() << "archetype came from:\n";
                 owningDC->dumpContext();

lib/AST/GenericEnvironment.cpp

@@ -122,7 +122,9 @@ Type GenericEnvironment::mapTypeIntoContext(GenericEnvironment *env,
 }
 
 Type MapTypeOutOfContext::operator()(SubstitutableType *type) const {
-  return cast<ArchetypeType>(type)->getInterfaceType();
+  auto archetype = cast<ArchetypeType>(type);
+  
+  return archetype->getInterfaceType();
 }
 
 Type TypeBase::mapTypeOutOfContext() {

lib/AST/GenericSignatureBuilder.cpp

@@ -2811,6 +2811,8 @@ static void concretizeNestedTypeFromConcreteParent(
         ->getTypeWitness(assocType, builder.getLazyResolver());
     if (!witnessType || witnessType->hasError())
       return; // FIXME: should we delay here?
+  } else if (auto archetype = concreteParent->getAs<ArchetypeType>()) {
+    witnessType = archetype->getNestedType(assocType->getName());
   } else {
     witnessType = DependentMemberType::get(concreteParent, assocType);
   }
@@ -2902,7 +2904,7 @@ PotentialArchetype *PotentialArchetype::updateNestedTypeForConformance(
 
 void ArchetypeType::resolveNestedType(
                                     std::pair<Identifier, Type> &nested) const {
-  auto genericEnv = getPrimary()->getGenericEnvironment();
+  auto genericEnv = getGenericEnvironment();
   auto &builder = *genericEnv->getGenericSignatureBuilder();
 
   Type interfaceType = getInterfaceType();
@@ -3886,14 +3888,17 @@ ResolvedType GenericSignatureBuilder::maybeResolveEquivalenceClass(
   }
 
   // The equivalence class of a dependent member type is determined by its
-  // base equivalence class.
+  // base equivalence class, if there is one.
   if (auto depMemTy = type->getAs<DependentMemberType>()) {
     // Find the equivalence class of the base.
     auto resolvedBase =
       maybeResolveEquivalenceClass(depMemTy->getBase(),
                                    resolutionKind,
                                    wantExactPotentialArchetype);
     if (!resolvedBase) return resolvedBase;
+    // If the base is concrete, so is this member.
+    if (resolvedBase.getAsConcreteType())
+      return ResolvedType::forConcrete(type);
 
     // Find the nested type declaration for this.
     auto baseEquivClass = resolvedBase.getEquivalenceClass(*this);

lib/AST/Type.cpp

@@ -2368,12 +2368,23 @@ ArchetypeType::ArchetypeType(TypeKind Kind,
                           getSubclassTrailingObjects<ProtocolDecl *>());
 }
 
-PrimaryArchetypeType *ArchetypeType::getPrimary() const {
-  ArchetypeType *archetype = const_cast<ArchetypeType*>(this);
-  while (auto child = dyn_cast<NestedArchetypeType>(archetype)) {
-    archetype = child->getParent();
+GenericEnvironment *ArchetypeType::getGenericEnvironment() const {
+  auto root = getRoot();
+  if (auto primary = dyn_cast<PrimaryArchetypeType>(root)) {
+    return primary->getGenericEnvironment();
   }
-  return dyn_cast<PrimaryArchetypeType>(archetype);
+  if (auto opened = dyn_cast<OpenedArchetypeType>(root)) {
+    return opened->getGenericEnvironment();
+  }
+  llvm_unreachable("unhandled root archetype kind?!");
+}
+
+ArchetypeType *ArchetypeType::getRoot() const {
+  auto parent = this;
+  while (auto nested = dyn_cast<NestedArchetypeType>(parent)) {
+    parent = nested->getParent();
+  }
+  return const_cast<ArchetypeType*>(parent);
 }
 
 PrimaryArchetypeType::PrimaryArchetypeType(const ASTContext &Ctx,
@@ -3368,10 +3379,10 @@ Type Type::transformRec(
     if (Optional<Type> transformed = fn(getPointer()))
       return *transformed;
 
-    // Recurse.
+    // Recur.
   }
 
-  // Recursive into children of this type.
+  // Recur into children of this type.
   TypeBase *base = getPointer();
   switch (base->getKind()) {
 #define ALWAYS_CANONICAL_TYPE(Id, Parent) \

lib/IRGen/GenArchetype.cpp

@@ -172,7 +172,7 @@ llvm::Value *irgen::emitArchetypeWitnessTableRef(IRGenFunction &IGF,
   auto wtable = IGF.tryGetLocalTypeData(archetype, localDataKind);
   if (wtable) return wtable;
 
-  // If we don't have an environment, this must be an implied witness table
+  // If we have an opened type, this must be an implied witness table
   // reference.
   // FIXME: eliminate this path when opened types have generic environments.
   if (auto opened = dyn_cast<OpenedArchetypeType>(archetype)) {
@@ -195,7 +195,9 @@ llvm::Value *irgen::emitArchetypeWitnessTableRef(IRGenFunction &IGF,
     });
   }
 
-  auto environment = archetype->getPrimary()->getGenericEnvironment();
+  assert(isa<PrimaryArchetypeType>(archetype->getRoot())
+         && "might not work yet for opened/opaque archetypes!");
+  auto environment = archetype->getGenericEnvironment();
 
   // Otherwise, ask the generic signature for the environment for the best
   // path to the conformance.

lib/IRGen/GenType.cpp

@@ -1508,7 +1508,7 @@ const TypeInfo &TypeConverter::getCompleteTypeInfo(CanType T) {
 
 ArchetypeType *TypeConverter::getExemplarArchetype(ArchetypeType *t) {
   // Get the primary archetype.
-  auto primary = t->getPrimary();
+  auto primary = dyn_cast<PrimaryArchetypeType>(t->getRoot());
 
   // If there is no primary (IOW, it's an opened archetype), the archetype is
   // an exemplar.
@@ -2219,7 +2219,7 @@ void IRGenFunction::setLocalSelfMetadata(llvm::Value *value,
 
 #ifndef NDEBUG
 bool TypeConverter::isExemplarArchetype(ArchetypeType *arch) const {
-  auto primary = arch->getPrimary();
+  auto primary = dyn_cast<PrimaryArchetypeType>(arch->getRoot());
   if (!primary) return true;
   auto genericEnv = primary->getGenericEnvironment();
 

lib/IRGen/Outlining.cpp

@@ -110,7 +110,10 @@ irgen::getTypeAndGenericSignatureForManglingOutlineFunction(SILType type) {
     GenericEnvironment *env = nullptr;
     loweredType.findIf([&env](Type t) -> bool {
         if (auto arch = t->getAs<ArchetypeType>()) {
-          env = arch->getPrimary()->getGenericEnvironment();
+          auto root = arch->getRoot();
+          if (!isa<PrimaryArchetypeType>(root))
+            return false;
+          env = root->getGenericEnvironment();
           return true;
         }
         return false;

lib/SIL/SILVerifier.cpp

@@ -77,16 +77,14 @@ static llvm::cl::opt<bool> SkipConvertEscapeToNoescapeAttributes(
 /// Returns true if A is an opened existential type or is equal to an
 /// archetype from F's generic context.
 static bool isArchetypeValidInFunction(ArchetypeType *A, const SILFunction *F) {
-  if (isa<OpenedArchetypeType>(A))
+  auto root = dyn_cast<PrimaryArchetypeType>(A->getRoot());
+  if (!root)
     return true;
 
-  // Find the primary archetype.
-  auto P = A->getPrimary();
-
   // Ok, we have a primary archetype, make sure it is in the nested generic
   // environment of our caller.
   if (auto *genericEnv = F->getGenericEnvironment())
-    if (P->getGenericEnvironment() == genericEnv)
+    if (root->getGenericEnvironment() == genericEnv)
       return true;
 
   return false;

lib/SILOptimizer/Utils/Generics.cpp

@@ -1036,7 +1036,9 @@ shouldBePartiallySpecialized(Type Replacement,
   llvm::SmallSetVector<ArchetypeType *, 2> UsedArchetypes;
   Replacement.visit([&](Type Ty) {
     if (auto Archetype = Ty->getAs<ArchetypeType>()) {
-      UsedArchetypes.insert(Archetype->getPrimary());
+      if (auto Primary = dyn_cast<PrimaryArchetypeType>(Archetype->getRoot())) {
+        UsedArchetypes.insert(Primary);
+      }
     }
   });
 
@@ -1264,7 +1266,9 @@ void FunctionSignaturePartialSpecializer::collectUsedCallerArchetypes(
     // Add used generic parameters/archetypes.
     Replacement.visit([&](Type Ty) {
       if (auto Archetype = Ty->getAs<ArchetypeType>()) {
-        UsedCallerArchetypes.insert(Archetype->getPrimary());
+        if (auto Primary = dyn_cast<PrimaryArchetypeType>(Archetype->getRoot())) {
+          UsedCallerArchetypes.insert(Primary);
+        }
       }
     });
   }

lib/Sema/CSSimplify.cpp

@@ -18,6 +18,8 @@
 #include "CSFix.h"
 #include "ConstraintSystem.h"
 #include "swift/AST/ExistentialLayout.h"
+#include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/GenericSignature.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/Basic/StringExtras.h"

lib/Sema/TypeCheckConstraints.cpp

@@ -2950,9 +2950,13 @@ static Type replaceArchetypesWithTypeVariables(ConstraintSystem &cs,
       if (found != types.end())
         return found->second;
 
-      if (isa<NestedArchetypeType>(origType))
-        return Type();
-      else if (auto archetypeType = dyn_cast<ArchetypeType>(origType)) {
+      if (auto archetypeType = dyn_cast<ArchetypeType>(origType)) {
+        auto root = archetypeType->getRoot();
+        // For other nested types, fail here so the default logic in subst()
+        // for nested types applies.
+        if (root != archetypeType)
+          return Type();
+        
         auto locator = cs.getConstraintLocator(nullptr);
         auto replacement = cs.createTypeVariable(locator);