[Unsafe] Teach for..in loops to let the sequence's 'unsafe' cover next()

lib/Sema/CSGen.cpp

@@ -3854,6 +3854,12 @@ generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
           AwaitExpr::createImplicit(ctx, nextCall->getLoc(), nextCall);
     }
 
+    // Wrap the 'next' call in 'unsafe', if there is one.
+    if (unsafeExpr) {
+      nextCall = new (ctx) UnsafeExpr(unsafeExpr->getLoc(), nextCall, Type(),
+                                      /*implicit=*/true);
+    }
+
     // The iterator type must conform to IteratorProtocol.
     {
       ProtocolDecl *iteratorProto = TypeChecker::getProtocol(

lib/Sema/TypeCheckEffects.cpp

@@ -4122,6 +4122,12 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
 
       if (auto parsedSequence = S->getParsedSequence()) {
         parentMap[typeCheckedExpr] = parsedSequence;
+
+        if (auto nextCall = S->getNextCall()) {
+          auto nextParentMap = nextCall->getParentMap();
+          parentMap.insert(nextParentMap.begin(), nextParentMap.end());
+          parentMap[nextCall] = parsedSequence;
+        }
       }
     }
 

test/Unsafe/safe.swift

@@ -87,15 +87,17 @@ func returnsExistentialP() -> any P {
   // expected-note@-1{{@unsafe conformance of 'Int' to protocol 'P' involves unsafe code}}
 }
 
-struct UnsafeAsSequence: @unsafe Sequence, IteratorProtocol {
-  mutating func next() -> Int? { nil }
+// FIXME: Should work even if the IteratorProtocol conformance is safe
+struct UnsafeAsSequence: @unsafe Sequence, @unsafe IteratorProtocol {
+  @unsafe mutating func next() -> Int? { nil }
 }
 
 func testUnsafeAsSequenceForEach() {
   let uas = UnsafeAsSequence()
 
   // expected-warning@+1{{expression uses unsafe constructs but is not marked with 'unsafe'}}{{12-12=unsafe }}
   for _ in uas { } // expected-note{{conformance}}
+  // expected-note@-1{{reference}}
 
   for _ in unsafe uas { } // okay
 }

test/Unsafe/unsafe-suppression.swift

@@ -143,3 +143,15 @@ var yieldUnsafeOkay: Int {
     yield unsafe &x
   }
 }
+
+struct UnsafeSequence: @unsafe IteratorProtocol, @unsafe Sequence {
+  @unsafe func next() -> Int? { nil }
+}
+
+func forEachLoop(us: UnsafeSequence) {
+  for _ in us { } // expected-warning{{expression uses unsafe constructs but is not marked with 'unsafe' [Unsafe]}}{{12-12=unsafe }}
+  // expected-note@-1{{@unsafe conformance of 'UnsafeSequence' to protocol 'Sequence' involves unsafe code}}
+  // expected-note@-2{{reference to unsafe instance method 'next()'}}
+
+  for _ in unsafe us { }
+}