Merge pull request #72234 from DougGregor/attr-implements-assoc-failure-type

[Associated type inference] Support `@_implements` on type witnesses and use it for async sequence `Failure`

Author: DougGregor

include/swift/Basic/Features.def

@@ -169,6 +169,7 @@ SUPPRESSIBLE_LANGUAGE_FEATURE(Extern, 0, "@_extern")
 LANGUAGE_FEATURE(ExpressionMacroDefaultArguments, 422, "Expression macro as caller-side default argument")
 LANGUAGE_FEATURE(BuiltinStoreRaw, 0, "Builtin.storeRaw")
 LANGUAGE_FEATURE(BuiltinCreateTask, 0, "Builtin.createTask and Builtin.createDiscardingTask")
+SUPPRESSIBLE_LANGUAGE_FEATURE(AssociatedTypeImplements, 0, "@_implements on associated types")
 
 // Swift 6
 UPCOMING_FEATURE(ConciseMagicFile, 274, 6)

lib/AST/ASTPrinter.cpp

@@ -3083,6 +3083,14 @@ static void suppressingFeatureIsolatedAny(PrintOptions &options,
   action();
 }
 
+static void suppressingFeatureAssociatedTypeImplements(PrintOptions &options,
+                                     llvm::function_ref<void()> action) {
+  unsigned originalExcludeAttrCount = options.ExcludeAttrList.size();
+  options.ExcludeAttrList.push_back(DeclAttrKind::Implements);
+  action();
+  options.ExcludeAttrList.resize(originalExcludeAttrCount);
+}
+
 /// Suppress the printing of a particular feature.
 static void suppressingFeature(PrintOptions &options, Feature feature,
                                llvm::function_ref<void()> action) {
@@ -5512,32 +5520,6 @@ void Decl::printInherited(ASTPrinter &Printer, const PrintOptions &Opts) const {
   printer.printInherited(this);
 }
 
-/// Determine whether this typealias is an inferred typealias "Failure" that
-/// would conflict with another entity named failure in the same type.
-static bool isConflictingFailureTypeWitness(
-    const TypeAliasDecl *typealias) {
-  if (!typealias->isImplicit())
-    return false;
-
-  ASTContext &ctx = typealias->getASTContext();
-  if (typealias->getName() != ctx.Id_Failure)
-    return false;
-
-  auto nominal = typealias->getDeclContext()->getSelfNominalTypeDecl();
-  if (!nominal)
-    return false;
-
-  // Look for another entity with the same name.
-  auto lookupResults = nominal->lookupDirect(
-      typealias->getName(), typealias->getLoc());
-  for (auto found : lookupResults) {
-    if (found != typealias)
-      return true;
-  }
-
-  return false;
-}
-
 bool Decl::shouldPrintInContext(const PrintOptions &PO) const {
   // Skip getters/setters. They are part of the variable or subscript.
   if (isa<AccessorDecl>(this))
@@ -5577,14 +5559,6 @@ bool Decl::shouldPrintInContext(const PrintOptions &PO) const {
     return PO.PrintIfConfig;
   }
 
-  // Prior to Swift 6, we shouldn't print the inferred associated type
-  // witness for AsyncSequence.Failure. It is always determined from the
-  // AsyncIteratorProtocol witness.
-  if (auto typealias = dyn_cast<TypeAliasDecl>(this)) {
-    if (isConflictingFailureTypeWitness(typealias))
-      return false;
-  }
-
   // Print everything else.
   return true;
 }

lib/AST/FeatureSet.cpp

@@ -342,6 +342,10 @@ static bool usesFeatureExtern(Decl *decl) {
   return decl->getAttrs().hasAttribute<ExternAttr>();
 }
 
+static bool usesFeatureAssociatedTypeImplements(Decl *decl) {
+  return isa<TypeDecl>(decl) && decl->getAttrs().hasAttribute<ImplementsAttr>();
+}
+
 static bool usesFeatureExpressionMacroDefaultArguments(Decl *decl) {
   if (auto func = dyn_cast<AbstractFunctionDecl>(decl)) {
     for (auto param : *func->getParameters()) {

lib/Sema/AssociatedTypeInference.cpp

@@ -231,6 +231,17 @@ static bool containsConcreteDependentMemberType(Type ty) {
   });
 }
 
+/// Determine whether this is the AsyncIteratorProtocol.Failure or
+/// AsyncSequence.Failure associated type.
+static bool isAsyncIteratorOrSequenceFailure(AssociatedTypeDecl *assocType) {
+  auto proto = assocType->getProtocol();
+  if (!proto->isSpecificProtocol(KnownProtocolKind::AsyncIteratorProtocol) &&
+      !proto->isSpecificProtocol(KnownProtocolKind::AsyncSequence))
+    return false;
+
+  return assocType->getName() == assocType->getASTContext().Id_Failure;
+}
+
 static void recordTypeWitness(NormalProtocolConformance *conformance,
                               AssociatedTypeDecl *assocType,
                               Type type,
@@ -254,14 +265,40 @@ static void recordTypeWitness(NormalProtocolConformance *conformance,
 
   // If there was no type declaration, synthesize one.
   if (typeDecl == nullptr) {
+    Identifier name;
+    bool needsImplementsAttr;
+    if (isAsyncIteratorOrSequenceFailure(assocType)) {
+      // Use __<protocol>_<assocType> as the name, to keep it out of the
+      // way of other names.
+      llvm::SmallString<32> nameBuffer;
+      nameBuffer += "__";
+      nameBuffer += assocType->getProtocol()->getName().str();
+      nameBuffer += "_";
+      nameBuffer += assocType->getName().str();
+
+      name = ctx.getIdentifier(nameBuffer);
+      needsImplementsAttr = true;
+    } else {
+      // Declare a typealias with the same name as the associated type.
+      name = assocType->getName();
+      needsImplementsAttr = false;
+    }
+
     auto aliasDecl = new (ctx) TypeAliasDecl(
-        SourceLoc(), SourceLoc(), assocType->getName(), SourceLoc(),
+        SourceLoc(), SourceLoc(), name, SourceLoc(),
         /*genericparams*/ nullptr, dc);
     aliasDecl->setUnderlyingType(type);
 
     aliasDecl->setImplicit();
     aliasDecl->setSynthesized();
 
+    // If needed, add an @_implements(Protocol, Name) attribute.
+    if (needsImplementsAttr) {
+      auto attr = ImplementsAttr::create(
+          dc, assocType->getProtocol(), assocType->getName());
+      aliasDecl->getAttrs().add(attr);
+    }
+
     // Inject the typealias into the nominal decl that conforms to the protocol.
     auto nominal = dc->getSelfNominalTypeDecl();
     auto requiredAccessScope = evaluateOrDefault(
@@ -395,17 +432,6 @@ static bool isAsyncSequenceFailure(AssociatedTypeDecl *assocType) {
   return assocType->getName() == assocType->getASTContext().Id_Failure;
 }
 
-/// Determine whether this is the AsyncIteratorProtocol.Failure or
-/// AsyncSequence.Failure associated type.
-static bool isAsyncIteratorOrSequenceFailure(AssociatedTypeDecl *assocType) {
-  auto proto = assocType->getProtocol();
-  if (!proto->isSpecificProtocol(KnownProtocolKind::AsyncIteratorProtocol) &&
-      !proto->isSpecificProtocol(KnownProtocolKind::AsyncSequence))
-    return false;
-
-  return assocType->getName() == assocType->getASTContext().Id_Failure;
-}
-
 /// Attempt to resolve a type witness via member name lookup.
 static ResolveWitnessResult resolveTypeWitnessViaLookup(
                        NormalProtocolConformance *conformance,
@@ -429,7 +455,8 @@ static ResolveWitnessResult resolveTypeWitnessViaLookup(
     abort();
   }
 
-  NLOptions subOptions = (NL_QualifiedDefault | NL_OnlyTypes | NL_ProtocolMembers);
+  NLOptions subOptions = (NL_QualifiedDefault | NL_OnlyTypes |
+                          NL_ProtocolMembers | NL_IncludeAttributeImplements);
 
   // Look for a member type with the same name as the associated type.
   SmallVector<ValueDecl *, 4> candidates;
@@ -455,6 +482,16 @@ static ResolveWitnessResult resolveTypeWitnessViaLookup(
     if (isa<AssociatedTypeDecl>(typeDecl))
       continue;
 
+    // If the name doesn't match and there's no appropriate @_implements
+    // attribute, skip this candidate.
+    //
+    // Also skip candidates in protocol extensions, because they tend to cause
+    // request cycles. We'll look at those during associated type inference.
+    if (assocType->getName() != typeDecl->getName() &&
+        !(witnessHasImplementsAttrForRequiredName(typeDecl, assocType) &&
+          !typeDecl->getDeclContext()->getSelfProtocolDecl()))
+      continue;
+
     auto *genericDecl = cast<GenericTypeDecl>(typeDecl);
 
     // If the declaration has generic parameters, it cannot witness an
@@ -2008,7 +2045,8 @@ AssociatedTypeInference::inferTypeWitnessesViaAssociatedType(
 
   NLOptions subOptions = (NL_QualifiedDefault |
                           NL_OnlyTypes |
-                          NL_ProtocolMembers);
+                          NL_ProtocolMembers |
+                          NL_IncludeAttributeImplements);
 
   // Look for types with the given default name that have appropriate
   // @_implements attributes.
@@ -2030,6 +2068,12 @@ AssociatedTypeInference::inferTypeWitnessesViaAssociatedType(
     if (!defaultProto)
       continue;
 
+    // If the name doesn't match and there's no appropriate @_implements
+    // attribute, skip this candidate.
+    if (defaultName.getBaseName() != typeDecl->getName() &&
+        !witnessHasImplementsAttrForRequiredName(typeDecl, assocType))
+      continue;
+
     // Determine the witness type.
     Type witnessType = getWitnessTypeForMatching(conformance, typeDecl);
     if (!witnessType) continue;

lib/Sema/TypeCheckAttr.cpp

@@ -3885,7 +3885,7 @@ void AttributeChecker::visitImplementsAttr(ImplementsAttr *attr) {
   // conforms to the specified protocol.
   NominalTypeDecl *NTD = DC->getSelfNominalTypeDecl();
   if (auto *OtherPD = dyn_cast<ProtocolDecl>(NTD)) {
-    if (!OtherPD->inheritsFrom(PD) &&
+    if (!(OtherPD == PD || OtherPD->inheritsFrom(PD)) &&
         !(OtherPD->isSpecificProtocol(KnownProtocolKind::DistributedActor) ||
           PD->isSpecificProtocol(KnownProtocolKind::Actor))) {
       diagnose(attr->getLocation(),

lib/Sema/TypeCheckProtocol.cpp

@@ -1281,18 +1281,18 @@ swift::matchWitness(WitnessChecker::RequirementEnvironmentCache &reqEnvCache,
   return matchWitness(dc, req, witness, setup, matchTypes, finalize);
 }
 
-static bool
-witnessHasImplementsAttrForRequiredName(ValueDecl *witness,
-                                        ValueDecl *requirement) {
+bool
+swift::witnessHasImplementsAttrForRequiredName(ValueDecl *witness,
+                                               ValueDecl *requirement) {
   if (auto A = witness->getAttrs().getAttribute<ImplementsAttr>()) {
     return A->getMemberName() == requirement->getName();
   }
   return false;
 }
 
-static bool
-witnessHasImplementsAttrForExactRequirement(ValueDecl *witness,
-                                            ValueDecl *requirement) {
+bool
+swift::witnessHasImplementsAttrForExactRequirement(ValueDecl *witness,
+                                                   ValueDecl *requirement) {
   assert(requirement->isProtocolRequirement());
   auto *PD = cast<ProtocolDecl>(requirement->getDeclContext());
   if (auto A = witness->getAttrs().getAttribute<ImplementsAttr>()) {

lib/Sema/TypeCheckProtocol.h

@@ -223,6 +223,16 @@ AssociatedTypeDecl *findDefaultedAssociatedType(
     DeclContext *dc, NominalTypeDecl *adoptee,
     AssociatedTypeDecl *assocType);
 
+/// Determine whether this witness has an `@_implements` attribute whose
+/// name matches that of the given requirement.
+bool witnessHasImplementsAttrForRequiredName(ValueDecl *witness,
+                                             ValueDecl *requirement);
+
+/// Determine whether this witness has an `@_implements` attribute whose name
+/// and protocol match that of the requirement exactly.
+bool witnessHasImplementsAttrForExactRequirement(ValueDecl *witness,
+                                                 ValueDecl *requirement);
+
 }
 
 #endif // SWIFT_SEMA_PROTOCOL_H

stdlib/public/Distributed/DistributedActor.swift

@@ -448,4 +448,4 @@ public protocol _DistributedActorStub where Self: DistributedActor {}
 @available(SwiftStdlib 6.0, *)
 public func _distributedStubFatalError(function: String = #function) -> Never {
   fatalError("Unexpected invocation of distributed method '\(function)' stub!")
-}
+}

stdlib/public/core/Sequence.swift

@@ -453,7 +453,7 @@ public protocol Sequence<Element> {
 // Provides a default associated type witness for Iterator when the
 // Self type is both a Sequence and an Iterator.
 extension Sequence where Self: IteratorProtocol {
-  // @_implements(Sequence, Iterator)
+  @_implements(Sequence, Iterator)
   public typealias _Default_Iterator = Self
 }
 

test/Concurrency/async_iterator_inference.swift

@@ -76,11 +76,11 @@ case fail
 
 @available(SwiftStdlib 5.1, *)
 func testAssocTypeInference(sf: S.Failure, tsf: TS.Failure, gtsf1: GenericTS<MyError>.Failure, adapter: SequenceAdapter<SpecificTS<MyError>>.Failure, ntas: NormalThrowingAsyncSequence<String, MyError>.Failure) {
-  let _: Int = sf // expected-error{{cannot convert value of type 'S.Failure' (aka 'Never') to specified type 'Int'}}
-  let _: Int = tsf // expected-error{{cannot convert value of type 'TS.Failure' (aka 'any Error') to specified type 'Int'}}
-  let _: Int = gtsf1 // expected-error{{cannot convert value of type 'GenericTS<MyError>.Failure' (aka 'any Error') to specified type 'Int'}}
-  let _: Int = adapter // expected-error{{cannot convert value of type 'SequenceAdapter<SpecificTS<MyError>>.Failure' (aka 'MyError') to specified type 'Int'}}
-  let _: Int = ntas // expected-error{{cannot convert value of type 'NormalThrowingAsyncSequence<String, MyError>.Failure' (aka 'any Error') to specified type 'Int'}}
+  let _: Int = sf // expected-error{{cannot convert value of type 'S.__AsyncSequence_Failure' (aka 'Never') to specified type 'Int'}}
+  let _: Int = tsf // expected-error{{cannot convert value of type 'TS.__AsyncSequence_Failure' (aka 'any Error') to specified type 'Int'}}
+  let _: Int = gtsf1 // expected-error{{cannot convert value of type 'GenericTS<MyError>.__AsyncSequence_Failure' (aka 'any Error') to specified type 'Int'}}
+  let _: Int = adapter // expected-error{{cannot convert value of type 'SequenceAdapter<SpecificTS<MyError>>.__AsyncSequence_Failure' (aka 'MyError') to specified type 'Int'}}
+  let _: Int = ntas // expected-error{{cannot convert value of type 'NormalThrowingAsyncSequence<String, MyError>.__AsyncSequence_Failure' (aka 'any Error') to specified type 'Int'}}
 }
 
 

test/ModuleInterface/async_sequence_conformance.swift

@@ -4,15 +4,21 @@
 
 // REQUIRES: concurrency, OS=macosx
 
-// CHECK: @available(
-// CHECK-NEXT: public struct SequenceAdapte
+// CHECK: public struct SequenceAdapte
 @available(SwiftStdlib 5.1, *)
 public struct SequenceAdapter<Base: AsyncSequence>: AsyncSequence {
   // CHECK-LABEL: public struct AsyncIterator
   // CHECK: @available{{.*}}macOS 10.15
   // CHECK-NEXT: public typealias Element = Base.Element
+
+  // CHECK: #if compiler(>=5.3) && $AssociatedTypeImplements
   // CHECK: @available(
-  // CHECK-NEXT: public typealias Failure = Base.Failure
+  // CHECK: @_implements(_Concurrency.AsyncIteratorProtocol, Failure)
+  // CHECK-SAME: public typealias __AsyncIteratorProtocol_Failure = Base.Failure
+  // CHECK-NEXT: #else
+  // CHECK-NOT: @_implements
+  // CHECK: public typealias __AsyncIteratorProtocol_Failure = Base.Failure
+  // CHECK-NEXT: #endif
   public typealias Element = Base.Element
 
   public struct AsyncIterator: AsyncIteratorProtocol {
@@ -23,11 +29,11 @@ public struct SequenceAdapter<Base: AsyncSequence>: AsyncSequence {
   public func makeAsyncIterator() -> AsyncIterator { AsyncIterator() }
 
   // CHECK: @available(
-  // CHECK-NEXT: public typealias Failure = Base.Failure
+  // CHECK: @_implements(_Concurrency.AsyncSequence, Failure)
+  // CHECK-SAME: public typealias __AsyncSequence_Failure = Base.Failure
 }
 
-// CHECK: @available(
-// CHECK-NEXT: public struct OtherSequenceAdapte
+// CHECK: public struct OtherSequenceAdapte
 @available(SwiftStdlib 5.1, *)
 public struct OtherSequenceAdapter<Base: AsyncSequence>: AsyncSequence {
   // CHECK: public typealias Element = Base.Element
@@ -37,7 +43,8 @@ public struct OtherSequenceAdapter<Base: AsyncSequence>: AsyncSequence {
   // CHECK-LABEL: public struct AsyncIterator
   // CHECK: @available{{.*}}macOS 10.15
   // CHECK: @available(
-  // CHECK-NEXT: public typealias Failure = Base.Failure
+  // CHECK: @_implements(_Concurrency.AsyncIteratorProtocol, Failure)
+  // CHECK-SAME: public typealias __AsyncIteratorProtocol_Failure = Base.Failure
   public typealias Element = Base.Element
 
   public struct Failure: Error { }
@@ -52,3 +59,16 @@ public struct OtherSequenceAdapter<Base: AsyncSequence>: AsyncSequence {
 
   // CHECK-NOT: public typealias Failure
 }
+
+// CHECK: public struct MineOwnIterator
+@available(SwiftStdlib 5.1, *)
+public struct MineOwnIterator<Element>: AsyncSequence, AsyncIteratorProtocol {
+  public mutating func next() async -> Element? { nil }
+  public func makeAsyncIterator() -> Self { self }
+
+  // CHECK:      @_implements(_Concurrency.AsyncIteratorProtocol, Failure)
+  // CHECK-SAME: public typealias __AsyncIteratorProtocol_Failure = Swift.Never
+
+  // CHECK:      @_implements(_Concurrency.AsyncSequence, Failure)
+  // CHECK-SAME: public typealias __AsyncSequence_Failure = Swift.Never
+}

test/attr/attr_implements.swift

@@ -97,3 +97,13 @@ func falseWhenSpecificType(_ x: SpecificType) -> Bool { return x == x }
 
 assert(trueWhenJustEquatable(SpecificType()))
 assert(!falseWhenSpecificType(SpecificType()))
+
+// @_implements on associated types
+protocol PWithAssoc {
+  associatedtype A
+}
+
+struct XWithAssoc: PWithAssoc {
+  @_implements(PWithAssoc, A)
+  typealias __P_A = Int
+}