Merge pull request #71500 from DougGregor/rethrows-inheriting-async-sequence

Treat `@rethrows` protocols inheriting AsyncSequence like AsyncSequence

Author: DougGregor

lib/Sema/TypeCheckEffects.cpp

@@ -96,6 +96,25 @@ PolymorphicEffectRequirementsRequest::evaluate(Evaluator &evaluator,
                                           ctx.AllocateCopy(conformances));
 }
 
+/// Determine whether the given protocol inherits from either
+/// AsyncIteratorProtocol or AsyncSequence.
+static bool inheritsFromAsyncSequenceProtocol(ProtocolDecl *proto) {
+  // If it's exactly one of these, shortcut.
+  if (proto->isSpecificProtocol(KnownProtocolKind::AsyncIteratorProtocol) ||
+      proto->isSpecificProtocol(KnownProtocolKind::AsyncSequence))
+    return false;
+
+  auto &ctx = proto->getASTContext();
+  if (auto iter = ctx.getProtocol(KnownProtocolKind::AsyncIteratorProtocol))
+    if (proto->inheritsFrom(iter))
+      return true;
+  if (auto seq = ctx.getProtocol(KnownProtocolKind::AsyncSequence))
+    if (proto->inheritsFrom(seq))
+      return true;
+
+  return false;
+}
+
 PolymorphicEffectKind
 PolymorphicEffectKindRequest::evaluate(Evaluator &evaluator,
                                        EffectKind kind,
@@ -117,6 +136,13 @@ PolymorphicEffectKindRequest::evaluate(Evaluator &evaluator,
       auto proto = req.getProtocolDecl();
 
       if (proto->hasPolymorphicEffect(kind)) {
+        // @rethrows protocols that inherit from AsyncIteratorProtocol or
+        // AsyncSequence should be categorized like AsyncIteratorProtocol or
+        // AsyncSequence.
+        if (kind == EffectKind::Throws &&
+            inheritsFromAsyncSequenceProtocol(proto))
+          return PolymorphicEffectKind::AsyncSequenceRethrows;
+
         return PolymorphicEffectKind::ByConformance;
       }
 

test/Concurrency/async_sequence_rethrows.swift

@@ -40,3 +40,14 @@ func testCalls(x: some AsyncSequence<Int, Never>, y: any AsyncSequence<Int, any
   await f3(x, x)
   try await f3(x, y)
 }
+
+// Treat @rethrows protocols that inherit AsyncSequence like they are
+// AsyncSequence for the purpose of rethrowing methods.
+@rethrows
+protocol InheritsAsyncSequence: AsyncSequence { }
+
+extension InheritsAsyncSequence {
+  func blah() async rethrows -> [Element] {
+    try await self.reduce(into: [Element]()) { $0.append($1) }
+  }
+}