[Sema] TypeWrappers: Allow type wrapper inference from protocol conformances

- Infer type wrappers only from direct (declared on type) protocols
- Inferences from protocol to protocol is not allowed
- If type specifies a type wrapper attribute explicitly it has to
  match the one associated with a declared protocol, otherwise the
  declaration is going to be rejected as having multiple wrappers.

Author: xedin

include/swift/AST/DiagnosticsSema.def

@@ -6616,6 +6616,10 @@ ERROR(cannot_use_multiple_type_wrappers,none,
       "%0 %1 cannot use more than one type wrapper",
       (DescriptiveDeclKind, DeclName))
 
+NOTE(type_wrapper_inferred_from,none,
+     "type wrapper %0 inferred from %1 %2",
+     (DeclName, DescriptiveDeclKind, DeclName))
+
 ERROR(type_wrapper_requires_memberwise_init,none,
       "type wrapper type %0 does not contain a required initializer"
       " - init(for:storage:)",

lib/Sema/TypeCheckTypeWrapper.cpp

@@ -83,11 +83,18 @@ NominalTypeDecl *NominalTypeDecl::getTypeWrapper() const {
                            GetTypeWrapper{mutableSelf}, nullptr);
 }
 
-static void getTypeWrappers(
-    NominalTypeDecl *decl,
-    SmallVectorImpl<std::pair<CustomAttr *, NominalTypeDecl *>> &typeWrappers) {
+struct TypeWrapperAttrInfo {
+  CustomAttr *Attr;
+  NominalTypeDecl *Wrapper;
+  NominalTypeDecl *AttachedTo;
+};
+
+static void
+getTypeWrappers(NominalTypeDecl *decl,
+                SmallVectorImpl<TypeWrapperAttrInfo> &typeWrappers) {
   auto &ctx = decl->getASTContext();
 
+  // Attributes applied directly to the type.
   for (auto *attr : decl->getAttrs().getAttributes<CustomAttr>()) {
     auto *mutableAttr = const_cast<CustomAttr *>(attr);
     auto *nominal = evaluateOrDefault(
@@ -98,7 +105,39 @@ static void getTypeWrappers(
 
     auto *typeWrapper = nominal->getAttrs().getAttribute<TypeWrapperAttr>();
     if (typeWrapper && typeWrapper->isValid())
-      typeWrappers.push_back({mutableAttr, nominal});
+      typeWrappers.push_back({mutableAttr, nominal, decl});
+  }
+
+  // Do not allow transitive protocol inference between protocols.
+  if (isa<ProtocolDecl>(decl))
+    return;
+
+  // Attributes inferred from (explicit) protocol conformances
+  // associated with the declaration of the type.
+  for (unsigned i : indices(decl->getInherited())) {
+    auto inheritedType = evaluateOrDefault(
+        ctx.evaluator,
+        InheritedTypeRequest{decl, i, TypeResolutionStage::Interface}, Type());
+
+    if (!(inheritedType && inheritedType->isConstraintType()))
+      continue;
+
+    auto *protocol = inheritedType->getAnyNominal();
+    if (!protocol)
+      continue;
+
+    SmallVector<TypeWrapperAttrInfo, 2> inferredAttrs;
+    getTypeWrappers(protocol, inferredAttrs);
+
+    // De-duplicate inferred type wrappers. This also makes sure
+    // that if both protocol and conforming type explicitly declare
+    // the same type wrapper there is no clash between them.
+    for (const auto &inferredAttr : inferredAttrs) {
+      if (llvm::find_if(typeWrappers, [&](const TypeWrapperAttrInfo &attr) {
+            return attr.Wrapper == inferredAttr.Wrapper;
+          }) == typeWrappers.end())
+        typeWrappers.push_back(inferredAttr);
+    }
   }
 }
 
@@ -108,7 +147,7 @@ NominalTypeDecl *GetTypeWrapper::evaluate(Evaluator &evaluator,
 
   // Note that we don't actually care whether there are duplicates,
   // using the same type wrapper multiple times is still an error.
-  SmallVector<std::pair<CustomAttr *, NominalTypeDecl *>, 2> typeWrappers;
+  SmallVector<TypeWrapperAttrInfo, 2> typeWrappers;
 
   getTypeWrappers(decl, typeWrappers);
 
@@ -119,27 +158,34 @@ NominalTypeDecl *GetTypeWrapper::evaluate(Evaluator &evaluator,
     ctx.Diags.diagnose(decl, diag::cannot_use_multiple_type_wrappers,
                        decl->getDescriptiveKind(), decl->getName());
 
-    for (const auto &attr : typeWrappers) {
-      ctx.Diags.diagnose(attr.first->getLocation(), diag::decl_declared_here,
-                         attr.second->getName());
+    for (const auto &entry : typeWrappers) {
+      if (entry.AttachedTo == decl) {
+        ctx.Diags.diagnose(entry.Attr->getLocation(), diag::decl_declared_here,
+                           entry.Wrapper->getName());
+      } else {
+        ctx.Diags.diagnose(decl, diag::type_wrapper_inferred_from,
+                           entry.Wrapper->getName(),
+                           entry.AttachedTo->getDescriptiveKind(),
+                           entry.AttachedTo->getName());
+      }
     }
 
     return nullptr;
   }
 
-  return typeWrappers.front().second;
+  return typeWrappers.front().Wrapper;
 }
 
 Type GetTypeWrapperType::evaluate(Evaluator &evaluator,
                                   NominalTypeDecl *decl) const {
-  SmallVector<std::pair<CustomAttr *, NominalTypeDecl *>, 2> typeWrappers;
+  SmallVector<TypeWrapperAttrInfo, 2> typeWrappers;
 
   getTypeWrappers(decl, typeWrappers);
 
   if (typeWrappers.size() != 1)
     return Type();
 
-  auto *typeWrapperAttr = typeWrappers.front().first;
+  auto *typeWrapperAttr = typeWrappers.front().Attr;
   auto type = evaluateOrDefault(
       evaluator,
       CustomAttrTypeRequest{typeWrapperAttr, decl->getDeclContext(),

test/Interpreter/Inputs/type_wrapper_defs.swift

@@ -394,3 +394,10 @@ public class TypeWithSomeDefaultedLetProperties<T> {
     print(self.c)
   }
 }
+
+@Wrapper
+public protocol WrappedProtocol {
+  associatedtype T : RangeReplaceableCollection
+
+  var v: T { get set }
+}

test/Interpreter/type_wrappers.swift

@@ -666,3 +666,26 @@ do {
   // CHECK-NEXT: in getter storage: \$Storage.b
   // CHECK-NEXT: 42
 }
+
+// Type wrapper inference from protocols
+do {
+  struct Test<T> : WrappedProtocol {
+    var v: [T?]
+
+    init(_ value: T? = nil) {
+      v = [value]
+    }
+  }
+
+  func run_test<T: WrappedProtocol>(_ test: T) {
+    print(test.v)
+    print(test.v.count)
+  }
+
+  run_test(Test(Optional("Hello")))
+  // CHECK: Wrapper.init(for: Test<String>, storage: $Storage(v: [Optional("Hello")]))
+  // CHECK-NEXT: in getter storage: \$Storage.v
+  // CHECK-NEXT: [Optional("Hello")]
+  // CHECK-NEXT: in getter storage: \$Storage.v
+  // CHECK-NEXT: 1
+}

test/Serialization/AllowErrors/invalid-inheritance.swift

@@ -32,9 +32,8 @@ extension undefined: SomeProto {} // expected-error {{cannot find type 'undefine
 
 public extension undefined { // expected-error {{cannot find type 'undefined' in scope}}
   protocol SomeProtoInner: undefined {} // expected-error {{cannot find type 'undefined' in scope}}
-  // TODO: Why don't these have errors?
-  class SomeClassInner: undefined {}
-  struct SomeStructInner: undefined {}
+  class SomeClassInner: undefined {} // expected-error {{cannot find type 'undefined' in scope}}
+  struct SomeStructInner: undefined {} // expected-error {{cannot find type 'undefined' in scope}}
   enum SomeEnumInner: undefined { // expected-error {{cannot find type 'undefined' in scope}}
     case a
   }

test/type/type_wrapper.swift

@@ -653,3 +653,37 @@ func test_multiple_wrapper_attributes() {
   // expected-note@-2 {{'NoopWrapper' declared here}}
   class Test2 {} // expected-error {{class 'Test2' cannot use more than one type wrapper}}
 }
+
+@NoopWrapper
+protocol WrappedProto {
+}
+
+do {
+  // @NoopWrapper is inferred from `WrappedProto`
+  struct InferenceFromProto : WrappedProto { // Ok
+    var x: Int
+
+    func test() {
+      _ = $storage // Ok (to make sure that NoopWrapper is indeed inferred)
+    }
+  }
+
+  @Parent.Wrapper // expected-note {{'Wrapper' declared here}}
+  struct ClashBetweenDirectAndInferred : WrappedProto {
+    // expected-error@-1 {{struct 'ClashBetweenDirectAndInferred' cannot use more than one type wrapper}}
+    // expected-note@-2 {{type wrapper 'NoopWrapper' inferred from protocol 'WrappedProto'}}
+    // expected-error@-3 {{type 'ClashBetweenDirectAndInferred' does not conform to protocol 'WrappedProto'}}
+  }
+
+  @NoopWrapper
+  final class NoClashDueToSameWrapper : WrappedProto { // Ok
+    var v: [Int?]
+
+    func test() {
+      let storage: NoopWrapper<NoClashDueToSameWrapper, $Storage> = $storage
+      _ = storage[propertyKeyPath: \.v, storageKeyPath: \$Storage.v]
+    }
+  }
+
+
+}