Merge pull request #35761 from phausler/async_sequence_conformance_fixes

[Sema] Corrections for for-await-in syntax to prevent specific bad code-gen scenarios and improve diagnostics

Author: DougGregor

lib/Sema/ConstraintSystem.cpp

@@ -2295,6 +2295,13 @@ FunctionType::ExtInfo ConstraintSystem::closureEffects(ClosureExpr *expr) {
         return { false, stmt };
       }
 
+      if (auto forEach = dyn_cast<ForEachStmt>(stmt)) {
+        if (forEach->getTryLoc().isValid()) {
+          FoundThrow = true;
+          return { false, nullptr };
+        }
+      }
+
       return { true, stmt };
     }
 
@@ -2332,6 +2339,17 @@ FunctionType::ExtInfo ConstraintSystem::closureEffects(ClosureExpr *expr) {
       return true;
     }
 
+    std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override { 
+      if (auto forEach = dyn_cast<ForEachStmt>(stmt)) {
+        if (forEach->getAwaitLoc().isValid()) {
+          FoundAsync = true;
+          return { false, nullptr };
+        }
+      }
+
+      return { true, stmt };
+    }
+
   public:
     bool foundAsync() { return FoundAsync; }
   };

lib/Sema/TypeCheckConstraints.cpp

@@ -592,37 +592,17 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
   // check to see if the sequence expr is throwing (and async), if so require 
   // the stmt to have a try loc
   if (stmt->getAwaitLoc().isValid()) {
-    auto Ty = sequence->getType();
-    if (Ty.isNull()) { 
-      auto DRE = dyn_cast<DeclRefExpr>(sequence);
-      if (DRE) {
-        Ty = DRE->getDecl()->getInterfaceType();
-      }
-      if (Ty.isNull()) {
-        return failed();
-      }
-    }
-    auto context = Ty->getNominalOrBoundGenericNominal();
-    if (!context) {
-      // if no nominal type can be determined then we must consider this to be
-      // a potential throwing source and concequently this must have a valid try
-      // location to account for that potential ambiguity.
-      if (stmt->getTryLoc().isInvalid()) {
-        auto &diags = dc->getASTContext().Diags;
-        diags.diagnose(stmt->getAwaitLoc(), diag::throwing_call_unhandled);
-        return failed();
-      } else {
-        return false;
-      }
-       
-    }
+    // fetch the sequence out of the statement
+    // else wise the value is potentially unresolved
+    auto Ty = stmt->getSequence()->getType();
     auto module = dc->getParentModule();
     auto conformanceRef = module->lookupConformance(Ty, sequenceProto);
 
     if (conformanceRef.classifyAsThrows() && 
         stmt->getTryLoc().isInvalid()) {
       auto &diags = dc->getASTContext().Diags;
-      diags.diagnose(stmt->getAwaitLoc(), diag::throwing_call_unhandled);
+      diags.diagnose(stmt->getAwaitLoc(), diag::throwing_call_unhandled)
+        .fixItInsert(stmt->getAwaitLoc(), "try");
 
       return failed();
     }

lib/Sema/TypeCheckEffects.cpp

@@ -528,6 +528,13 @@ class ApplyClassifier {
       }
 
       if (classifiedAsThrows) {
+        // multiple passes can occur, so ensure that any sub-expressions of this
+        // call are marked as throws to mimic the closure variant.
+        if (auto subExpr = dyn_cast<ApplyExpr>(E->getFn())) {
+          if (!subExpr->isThrowsSet()) {
+            subExpr->setThrows(true);
+          }
+        }
         return Classification::forRethrowingOnly(
           PotentialThrowReason::forRethrowsConformance(E), isAsync);
       }
@@ -2096,6 +2103,12 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
                         Classification::forThrow(PotentialThrowReason::forThrow(),
                                                  /*async*/false));
     }
+    if (S->getAwaitLoc().isValid() && 
+        !Flags.has(ContextFlags::HasAnyAsyncSite)) {
+      if (!CurContext.handlesAsync()) {
+        CurContext.diagnoseUnhandledAsyncSite(Ctx.Diags, S);
+      }
+    }
     return ShouldRecurse;
   }
 };

test/Concurrency/async_sequence_syntax.swift

@@ -0,0 +1,57 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-concurrency
+// REQUIRES: concurrency
+
+// expected-note@+1{{add 'async' to function 'missingAsync' to make it asynchronous}}
+func missingAsync<T : AsyncSequence>(_ seq: T) throws { 
+  for try await _ in seq { } // expected-error{{'async' in a function that does not support concurrency}}
+}
+
+func missingThrows<T : AsyncSequence>(_ seq: T) async {
+  for try await _ in seq { } // expected-error{{error is not handled because the enclosing function is not declared 'throws'}}
+}
+
+func executeAsync(_ work: () async -> Void) { }
+func execute(_ work: () -> Void) { }
+
+func missingThrowingInBlock<T : AsyncSequence>(_ seq: T) { 
+  executeAsync { // expected-error{{invalid conversion from throwing function of type '() async throws -> Void' to non-throwing function type '() async -> Void'}}
+    for try await _ in seq { }
+  }
+}
+
+func missingTryInBlock<T : AsyncSequence>(_ seq: T) { 
+  executeAsync { 
+    for await _ in seq { } // expected-error{{call can throw, but the error is not handled}}
+  }
+}
+
+func missingAsyncInBlock<T : AsyncSequence>(_ seq: T) { 
+  execute { // expected-error{{invalid conversion from 'async' function of type '() async -> Void' to synchronous function type '() -> Void'}}
+    do { 
+      for try await _ in seq { }
+    } catch { }
+  }
+}
+
+func doubleDiagCheckGeneric<T : AsyncSequence>(_ seq: T) async {
+  var it = seq.makeAsyncIterator()
+  // expected-note@+2{{call is to 'rethrows' function, but a conformance has a throwing witness}}
+  // expected-error@+1{{call can throw, but it is not marked with 'try' and the error is not handled}}
+  let _ = await it.next()
+}
+
+struct ThrowingAsyncSequence: AsyncSequence, AsyncIteratorProtocol {
+  typealias Element = Int
+  typealias AsyncIterator = Self
+  mutating func next() async throws -> Int? {
+    return nil
+  }
+
+  func makeAsyncIterator() -> Self { return self }
+}
+
+func doubleDiagCheckConcrete(_ seq: ThrowingAsyncSequence) async {
+  var it = seq.makeAsyncIterator()
+  // expected-error@+1{{call can throw, but it is not marked with 'try' and the error is not handled}}
+  let _ = await it.next()
+}