AST: Introduce ElementArchetypeType

Author: slavapestov

include/swift/AST/GenericEnvironment.h

@@ -51,13 +51,24 @@ class QueryInterfaceTypeSubstitutions {
   Type operator()(SubstitutableType *type) const;
 };
 
+/// Extra data in a generic environment for an opaque type.
+struct OpaqueEnvironmentData {
+  OpaqueTypeDecl *decl;
+  SubstitutionMap subMap;
+};
+
 /// Extra data in a generic environment for an opened existential.
-struct OpenedGenericEnvironmentData {
+struct OpenedExistentialEnvironmentData {
   Type existential;
   GenericSignature parentSig;
   UUID uuid;
 };
 
+/// Extra data in a generic environment for an opened element.
+struct OpenedElementEnvironmentData {
+  UUID uuid;
+};
+
 /// Describes the mapping between archetypes and interface types for the
 /// generic parameters of a DeclContext.
 ///
@@ -68,17 +79,23 @@ struct OpenedGenericEnvironmentData {
 ///
 class alignas(1 << DeclAlignInBits) GenericEnvironment final
     : private llvm::TrailingObjects<
-        GenericEnvironment, OpaqueTypeDecl *, SubstitutionMap,
-        OpenedGenericEnvironmentData, Type> {
+        GenericEnvironment,
+        OpaqueEnvironmentData,
+        OpenedExistentialEnvironmentData,
+        OpenedElementEnvironmentData,
+        Type> {
 public:
   enum class Kind {
     /// A normal generic environment, determined only by its generic
     /// signature.
     Primary,
-    /// A generic environment describing an opened existential archetype.
-    OpenedExistential,
     /// A generic environment describing an opaque type archetype.
     Opaque,
+    /// A generic environment describing an opened existential archetype.
+    OpenedExistential,
+    /// A generic environment describing an opened element type of a
+    /// pack archetype inside a pack expansion expression.
+    OpenedElement,
   };
 
   class NestedTypeStorage;
@@ -91,10 +108,10 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   friend TrailingObjects;
   friend OpaqueTypeArchetypeType;
 
-  size_t numTrailingObjects(OverloadToken<OpaqueTypeDecl *>) const;
-  size_t numTrailingObjects(OverloadToken<SubstitutionMap>) const;
+  size_t numTrailingObjects(OverloadToken<OpaqueEnvironmentData>) const;
+  size_t numTrailingObjects(OverloadToken<OpenedExistentialEnvironmentData>) const;
+  size_t numTrailingObjects(OverloadToken<OpenedElementEnvironmentData>) const;
   size_t numTrailingObjects(OverloadToken<Type>) const;
-  size_t numTrailingObjects(OverloadToken<OpenedGenericEnvironmentData>) const;
 
   /// Retrieve the array containing the context types associated with the
   /// generic parameters, stored in parallel with the generic parameters of the
@@ -109,12 +126,20 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// Get the nested type storage, allocating it if required.
   NestedTypeStorage &getOrCreateNestedTypeStorage();
 
+  /// Private constructor for primary environments.
   explicit GenericEnvironment(GenericSignature signature);
+
+  /// Private constructor for opaque type environments.
+  explicit GenericEnvironment(
+      GenericSignature signature, OpaqueTypeDecl *opaque, SubstitutionMap subs);
+
+  /// Private constructor for opened existential environments.
   explicit GenericEnvironment(
       GenericSignature signature,
       Type existential, GenericSignature parentSig, UUID uuid);
-  explicit GenericEnvironment(
-      GenericSignature signature, OpaqueTypeDecl *opaque, SubstitutionMap subs);
+
+  /// Private constructor for opened element environments.
+  explicit GenericEnvironment(GenericSignature signature, UUID uuid);
 
   friend ArchetypeType;
   friend QueryInterfaceTypeSubstitutions;
@@ -156,9 +181,17 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// create a generic environment.
   SubstitutionMap getOpaqueSubstitutions() const;
 
+  /// Retrieve the UUID for an opened element environment.
+  UUID getOpenedElementUUID() const;
+
   /// Create a new, primary generic environment.
   static GenericEnvironment *forPrimary(GenericSignature signature);
 
+  /// Create a new generic environment for an opaque type with the given set of
+  /// outer substitutions.
+  static GenericEnvironment *forOpaqueType(
+      OpaqueTypeDecl *opaque, SubstitutionMap subs, AllocationArena arena);
+
   /// Create a new generic environment for an opened existential.
   ///
   /// \param existential The subject existential type
@@ -167,10 +200,14 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   static GenericEnvironment *
   forOpenedExistential(Type existential, GenericSignature parentSig, UUID uuid);
 
-  /// Create a new generic environment for an opaque type with the given set of
-  /// outer substitutions.
-  static GenericEnvironment *forOpaqueType(
-      OpaqueTypeDecl *opaque, SubstitutionMap subs, AllocationArena arena);
+  /// Create a new generic environment for an opened element.
+  ///
+  /// \param signature The opened element signature, which is the same as the
+  /// signature of the context whose element type is being opened, but with
+  /// the pack parameter bit erased from one or more generic parameters
+  /// \param uuid The unique identifier for this opened element
+  static GenericEnvironment *
+  forOpenedElement(GenericSignature signature, UUID uuid);
 
   /// Make vanilla new/delete illegal.
   void *operator new(size_t Bytes) = delete;

include/swift/AST/TypeNodes.def

@@ -148,6 +148,7 @@ ABSTRACT_TYPE(Substitutable, Type)
     ALWAYS_CANONICAL_TYPE(PrimaryArchetype, ArchetypeType)
     ALWAYS_CANONICAL_TYPE(OpaqueTypeArchetype, ArchetypeType)
     ALWAYS_CANONICAL_TYPE(OpenedArchetype, ArchetypeType)
+    ALWAYS_CANONICAL_TYPE(ElementArchetype, ArchetypeType)
     ALWAYS_CANONICAL_TYPE(PackArchetype, ArchetypeType)
     TYPE_RANGE(Archetype, PrimaryArchetype, PackArchetype)
   TYPE(GenericTypeParam, SubstitutableType)

include/swift/AST/Types.h

@@ -125,8 +125,8 @@ class RecursiveTypeProperties {
     HasTypeVariable      = 0x01,
 
     /// This type expression contains a context-dependent archetype, either a
-    /// \c PrimaryArchetypeType, \c OpenedArchetypeType, or
-    /// \c PackArchetype.
+    /// \c PrimaryArchetypeType, \c OpenedArchetypeType,
+    /// \c ElementArchetypeType, or \c PackArchetype.
     HasArchetype         = 0x02,
 
     /// This type expression contains a GenericTypeParamType.
@@ -6040,6 +6040,52 @@ class PackArchetypeType final
 BEGIN_CAN_TYPE_WRAPPER(PackArchetypeType, ArchetypeType)
 END_CAN_TYPE_WRAPPER(PackArchetypeType, ArchetypeType)
 
+/// An archetype that represents the element type of a pack archetype.
+class ElementArchetypeType final : public ArchetypeType,
+    private ArchetypeTrailingObjects<ElementArchetypeType>
+{
+  friend TrailingObjects;
+  friend ArchetypeType;
+  friend GenericEnvironment;
+
+  UUID ID;
+
+  /// Create a new element archetype in the given environment representing
+  /// the interface type.
+  ///
+  /// This is only invoked by the generic environment when mapping the
+  /// interface type into context.
+  static CanTypeWrapper<ElementArchetypeType>
+  getNew(GenericEnvironment *environment, Type interfaceType,
+         ArrayRef<ProtocolDecl *> conformsTo, Type superclass,
+         LayoutConstraint layout);
+
+public:
+  /// Retrieve the ID number of this opened element.
+  UUID getOpenedElementID() const;
+
+  /// Return the archetype that represents the root generic parameter of its
+  /// interface type.
+  ElementArchetypeType *getRoot() const {
+    return cast<ElementArchetypeType>(ArchetypeType::getRoot());
+  }
+
+  static bool classof(const TypeBase *T) {
+    return T->getKind() == TypeKind::ElementArchetype;
+  }
+  
+private:
+  ElementArchetypeType(const ASTContext &ctx,
+                       GenericEnvironment *environment, Type interfaceType,
+                       ArrayRef<ProtocolDecl *> conformsTo, Type superclass,
+                       LayoutConstraint layout);
+};
+BEGIN_CAN_TYPE_WRAPPER(ElementArchetypeType, ArchetypeType)
+  CanElementArchetypeType getRoot() const {
+    return CanElementArchetypeType(getPointer()->getRoot());
+  }
+END_CAN_TYPE_WRAPPER(ElementArchetypeType, ArchetypeType)
+
 template<typename Type>
 const Type *ArchetypeType::getSubclassTrailingObjects() const {
   if (auto contextTy = dyn_cast<PrimaryArchetypeType>(this)) {
@@ -6054,6 +6100,9 @@ const Type *ArchetypeType::getSubclassTrailingObjects() const {
   if (auto childTy = dyn_cast<PackArchetypeType>(this)) {
     return childTy->getTrailingObjects<Type>();
   }
+  if (auto childTy = dyn_cast<ElementArchetypeType>(this)) {
+    return childTy->getTrailingObjects<Type>();
+  }
   llvm_unreachable("unhandled ArchetypeType subclass?");
 }
 

lib/AST/ASTContext.cpp

@@ -495,6 +495,7 @@ struct ASTContext::Implementation {
   llvm::FoldingSet<BuiltinVectorType> BuiltinVectorTypes;
   llvm::FoldingSet<DeclName::CompoundDeclName> CompoundNames;
   llvm::DenseMap<UUID, GenericEnvironment *> OpenedExistentialEnvironments;
+  llvm::DenseMap<UUID, GenericEnvironment *> OpenedElementEnvironments;
   llvm::FoldingSet<IndexSubset> IndexSubsets;
   llvm::FoldingSet<AutoDiffDerivativeFunctionIdentifier>
       AutoDiffDerivativeFunctionIdentifiers;
@@ -4928,13 +4929,32 @@ GenericEnvironment *GenericEnvironment::forPrimary(GenericSignature signature) {
 
   // Allocate and construct the new environment.
   unsigned numGenericParams = signature.getGenericParams().size();
-  size_t bytes = totalSizeToAlloc<OpaqueTypeDecl *, SubstitutionMap,
-                                  OpenedGenericEnvironmentData, Type>(
+  size_t bytes = totalSizeToAlloc<OpaqueEnvironmentData,
+                                  OpenedExistentialEnvironmentData,
+                                  OpenedElementEnvironmentData, Type>(
       0, 0, 0, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
   return new (mem) GenericEnvironment(signature);
 }
 
+/// Create a new generic environment for an opaque type with the given set of
+/// outer substitutions.
+GenericEnvironment *GenericEnvironment::forOpaqueType(
+    OpaqueTypeDecl *opaque, SubstitutionMap subs, AllocationArena arena) {
+  auto &ctx = opaque->getASTContext();
+
+  // Allocate and construct the new environment.
+  auto signature = opaque->getOpaqueInterfaceGenericSignature();
+  unsigned numGenericParams = signature.getGenericParams().size();
+  size_t bytes = totalSizeToAlloc<OpaqueEnvironmentData,
+                                  OpenedExistentialEnvironmentData,
+                                  OpenedElementEnvironmentData, Type>(
+      1, 0, 0, numGenericParams);
+  void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment), arena);
+  auto env = new (mem) GenericEnvironment(signature, opaque, subs);
+  return env;
+}
+
 /// Create a new generic environment for an opened archetype.
 GenericEnvironment *
 GenericEnvironment::forOpenedExistential(
@@ -4969,9 +4989,10 @@ GenericEnvironment::forOpenedExistential(
 
   // Allocate and construct the new environment.
   unsigned numGenericParams = signature.getGenericParams().size();
-  size_t bytes = totalSizeToAlloc<OpaqueTypeDecl *, SubstitutionMap,
-                                  OpenedGenericEnvironmentData, Type>(
-      0, 0, 1, numGenericParams);
+  size_t bytes = totalSizeToAlloc<OpaqueEnvironmentData,
+                                  OpenedExistentialEnvironmentData,
+                                  OpenedElementEnvironmentData, Type>(
+      0, 1, 0, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
   auto *genericEnv =
       new (mem) GenericEnvironment(signature, existential, parentSig, uuid);
@@ -4981,21 +5002,35 @@ GenericEnvironment::forOpenedExistential(
   return genericEnv;
 }
 
-/// Create a new generic environment for an opaque type with the given set of
-/// outer substitutions.
-GenericEnvironment *GenericEnvironment::forOpaqueType(
-    OpaqueTypeDecl *opaque, SubstitutionMap subs, AllocationArena arena) {
-  auto &ctx = opaque->getASTContext();
+/// Create a new generic environment for an element archetype.
+GenericEnvironment *
+GenericEnvironment::forOpenedElement(GenericSignature signature, UUID uuid) {
+  auto &ctx = signature->getASTContext();
+
+  auto &openedElementEnvironments =
+      ctx.getImpl().OpenedElementEnvironments;
+  auto found = openedElementEnvironments.find(uuid);
+
+  if (found != openedElementEnvironments.end()) {
+    auto *existingEnv = found->second;
+    assert(existingEnv->getGenericSignature().getPointer() == signature.getPointer());
+    assert(existingEnv->getOpenedElementUUID() == uuid);
+
+    return existingEnv;
+  }
 
   // Allocate and construct the new environment.
-  auto signature = opaque->getOpaqueInterfaceGenericSignature();
   unsigned numGenericParams = signature.getGenericParams().size();
-  size_t bytes = totalSizeToAlloc<OpaqueTypeDecl *, SubstitutionMap,
-                                  OpenedGenericEnvironmentData, Type>(
-      1, 1, 0, numGenericParams);
-  void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment), arena);
-  auto env = new (mem) GenericEnvironment(signature, opaque, subs);
-  return env;
+  size_t bytes = totalSizeToAlloc<OpaqueEnvironmentData,
+                                  OpenedExistentialEnvironmentData,
+                                  OpenedElementEnvironmentData, Type>(
+      0, 0, 1, numGenericParams);
+  void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
+  auto *genericEnv = new (mem) GenericEnvironment(signature, uuid);
+
+  openedElementEnvironments[uuid] = genericEnv;
+
+  return genericEnv;
 }
 
 void DeclName::CompoundDeclName::Profile(llvm::FoldingSetNodeID &id,

lib/AST/ASTDumper.cpp

@@ -3881,6 +3881,11 @@ namespace {
       OS << "\n";
       PrintWithColorRAII(OS, ParenthesisColor) << ')';
     }
+    void visitElementArchetypeType(ElementArchetypeType *T, StringRef label) {
+      printArchetypeCommon(T, "element_archetype_type", label);
+      printField("opened_element_id", T->getOpenedElementID());
+      PrintWithColorRAII(OS, ParenthesisColor) << ')';
+    }
 
     void visitGenericTypeParamType(GenericTypeParamType *T, StringRef label) {
       printCommon(label, "generic_type_param_type");

lib/AST/ASTMangler.cpp

@@ -1391,6 +1391,7 @@ void ASTMangler::appendType(Type type, GenericSignature sig,
       // type ::= archetype
     case TypeKind::PrimaryArchetype:
     case TypeKind::PackArchetype:
+    case TypeKind::ElementArchetype:
       llvm_unreachable("Cannot mangle free-standing archetypes");
 
     case TypeKind::OpenedArchetype: {

lib/AST/ASTPrinter.cpp

@@ -6370,6 +6370,19 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     }
   }
 
+  void visitElementArchetypeType(ElementArchetypeType *T) {
+    if (Options.PrintForSIL) {
+      Printer << "@element(\"" << T->getOpenedElementID() << ") ";
+      visit(T->getGenericEnvironment()->getOpenedExistentialType());
+      Printer << ") ";
+
+      auto interfaceTy = T->getInterfaceType();
+      visit(interfaceTy);
+    } else {
+      visit(T->getInterfaceType());
+    }
+  }
+
   void printDependentMember(DependentMemberType *T) {
     if (auto *const Assoc = T->getAssocType()) {
       if (Options.ProtocolQualifiedDependentMemberTypes) {