AST: ExistentialLayout support for ParametrizedProtocolType

Author: slavapestov

include/swift/AST/ExistentialLayout.h

@@ -26,6 +26,11 @@ namespace swift {
   class ProtocolType;
   class ProtocolCompositionType;
 
+struct PrimaryAssociatedTypeRequirement {
+  AssociatedTypeDecl *AssocType;
+  Type Argument;
+};
+
 struct ExistentialLayout {
   enum Kind { Class, Error, Opaque };
 
@@ -37,6 +42,7 @@ struct ExistentialLayout {
 
   ExistentialLayout(ProtocolType *type);
   ExistentialLayout(ProtocolCompositionType *type);
+  ExistentialLayout(ParametrizedProtocolType *type);
 
   /// The explicit superclass constraint, if any.
   Type explicitSuperclass;
@@ -108,6 +114,10 @@ struct ExistentialLayout {
 
   /// Zero or more protocol constraints from a ProtocolCompositionType
   ArrayRef<Type> protocols;
+
+  /// Zero or more primary associated type requirements from a
+  /// ParametrizedProtocolType
+  ArrayRef<PrimaryAssociatedTypeRequirement> sameTypeRequirements;
 };
 
 }

include/swift/AST/Types.h

@@ -5261,7 +5261,10 @@ END_CAN_TYPE_WRAPPER(ProtocolCompositionType, Type)
 /// \endcode
 class ParametrizedProtocolType final : public TypeBase,
     public llvm::FoldingSetNode {
+  friend struct ExistentialLayout;
+
   ProtocolType *Base;
+  AssociatedTypeDecl *AssocType;
   Type Arg;
 
 public:
@@ -5274,6 +5277,10 @@ class ParametrizedProtocolType final : public TypeBase,
     return Base;
   }
 
+  AssociatedTypeDecl *getAssocType() const {
+    return AssocType;
+  }
+
   Type getArgumentType() const {
     return Arg;
   }
@@ -5293,9 +5300,7 @@ class ParametrizedProtocolType final : public TypeBase,
 private:
   ParametrizedProtocolType(const ASTContext *ctx,
                            ProtocolType *base, Type arg,
-                           RecursiveTypeProperties properties)
-    : TypeBase(TypeKind::ParametrizedProtocol, /*Context=*/ctx, properties),
-      Base(base), Arg(arg) { }
+                           RecursiveTypeProperties properties);
 };
 BEGIN_CAN_TYPE_WRAPPER(ParametrizedProtocolType, Type)
   PROXY_CAN_TYPE_SIMPLE_GETTER(getBaseType)

lib/AST/Type.cpp

@@ -295,6 +295,15 @@ ExistentialLayout::ExistentialLayout(ProtocolCompositionType *type) {
   protocols = { members.data(), members.size() };
 }
 
+ExistentialLayout::ExistentialLayout(ParametrizedProtocolType *type) {
+  assert(type->isCanonical());
+
+  *this = ExistentialLayout(type->getBaseType());
+  sameTypeRequirements = {
+      reinterpret_cast<PrimaryAssociatedTypeRequirement *>(&type->AssocType),
+      1
+  };
+}
 
 ExistentialLayout TypeBase::getExistentialLayout() {
   return getCanonicalType().getExistentialLayout();
@@ -310,6 +319,9 @@ ExistentialLayout CanType::getExistentialLayout() {
   if (auto proto = dyn_cast<ProtocolType>(*this))
     return ExistentialLayout(proto);
 
+  if (auto param = dyn_cast<ParametrizedProtocolType>(*this))
+    return ExistentialLayout(param);
+
   auto comp = cast<ProtocolCompositionType>(*this);
   return ExistentialLayout(comp);
 }
@@ -3808,6 +3820,17 @@ void ProtocolCompositionType::Profile(llvm::FoldingSetNodeID &ID,
     ID.AddPointer(T.getPointer());
 }
 
+ParametrizedProtocolType::ParametrizedProtocolType(
+    const ASTContext *ctx,
+    ProtocolType *base, Type arg,
+    RecursiveTypeProperties properties)
+  : TypeBase(TypeKind::ParametrizedProtocol, /*Context=*/ctx, properties),
+    Base(base), AssocType(base->getDecl()->getPrimaryAssociatedType()),
+    Arg(arg) {
+  assert(AssocType != nullptr &&
+         "Protocol doesn't have a primary associated type");
+}
+
 void ParametrizedProtocolType::Profile(llvm::FoldingSetNodeID &ID,
                                        ProtocolType *baseTy,
                                        Type argTy) {