Merge pull request #40714 from slavapestov/parametrized-protocol-type

Parametrized protocol types

Author: slavapestov

include/swift/AST/Attr.def

@@ -657,7 +657,10 @@ CONTEXTUAL_SIMPLE_DECL_ATTR(distributed, DistributedActor,
   APIBreakingToAdd | APIBreakingToRemove,
   118)
 
-// 119 is unused
+SIMPLE_DECL_ATTR(_primaryAssociatedType,
+  PrimaryAssociatedType, OnAssociatedType | UserInaccessible |
+  APIStableToAdd | ABIStableToAdd | APIBreakingToRemove | ABIStableToRemove,
+  119)
 
 SIMPLE_DECL_ATTR(_assemblyVision, EmitAssemblyVisionRemarks,
   OnFunc | UserInaccessible | NotSerialized | OnNominalType |

include/swift/AST/Decl.h

@@ -4372,6 +4372,11 @@ class ProtocolDecl final : public NominalTypeDecl {
   /// a protocol having nested types (ObjC protocols).
   ArrayRef<AssociatedTypeDecl *> getAssociatedTypeMembers() const;
 
+  /// Returns the primary associated type, or nullptr if there isn't one. This is
+  /// the associated type that is parametrized with a same-type requirement in a
+  /// parametrized protocol type of the form SomeProtocol<SomeArgType>.
+  AssociatedTypeDecl *getPrimaryAssociatedType() const;
+
   /// Returns a protocol requirement with the given name, or nullptr if the
   /// name has multiple overloads, or no overloads at all.
   ValueDecl *getSingleRequirement(DeclName name) const;

include/swift/AST/DiagnosticsSema.def

@@ -2773,6 +2773,8 @@ ERROR(inheritance_from_non_protocol_or_class,none,
   "inheritance from non-protocol, non-class type %0", (Type))
 ERROR(inheritance_from_non_protocol,none,
   "inheritance from non-protocol type %0", (Type))
+ERROR(inheritance_from_parametrized_protocol,none,
+  "cannot inherit from parametrized protocol type %0", (Type))
 ERROR(superclass_not_first,none,
   "superclass %0 must appear first in the inheritance clause", (Type))
 ERROR(superclass_not_open,none,
@@ -3714,6 +3716,12 @@ ERROR(not_a_generic_definition,none,
       "cannot specialize a non-generic definition", ())
 ERROR(not_a_generic_type,none,
       "cannot specialize non-generic type %0", (Type))
+ERROR(parametrized_protocol_not_supported,none,
+      "protocol type with generic argument can only be used as a generic constraint", ())
+ERROR(protocol_does_not_have_primary_assoc_type,none,
+      "cannot specialize protocol type %0", (Type))
+ERROR(protocol_cannot_have_multiple_generic_arguments,none,
+      "protocol type %0 can only be specialized with exactly one argument", (Type))
 ERROR(cannot_specialize_self,none,
       "cannot specialize 'Self'", ())
 NOTE(specialize_explicit_type_instead,none,

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/TypeCheckRequests.h

@@ -313,6 +313,26 @@ class ExistentialRequiresAnyRequest :
   void cacheResult(bool value) const;
 };
 
+/// Find the primary associated type of the given protocol.
+class PrimaryAssociatedTypeRequest :
+    public SimpleRequest<PrimaryAssociatedTypeRequest,
+                         AssociatedTypeDecl *(ProtocolDecl *),
+                         RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  AssociatedTypeDecl *
+  evaluate(Evaluator &evaluator, ProtocolDecl *decl) const;
+
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
+
 class PolymorphicEffectRequirementsRequest :
     public SimpleRequest<PolymorphicEffectRequirementsRequest,
                          PolymorphicEffectRequirementList(EffectKind, ProtocolDecl *),

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -233,6 +233,9 @@ SWIFT_REQUEST(TypeChecker, PropertyWrapperTypeInfoRequest,
               NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ProtocolRequiresClassRequest, bool(ProtocolDecl *),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PrimaryAssociatedTypeRequest,
+              AssociatedTypeDecl *(ProtocolDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, RequirementRequest,
               Requirement(WhereClauseOwner, unsigned, TypeResolutionStage),
               Cached, HasNearestLocation)

include/swift/AST/TypeDifferenceVisitor.h

@@ -339,6 +339,15 @@ class CanTypeDifferenceVisitor : public CanTypePairVisitor<Impl, bool> {
                                type1->getMembers(), type2->getMembers());
   }
 
+  bool visitParametrizedProtocolType(CanParametrizedProtocolType type1,
+                                     CanParametrizedProtocolType type2) {
+    if (asImpl().visit(type1.getBaseType(), type2.getBaseType()))
+      return true;
+
+    return asImpl().visit(type1.getArgumentType(),
+                          type2.getArgumentType());
+  }
+
   bool visitExistentialType(CanExistentialType type1,
                             CanExistentialType type2) {
     return asImpl().visit(type1.getConstraintType(),

include/swift/AST/TypeMatcher.h

@@ -305,6 +305,27 @@ class TypeMatcher {
     TRIVIAL_CASE(SILBoxType)
     TRIVIAL_CASE(ProtocolCompositionType)
 
+    bool visitParametrizedProtocolType(CanParametrizedProtocolType firstParametrizedProto,
+                                       Type secondType,
+                                       Type sugaredFirstType) {
+      if (auto secondParametrizedProto = secondType->getAs<ParametrizedProtocolType>()) {
+        if (!this->visit(firstParametrizedProto.getBaseType(),
+                         secondParametrizedProto->getBaseType(),
+                         sugaredFirstType->castTo<ParametrizedProtocolType>()
+                             ->getBaseType())) {
+          return false;
+        }
+
+        return this->visit(firstParametrizedProto.getArgumentType(),
+                           secondParametrizedProto->getArgumentType(),
+                           sugaredFirstType->castTo<ParametrizedProtocolType>()
+                               ->getArgumentType());
+      }
+
+      return mismatch(firstParametrizedProto.getPointer(), secondType,
+                      sugaredFirstType);
+    }
+
     bool visitExistentialType(CanExistentialType firstExistential,
                               Type secondType,
                               Type sugaredFirstType) {

include/swift/AST/TypeNodes.def

@@ -161,6 +161,7 @@ ARTIFICIAL_TYPE(SILBlockStorage, Type)
 ARTIFICIAL_TYPE(SILBox, Type)
 ARTIFICIAL_TYPE(SILToken, Type)
 TYPE(ProtocolComposition, Type)
+TYPE(ParametrizedProtocol, Type)
 TYPE(Existential, Type)
 TYPE(LValue, Type)
 TYPE(InOut, Type)

include/swift/AST/Types.h

@@ -5243,6 +5243,70 @@ class ProtocolCompositionType final : public TypeBase,
 BEGIN_CAN_TYPE_WRAPPER(ProtocolCompositionType, Type)
 END_CAN_TYPE_WRAPPER(ProtocolCompositionType, Type)
 
+/// ParametrizedProtocolType - A type that constrains the primary associated
+/// type of a protocol to an argument type.
+///
+/// Written like a bound generic type, eg Sequence<Int>.
+///
+/// For now, these are only supported in generic requirement-like contexts:
+/// - Inheritance clauses of protocols, generic parameters, associated types
+/// - Conformance requirements in where clauses
+/// - Extensions
+///
+/// Assuming that the primary associated type of Sequence is Element, the
+/// desugaring is that T : Sequence<Int> is equivalent to
+///
+/// \code
+/// T : Sequence where T.Element == Int.
+/// \endcode
+class ParametrizedProtocolType final : public TypeBase,
+    public llvm::FoldingSetNode {
+  friend struct ExistentialLayout;
+
+  ProtocolType *Base;
+  AssociatedTypeDecl *AssocType;
+  Type Arg;
+
+public:
+  /// Retrieve an instance of a protocol composition type with the
+  /// given set of members.
+  static Type get(const ASTContext &C, ProtocolType *base,
+                  Type arg);
+
+  ProtocolType *getBaseType() const {
+    return Base;
+  }
+
+  AssociatedTypeDecl *getAssocType() const {
+    return AssocType;
+  }
+
+  Type getArgumentType() const {
+    return Arg;
+  }
+
+  void Profile(llvm::FoldingSetNodeID &ID) {
+    Profile(ID, Base, Arg);
+  }
+  static void Profile(llvm::FoldingSetNodeID &ID,
+                      ProtocolType *base,
+                      Type arg);
+
+  // Implement isa/cast/dyncast/etc.
+  static bool classof(const TypeBase *T) {
+    return T->getKind() == TypeKind::ParametrizedProtocol;
+  }
+  
+private:
+  ParametrizedProtocolType(const ASTContext *ctx,
+                           ProtocolType *base, Type arg,
+                           RecursiveTypeProperties properties);
+};
+BEGIN_CAN_TYPE_WRAPPER(ParametrizedProtocolType, Type)
+  PROXY_CAN_TYPE_SIMPLE_GETTER(getBaseType)
+  PROXY_CAN_TYPE_SIMPLE_GETTER(getArgumentType)
+END_CAN_TYPE_WRAPPER(ParametrizedProtocolType, Type)
+
 /// An existential type, spelled with \c any .
 ///
 /// In Swift 5 mode, a plain protocol name in type