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

include/swift/Basic/LangOptions.h

@@ -317,6 +317,10 @@ namespace swift {
     /// keyword.
     bool EnableExplicitExistentialTypes = true;
 
+    /// Enable support for protocol types parametrized by primary
+    /// associated type.
+    bool EnableParametrizedProtocolTypes = false;
+
     /// Enable experimental flow-sensitive concurrent captures.
     bool EnableExperimentalFlowSensitiveConcurrentCaptures = false;
 

include/swift/Option/FrontendOptions.td

@@ -506,6 +506,10 @@ def enable_explicit_existential_types :
   Flag<["-"], "enable-explicit-existential-types">,
   HelpText<"Enable experimental support for explicit existential types">;
 
+def enable_parametrized_protocol_types :
+  Flag<["-"], "enable-parametrized-protocol-types">,
+  HelpText<"Enable experimental support for primary associated types and parametrized protocols">;
+
 def enable_deserialization_recovery :
   Flag<["-"], "enable-deserialization-recovery">,
   HelpText<"Attempt to recover from missing xrefs (etc) in swiftmodules">;

lib/AST/ASTContext.cpp

@@ -405,6 +405,7 @@ struct ASTContext::Implementation {
     llvm::FoldingSet<UnboundGenericType> UnboundGenericTypes;
     llvm::FoldingSet<BoundGenericType> BoundGenericTypes;
     llvm::FoldingSet<ProtocolCompositionType> ProtocolCompositionTypes;
+    llvm::FoldingSet<ParametrizedProtocolType> ParametrizedProtocolTypes;
     llvm::FoldingSet<LayoutConstraintInfo> LayoutConstraints;
     llvm::DenseMap<std::pair<OpaqueTypeDecl *, SubstitutionMap>,
                    GenericEnvironment *> OpaqueArchetypeEnvironments;
@@ -3240,10 +3241,37 @@ ProtocolCompositionType::build(const ASTContext &C, ArrayRef<Type> Members,
                                                   Members,
                                                   HasExplicitAnyObject,
                                                   properties);
-  C.getImpl().getArena(arena).ProtocolCompositionTypes.InsertNode(compTy, InsertPos);
+  C.getImpl().getArena(arena).ProtocolCompositionTypes.InsertNode(
+      compTy, InsertPos);
   return compTy;
 }
 
+Type ParametrizedProtocolType::get(const ASTContext &C,
+                                   ProtocolType *baseTy,
+                                   Type argTy) {
+  bool isCanonical = baseTy->isCanonical();
+  RecursiveTypeProperties properties = baseTy->getRecursiveProperties();
+  properties |= argTy->getRecursiveProperties();
+  isCanonical &= argTy->isCanonical();
+
+  auto arena = getArena(properties);
+
+  void *InsertPos = nullptr;
+  llvm::FoldingSetNodeID ID;
+  ParametrizedProtocolType::Profile(ID, baseTy, argTy);
+
+  if (auto paramTy
+      = C.getImpl().getArena(arena).ParametrizedProtocolTypes
+          .FindNodeOrInsertPos(ID, InsertPos))
+    return paramTy;
+
+  auto paramTy = new (C, arena) ParametrizedProtocolType(
+        isCanonical ? &C : nullptr, baseTy, argTy, properties);
+  C.getImpl().getArena(arena).ParametrizedProtocolTypes.InsertNode(
+      paramTy, InsertPos);
+  return paramTy;
+}
+
 ReferenceStorageType *ReferenceStorageType::get(Type T,
                                                 ReferenceOwnership ownership,
                                                 const ASTContext &C) {

lib/AST/ASTDumper.cpp

@@ -3942,6 +3942,14 @@ namespace {
       PrintWithColorRAII(OS, ParenthesisColor) << ')';
     }
 
+    void visitParametrizedProtocolType(ParametrizedProtocolType *T,
+                                       StringRef label) {
+      printCommon(label, "parametrized_protocol_type");
+      printRec("base", T->getBaseType());
+      printRec("arg", T->getArgumentType());
+      PrintWithColorRAII(OS, ParenthesisColor) << ')';
+    }
+
     void visitExistentialType(ExistentialType *T,
                               StringRef label) {
       printCommon(label, "existential_type");

lib/AST/ASTMangler.cpp

@@ -1266,6 +1266,11 @@ void ASTMangler::appendType(Type type, GenericSignature sig,
       return appendExistentialLayout(layout, sig, forDecl);
     }
 
+    case TypeKind::ParametrizedProtocol: {
+      llvm::errs() << "Not implemented\n";
+      abort();
+    }
+
     case TypeKind::Existential: {
       auto constraint = cast<ExistentialType>(tybase)->getConstraintType();
       return appendType(constraint, sig, forDecl);

lib/AST/ASTPrinter.cpp

@@ -5877,6 +5877,13 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     }
   }
 
+  void visitParametrizedProtocolType(ParametrizedProtocolType *T) {
+    visit(T->getBaseType());
+    Printer << "<";
+    visit(T->getArgumentType());
+    Printer << ">";
+  }
+
   void visitExistentialType(ExistentialType *T) {
     if (Options.PrintExplicitAny)
       Printer << "any ";

lib/AST/Decl.cpp

@@ -5284,6 +5284,12 @@ bool ProtocolDecl::existentialRequiresAny() const {
     ExistentialRequiresAnyRequest{const_cast<ProtocolDecl *>(this)}, true);
 }
 
+AssociatedTypeDecl *ProtocolDecl::getPrimaryAssociatedType() const {
+  return evaluateOrDefault(getASTContext().evaluator,
+    PrimaryAssociatedTypeRequest{const_cast<ProtocolDecl *>(this)},
+    nullptr);
+}
+
 StringRef ProtocolDecl::getObjCRuntimeName(
                           llvm::SmallVectorImpl<char> &buffer) const {
   // If there is an 'objc' attribute with a name, use that name.

lib/AST/GenericSignatureBuilder.cpp

@@ -4502,6 +4502,7 @@ ConstraintResult GenericSignatureBuilder::addTypeRequirement(
   // Check whether we have a reasonable constraint type at all.
   if (!constraintType->is<ProtocolType>() &&
       !constraintType->is<ProtocolCompositionType>() &&
+      !constraintType->is<ParametrizedProtocolType>() &&
       !constraintType->getClassOrBoundGenericClass()) {
     if (source.getLoc().isValid() && !constraintType->hasError()) {
       Impl->HadAnyError = true;
@@ -4514,8 +4515,30 @@ ConstraintResult GenericSignatureBuilder::addTypeRequirement(
     return ConstraintResult::Conflicting;
   }
 
+  // Parametrized protocol requirements.
+  if (auto *paramProtoType = constraintType->getAs<ParametrizedProtocolType>()) {
+    bool anyErrors = false;
+
+    auto *protoDecl = paramProtoType->getBaseType()->getDecl();
+
+    if (isErrorResult(addConformanceRequirement(resolvedSubject, protoDecl,
+                                                source)))
+        anyErrors = true;
+
+    auto *assocType = paramProtoType->getAssocType();
+    auto depType = DependentMemberType::get(
+        resolvedSubject.getDependentType(*this), assocType);
+    if (isErrorResult(addSameTypeRequirement(Type(depType),
+                                             paramProtoType->getArgumentType(),
+                                             source,
+                                             UnresolvedHandlingKind::GenerateConstraints)))
+      anyErrors = true;
+
+    return anyErrors ? ConstraintResult::Conflicting
+                     : ConstraintResult::Resolved;
+
   // Protocol requirements.
-  if (constraintType->isExistentialType()) {
+  } else if (constraintType->isExistentialType()) {
     bool anyErrors = false;
     auto layout = constraintType->getExistentialLayout();