[AST/Sema] TypeWrappers: Implement type wrappers on protocols

Adding a type wrapper attribute on a protocol does two things:

- Synthesizes `associatedtype $Storage` declaration with `internal` access
- Synthesizes `var $storage: <#Wrapper#><Self, Self.$Storage>` value requirement

Author: xedin

include/swift/AST/Decl.h

@@ -3909,8 +3909,9 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
 
   /// If this declaration has a type wrapper, return `$Storage`
   /// declaration that contains all the stored properties managed
-  /// by the wrapper.
-  NominalTypeDecl *getTypeWrapperStorageDecl() const;
+  /// by the wrapper. Note that if this type is a protocol them
+  /// this method returns an associated type for $Storage.
+  TypeDecl *getTypeWrapperStorageDecl() const;
 
   /// If this declaration is a type wrapper, retrieve
   /// its required initializer - `init(storageWrapper:)`.

include/swift/AST/DiagnosticsSema.def

@@ -6605,7 +6605,7 @@ ERROR(type_wrappers_are_experimental,none,
       "type wrappers are an experimental feature", ())
 
 ERROR(type_wrapper_attribute_not_allowed_here,none,
-      "type wrapper attribute %0 can only be applied to a class, struct",
+      "type wrapper attribute %0 can only be applied to a class, struct, or protocol",
       (Identifier))
 
 ERROR(type_wrapper_requires_two_generic_params,none,

include/swift/AST/TypeCheckRequests.h

@@ -3552,16 +3552,15 @@ class GetTypeWrapperType
 /// Inject or get `$Storage` type which has all of the stored properties
 /// of the given type with a type wrapper.
 class GetTypeWrapperStorage
-    : public SimpleRequest<GetTypeWrapperStorage,
-                           NominalTypeDecl *(NominalTypeDecl *),
+    : public SimpleRequest<GetTypeWrapperStorage, TypeDecl *(NominalTypeDecl *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  NominalTypeDecl *evaluate(Evaluator &evaluator, NominalTypeDecl *) const;
+  TypeDecl *evaluate(Evaluator &evaluator, NominalTypeDecl *) const;
 
 public:
   bool isCached() const { return true; }

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -408,7 +408,7 @@ SWIFT_REQUEST(TypeChecker, GetTypeWrapperType,
               Type(NominalTypeDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetTypeWrapperStorage,
-              NominalTypeDecl *(NominalTypeDecl *),
+              TypeDecl *(NominalTypeDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetTypeWrapperProperty,
               VarDecl *(NominalTypeDecl *),

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -480,7 +480,8 @@ bool DIMemoryObjectInfo::isElementLetProperty(unsigned Element) const {
     assert(Element < storageVarType->getNumElements());
     auto propertyName = storageVarType->getElement(Element).getName();
 
-    auto *storageDecl = wrappedType->getTypeWrapperStorageDecl();
+    auto *storageDecl =
+        cast<NominalTypeDecl>(wrappedType->getTypeWrapperStorageDecl());
     auto *property = storageDecl->lookupDirect(propertyName).front();
     return cast<VarDecl>(property)->isLet();
   }

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -1093,7 +1093,8 @@ void LifetimeChecker::injectTypeWrapperStorageInitalization() {
 
   auto *ctor = cast<ConstructorDecl>(storageVar->getDeclContext()->getAsDecl());
   auto *parentType = ctor->getDeclContext()->getSelfNominalTypeDecl();
-  auto *storageDecl = parentType->getTypeWrapperStorageDecl();
+  auto *storageDecl =
+      cast<NominalTypeDecl>(parentType->getTypeWrapperStorageDecl());
 
   for (auto *point : points) {
     SILBuilderWithScope::insertAfter(point, [&](SILBuilder &b) {

lib/Sema/CodeSynthesis.cpp

@@ -1646,6 +1646,9 @@ synthesizeTypeWrappedTypeStorageWrapperInitializerBody(
 
 ConstructorDecl *SynthesizeTypeWrappedTypeStorageWrapperInitializer::evaluate(
     Evaluator &evaluator, NominalTypeDecl *wrappedType) const {
+  if (isa<ProtocolDecl>(wrappedType))
+    return nullptr;
+
   if (!wrappedType->hasTypeWrapper())
     return nullptr;
 
@@ -1673,8 +1676,11 @@ synthesizeTypeWrappedTypeMemberwiseInitializerBody(AbstractFunctionDecl *decl,
   auto &ctx = ctor->getASTContext();
   auto *parent = ctor->getDeclContext()->getSelfNominalTypeDecl();
 
+  assert(!isa<ProtocolDecl>(parent));
+
   // self.$storage = .init(storage: $Storage(...))
-  auto *storageType = parent->getTypeWrapperStorageDecl();
+  auto *storageType =
+      cast<NominalTypeDecl>(parent->getTypeWrapperStorageDecl());
   assert(storageType);
 
   auto *typeWrapperVar = parent->getTypeWrapperProperty();
@@ -1775,6 +1781,9 @@ synthesizeTypeWrappedTypeMemberwiseInitializerBody(AbstractFunctionDecl *decl,
 
 ConstructorDecl *SynthesizeTypeWrappedTypeMemberwiseInitializer::evaluate(
     Evaluator &evaluator, NominalTypeDecl *wrappedType) const {
+  if (isa<ProtocolDecl>(wrappedType))
+    return nullptr;
+
   if (!wrappedType->hasTypeWrapper())
     return nullptr;
 

lib/Sema/TypeCheckAttr.cpp

@@ -3597,7 +3597,7 @@ void AttributeChecker::visitCustomAttr(CustomAttr *attr) {
   }
 
   if (nominal->getAttrs().hasAttribute<TypeWrapperAttr>()) {
-    if (!(isa<ClassDecl>(D) || isa<StructDecl>(D))) {
+    if (!(isa<ClassDecl>(D) || isa<StructDecl>(D) || isa<ProtocolDecl>(D))) {
       diagnose(attr->getLocation(),
                diag::type_wrapper_attribute_not_allowed_here,
                nominal->getName());

lib/Sema/TypeCheckStorage.cpp

@@ -3424,6 +3424,12 @@ StorageImplInfoRequest::evaluate(Evaluator &evaluator,
       readWriteImpl = ReadWriteImplKind::MaterializeToTemporary;
     } else if (storage->getParsedAccessor(AccessorKind::Get)) {
       readImpl = ReadImplKind::Get;
+    } else if (storage->getName() ==
+               storage->getASTContext().Id_TypeWrapperProperty) {
+      // Type wrapper `$storage` property is `get set` requirement.
+      readImpl = ReadImplKind::Get;
+      writeImpl = WriteImplKind::Set;
+      readWriteImpl = ReadWriteImplKind::Modify;
     }
 
     StorageImplInfo info(readImpl, writeImpl, readWriteImpl);

lib/Sema/TypeCheckTypeWrapper.cpp

@@ -173,25 +173,36 @@ NominalTypeDecl::getTypeWrappedTypeMemberwiseInitializer() const {
       SynthesizeTypeWrappedTypeMemberwiseInitializer{mutableSelf}, nullptr);
 }
 
-NominalTypeDecl *NominalTypeDecl::getTypeWrapperStorageDecl() const {
+TypeDecl *NominalTypeDecl::getTypeWrapperStorageDecl() const {
   auto *mutableSelf = const_cast<NominalTypeDecl *>(this);
   return evaluateOrDefault(getASTContext().evaluator,
                            GetTypeWrapperStorage{mutableSelf}, nullptr);
 }
 
-NominalTypeDecl *
-GetTypeWrapperStorage::evaluate(Evaluator &evaluator,
-                                NominalTypeDecl *parent) const {
+TypeDecl *GetTypeWrapperStorage::evaluate(Evaluator &evaluator,
+                                          NominalTypeDecl *parent) const {
   if (!parent->hasTypeWrapper())
     return nullptr;
 
   auto &ctx = parent->getASTContext();
 
-  auto *storage =
-      new (ctx) StructDecl(/*StructLoc=*/SourceLoc(), ctx.Id_TypeWrapperStorage,
-                           /*NameLoc=*/SourceLoc(),
-                           /*Inheritted=*/{},
-                           /*GenericParams=*/nullptr, parent);
+  TypeDecl *storage = nullptr;
+  if (isa<ProtocolDecl>(parent)) {
+    // If type wrapper is associated with a protocol, we need to
+    // inject a new associated type - $Storage.
+    storage = new (ctx)
+        AssociatedTypeDecl(parent, /*keywordLoc=*/SourceLoc(),
+                           ctx.Id_TypeWrapperStorage, /*nameLoc=*/SourceLoc(),
+                           /*defaultDefinition=*/nullptr,
+                           /*trailingWhere=*/nullptr);
+  } else {
+    // For classes and structs we inject a new member struct - $Storage.
+    storage = new (ctx)
+        StructDecl(/*StructLoc=*/SourceLoc(), ctx.Id_TypeWrapperStorage,
+                   /*NameLoc=*/SourceLoc(),
+                   /*Inheritted=*/{},
+                   /*GenericParams=*/nullptr, parent);
+  }
 
   storage->setImplicit();
   storage->setSynthesized();
@@ -223,7 +234,7 @@ GetTypeWrapperProperty::evaluate(Evaluator &evaluator,
   auto propertyTy = BoundGenericType::get(
       typeWrapper, /*Parent=*/typeWrapperType->getParent(),
       /*genericArgs=*/
-      {parent->getDeclaredInterfaceType(),
+      {parent->getSelfInterfaceType(),
        storage->getDeclaredInterfaceType()});
 
   return injectProperty(parent, ctx.Id_TypeWrapperProperty, propertyTy,
@@ -240,7 +251,10 @@ VarDecl *GetTypeWrapperStorageForProperty::evaluate(Evaluator &evaluator,
   if (!property->isAccessedViaTypeWrapper())
     return nullptr;
 
-  auto *storage = wrappedType->getTypeWrapperStorageDecl();
+  assert(!isa<ProtocolDecl>(wrappedType));
+
+  auto *storage =
+      cast<NominalTypeDecl>(wrappedType->getTypeWrapperStorageDecl());
   assert(storage);
 
   // Type wrapper variables are never initialized directly,
@@ -456,7 +470,11 @@ VarDecl *SynthesizeLocalVariableForTypeWrapperStorage::evaluate(
   if (!(DC && DC->hasTypeWrapper()))
     return nullptr;
 
-  auto *storageDecl = DC->getTypeWrapperStorageDecl();
+  // Default protocol initializers do not get transformed.
+  if (isa<ProtocolDecl>(DC))
+    return nullptr;
+
+  auto *storageDecl = cast<NominalTypeDecl>(DC->getTypeWrapperStorageDecl());
   assert(storageDecl);
 
   SmallVector<TupleTypeElt, 4> members;