Merge pull request #73570 from slavapestov/pack-expansion-closures-part-4

Allow local archetypes in interface types of local declarations

Author: slavapestov

include/swift/AST/Decl.h

@@ -6010,10 +6010,6 @@ enum class PropertyWrapperSynthesizedPropertyKind {
 class VarDecl : public AbstractStorageDecl {
   friend class NamingPatternRequest;
   NamedPattern *NamingPattern = nullptr;
-  /// When the variable is declared in context of a for-in loop over the elements of 
-  /// a parameter pack, this is the opened element environment of the pack expansion
-  /// to use as the variable's context generic environment.
-  GenericEnvironment *OpenedElementEnvironment = nullptr;
 
 public:
   enum class Introducer : uint8_t {
@@ -6155,13 +6151,6 @@ class VarDecl : public AbstractStorageDecl {
   NamedPattern *getNamingPattern() const;
   void setNamingPattern(NamedPattern *Pat);
 
-  GenericEnvironment *getOpenedElementEnvironment() const {
-    return OpenedElementEnvironment;
-  }
-  void setOpenedElementEnvironment(GenericEnvironment *Env) {
-    OpenedElementEnvironment = Env;
-  }
-
   /// If this is a VarDecl that does not belong to a CaseLabelItem's pattern,
   /// return this. Otherwise, this VarDecl must belong to a CaseStmt's
   /// CaseLabelItem. In that case, return the first case label item of the first

include/swift/AST/Types.h

@@ -128,10 +128,9 @@ class RecursiveTypeProperties {
     /// This type expression contains a TypeVariableType.
     HasTypeVariable      = 0x01,
 
-    /// This type expression contains a context-dependent archetype, either a
-    /// \c PrimaryArchetypeType, \c OpenedArchetypeType,
-    /// \c ElementArchetypeType, or \c PackArchetype.
-    HasArchetype         = 0x02,
+    /// This type expression contains a PrimaryArchetypeType
+    /// or PackArchetypeType.
+    HasPrimaryArchetype  = 0x02,
 
     /// This type expression contains a GenericTypeParamType.
     HasTypeParameter     = 0x04,
@@ -171,7 +170,7 @@ class RecursiveTypeProperties {
     /// This type contains a parameterized existential type \c any P<T>.
     HasParameterizedExistential = 0x2000,
 
-    /// This type contains an ElementArchetype.
+    /// This type contains an ElementArchetypeType.
     HasElementArchetype = 0x4000,
 
     /// Whether the type is allocated in the constraint solver arena. This can
@@ -205,9 +204,9 @@ class RecursiveTypeProperties {
   /// variable?
   bool hasTypeVariable() const { return Bits & HasTypeVariable; }
 
-  /// Does a type with these properties structurally contain a
-  /// context-dependent archetype (that is, a Primary- or OpenedArchetype)?
-  bool hasArchetype() const { return Bits & HasArchetype; }
+  /// Does a type with these properties structurally contain a primary
+  /// archetype?
+  bool hasPrimaryArchetype() const { return Bits & HasPrimaryArchetype; }
 
   /// Does a type with these properties structurally contain an
   /// archetype from an opaque type declaration?
@@ -696,9 +695,21 @@ class alignas(1 << TypeAlignInBits) TypeBase
     return getRecursiveProperties().hasPlaceholder();
   }
 
-  /// Determine whether the type involves a context-dependent archetype.
+  /// Determine whether the type involves a primary archetype.
+  bool hasPrimaryArchetype() const {
+    return getRecursiveProperties().hasPrimaryArchetype();
+  }
+
+  /// Whether the type contains a PackArchetypeType.
+  bool hasPackArchetype() const {
+    return getRecursiveProperties().hasPackArchetype();
+  }
+
+  /// Determine whether the type involves a primary, pack or local archetype.
+  ///
+  /// FIXME: Replace all remaining callers with a more precise check.
   bool hasArchetype() const {
-    return getRecursiveProperties().hasArchetype();
+    return hasPrimaryArchetype() || hasLocalArchetype();
   }
 
   /// Determine whether the type involves an opened existential archetype.
@@ -727,11 +738,6 @@ class alignas(1 << TypeAlignInBits) TypeBase
     return getRecursiveProperties().hasPack();
   }
 
-  /// Whether the type contains a PackArchetypeType.
-  bool hasPackArchetype() const {
-    return getRecursiveProperties().hasPackArchetype();
-  }
-
   /// Whether the type has any flavor of pack.
   bool hasAnyPack() const {
     return hasParameterPack() || hasPack() || hasPackArchetype();

include/swift/SIL/SILType.h

@@ -461,7 +461,7 @@ class SILType {
   }
 
   /// Returns true if the referenced type is expressed in terms of one
-  /// or more opened existential types.
+  /// or more opened existential archetypes.
   bool hasOpenedExistential() const {
     return getASTType()->hasOpenedExistential();
   }
@@ -470,6 +470,12 @@ class SILType {
     return getASTType()->canBeClass();
   }
 
+  /// Returns true if the referenced type is expressed in terms of one
+  /// or more element archetypes.
+  bool hasElementArchetype() const {
+    return getASTType()->hasElementArchetype();
+  }
+
   /// Returns true if the referenced type is expressed in terms of one
   /// or more local archetypes.
   bool hasLocalArchetype() const {
@@ -570,9 +576,14 @@ class SILType {
     return false;
   }
 
-  /// True if the type involves any archetypes.
+  /// True if the type involves any primary or local archetypes.
   bool hasArchetype() const { return getASTType()->hasArchetype(); }
 
+  /// True if the type involves any primary archetypes.
+  bool hasPrimaryArchetype() const {
+    return getASTType()->hasPrimaryArchetype();
+  }
+
   /// True if the type involves any opaque archetypes.
   bool hasOpaqueArchetype() const {
     return getASTType()->hasOpaqueArchetype();

lib/AST/Decl.cpp

@@ -5423,8 +5423,8 @@ Type TypeAliasDecl::getUnderlyingType() const {
 void TypeAliasDecl::setUnderlyingType(Type underlying) {
   // lldb creates global typealiases containing archetypes
   // sometimes...
-  if (underlying->hasArchetype() && isGenericContext())
-    underlying = underlying->mapTypeOutOfContext();
+  assert(!underlying->hasArchetype() || !isGenericContext());
+
   getASTContext().evaluator.cacheOutput(
           StructuralTypeRequest{const_cast<TypeAliasDecl *>(this)},
           std::move(underlying));
@@ -7260,12 +7260,6 @@ VarDecl::VarDecl(DeclKind kind, bool isStatic, VarDecl::Introducer introducer,
 }
 
 Type VarDecl::getTypeInContext() const {
-  // If the variable is declared in context of a for-in loop over the elements
-  // of a parameter pack, its interface type must be mapped into context using
-  // the opened element environment of the pack expansion.
-  if (auto *env = getOpenedElementEnvironment())
-    return GenericEnvironment::mapTypeIntoContext(env, getInterfaceType());
-
   return getDeclContext()->mapTypeIntoContext(getInterfaceType());
 }
 
@@ -8144,7 +8138,7 @@ Type VarDecl::getPropertyWrapperInitValueInterfaceType() const {
     return Type();
 
   Type valueInterfaceTy = initInfo.getWrappedValuePlaceholder()->getType();
-  if (valueInterfaceTy->hasArchetype())
+  if (valueInterfaceTy->hasPrimaryArchetype())
     valueInterfaceTy = valueInterfaceTy->mapTypeOutOfContext();
 
   return valueInterfaceTy;

lib/AST/GenericEnvironment.cpp

@@ -404,28 +404,24 @@ GenericEnvironment::maybeApplyOuterContextSubstitutions(Type type) const {
 
 Type GenericEnvironment::mapTypeIntoContext(GenericEnvironment *env,
                                             Type type) {
-  assert((!type->hasArchetype() || type->hasLocalArchetype()) &&
-         "already have a contextual type");
-  assert((env || !type->hasTypeParameter()) &&
-         "no generic environment provided for type with type parameters");
+  assert(!type->hasPrimaryArchetype() && "already have a contextual type");
 
   if (!env) {
+    assert(!type->hasTypeParameter() &&
+           "no generic environment provided for type with type parameters");
     return type;
   }
 
   return env->mapTypeIntoContext(type);
 }
 
 Type MapTypeOutOfContext::operator()(SubstitutableType *type) const {
-  auto archetype = cast<ArchetypeType>(type);
-  if (isa<OpaqueTypeArchetypeType>(archetype->getRoot()))
-    return Type();
-
-  // Leave opened archetypes alone; they're handled contextually.
-  if (isa<OpenedArchetypeType>(archetype))
-    return Type(type);
+  if (isa<PrimaryArchetypeType>(type) ||
+      isa<PackArchetypeType>(type)) {
+    return cast<ArchetypeType>(type)->getInterfaceType();
+  }
 
-  return archetype->getInterfaceType();
+  return type;
 }
 
 Type TypeBase::mapTypeOutOfContext() {
@@ -632,8 +628,7 @@ Type QueryInterfaceTypeSubstitutions::operator()(SubstitutableType *type) const{
 Type GenericEnvironment::mapTypeIntoContext(
                                 Type type,
                                 LookupConformanceFn lookupConformance) const {
-  assert((!type->hasArchetype() || type->hasLocalArchetype()) &&
-         "already have a contextual type");
+  assert(!type->hasPrimaryArchetype() && "already have a contextual type");
 
   Type result = type.subst(QueryInterfaceTypeSubstitutions(this),
                            lookupConformance,
@@ -668,7 +663,7 @@ GenericEnvironment::mapContextualPackTypeIntoElementContext(Type type) const {
   assert(getKind() == Kind::OpenedElement);
   assert(!type->hasTypeParameter() && "expected contextual type");
 
-  if (!type->hasArchetype()) return type;
+  if (!type->hasPackArchetype()) return type;
 
   auto sig = getGenericSignature();
   auto shapeClass = getOpenedElementShapeClass();
@@ -698,9 +693,9 @@ GenericEnvironment::mapContextualPackTypeIntoElementContext(CanType type) const
 Type
 GenericEnvironment::mapPackTypeIntoElementContext(Type type) const {
   assert(getKind() == Kind::OpenedElement);
-  assert(!type->hasArchetype());
+  assert(!type->hasPackArchetype());
 
-  if (!type->hasTypeParameter()) return type;
+  if (!type->hasParameterPack()) return type;
 
   // Get a contextual type in the original generic environment, not the
   // substituted one, which is what mapContextualPackTypeIntoElementContext()
@@ -720,52 +715,58 @@ GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
   // generic environment.
   assert(type->hasElementArchetype());
 
-  ElementArchetypeType *element = nullptr;
-  type.visit([&](Type type) {
-    auto archetype = type->getAs<ElementArchetypeType>();
-    if (!element && archetype)
-      element = archetype;
-  });
+  GenericEnvironment *elementEnv = nullptr;
 
-  auto sig = getGenericSignature();
-  auto *elementEnv = element->getGenericEnvironment();
-  auto shapeClass = elementEnv->getOpenedElementShapeClass();
-  QueryInterfaceTypeSubstitutions substitutions(this);
+  // Map element archetypes to interface types in the element generic
+  // environment's signature.
+  type = type.subst(
+    [&](SubstitutableType *type) -> Type {
+      auto *archetype = cast<ArchetypeType>(type);
 
-  type = type->mapTypeOutOfContext();
+      if (isa<OpenedArchetypeType>(archetype))
+        return archetype;
 
-  auto interfaceType = element->getInterfaceType();
+      if (isa<ElementArchetypeType>(archetype)) {
+        assert(!elementEnv ||
+               elementEnv == archetype->getGenericEnvironment());
+        elementEnv = archetype->getGenericEnvironment();
+      }
 
-  llvm::SmallDenseMap<GenericParamKey, GenericTypeParamType *>
-      packParamForElement;
-  auto elementDepth = interfaceType->getRootGenericParam()->getDepth();
+      return archetype->getInterfaceType();
+    },
+    MakeAbstractConformanceForGenericType(),
+    SubstFlags::AllowLoweredTypes |
+    SubstFlags::PreservePackExpansionLevel);
+
+  auto shapeClass = elementEnv->getOpenedElementShapeClass();
+
+  llvm::SmallVector<GenericTypeParamType *, 2> members;
+  auto elementDepth = elementEnv->getGenericSignature()->getMaxDepth();
 
+  auto sig = getGenericSignature();
   for (auto *genericParam : sig.getGenericParams()) {
     if (!genericParam->isParameterPack())
       continue;
 
     if (!sig->haveSameShape(genericParam, shapeClass))
       continue;
 
-    GenericParamKey elementKey(/*isParameterPack*/false,
-                               /*depth*/elementDepth,
-                               /*index*/packParamForElement.size());
-    packParamForElement[elementKey] = genericParam;
+    members.push_back(genericParam);
   }
 
-  // Map element archetypes to the pack archetypes by converting
-  // element types to interface types and adding the isParameterPack
-  // bit. Then, map type parameters to archetypes.
+  // Map element interface types to pack archetypes.
+  QueryInterfaceTypeSubstitutions mapIntoContext(this);
   return type.subst(
       [&](SubstitutableType *type) {
         auto *genericParam = type->getAs<GenericTypeParamType>();
         if (!genericParam)
           return Type();
 
-        if (auto *packParam = packParamForElement[{genericParam}])
-          return substitutions(packParam);
-
-        return substitutions(genericParam);
+        if (genericParam->getDepth() == elementDepth) {
+          genericParam = members[genericParam->getIndex()];
+          assert(genericParam->isParameterPack());
+        }
+        return mapIntoContext(genericParam);
       },
       LookUpConformanceInSignature(sig.getPointer()),
       SubstFlags::PreservePackExpansionLevel);

lib/AST/Type.cpp

@@ -3454,7 +3454,7 @@ PrimaryArchetypeType::PrimaryArchetypeType(const ASTContext &Ctx,
                                      ArrayRef<ProtocolDecl *> ConformsTo,
                                      Type Superclass, LayoutConstraint Layout)
   : ArchetypeType(TypeKind::PrimaryArchetype, Ctx,
-                  RecursiveTypeProperties::HasArchetype,
+                  RecursiveTypeProperties::HasPrimaryArchetype,
                   InterfaceType, ConformsTo, Superclass, Layout, GenericEnv)
 {
   assert(!InterfaceType->isParameterPack());
@@ -3518,8 +3518,7 @@ OpenedArchetypeType::OpenedArchetypeType(
     LayoutConstraint layout)
   : LocalArchetypeType(TypeKind::OpenedArchetype,
                        interfaceType->getASTContext(),
-                       RecursiveTypeProperties::HasArchetype
-                         | RecursiveTypeProperties::HasOpenedExistential,
+                       RecursiveTypeProperties::HasOpenedExistential,
                        interfaceType, conformsTo, superclass, layout,
                        environment)
 {
@@ -3535,8 +3534,8 @@ PackArchetypeType::PackArchetypeType(
     ArrayRef<ProtocolDecl *> ConformsTo, Type Superclass,
     LayoutConstraint Layout, PackShape Shape)
     : ArchetypeType(TypeKind::PackArchetype, Ctx,
-                    RecursiveTypeProperties::HasArchetype |
-                        RecursiveTypeProperties::HasPackArchetype,
+                    RecursiveTypeProperties::HasPrimaryArchetype |
+                    RecursiveTypeProperties::HasPackArchetype,
                     InterfaceType, ConformsTo, Superclass, Layout, GenericEnv) {
   assert(InterfaceType->isParameterPack());
   *getTrailingObjects<PackShape>() = Shape;
@@ -3586,7 +3585,6 @@ ElementArchetypeType::ElementArchetypeType(
     ArrayRef<ProtocolDecl *> ConformsTo, Type Superclass,
     LayoutConstraint Layout)
     : LocalArchetypeType(TypeKind::ElementArchetype, Ctx,
-                         RecursiveTypeProperties::HasArchetype |
                          RecursiveTypeProperties::HasElementArchetype,
                          InterfaceType,
                          ConformsTo, Superclass, Layout, GenericEnv) {

lib/AST/TypeCheckRequests.cpp

@@ -1074,7 +1074,9 @@ void InterfaceTypeRequest::cacheResult(Type type) const {
   if (type) {
     assert(!type->hasTypeVariable() && "Type variable in interface type");
     assert(!type->is<InOutType>() && "Interface type must be materializable");
-    assert(!type->hasArchetype() && "Archetype in interface type");
+    assert(!type->hasPrimaryArchetype() && "Archetype in interface type");
+    assert(decl->getDeclContext()->isLocalContext() || !type->hasLocalArchetype() &&
+           "Local archetype in interface type of non-local declaration");
   }
   decl->TypeAndAccess.setPointer(type);
 }