From 260ade8076e7d4e5bb7b0b060ae259b7de1e8b2d Mon Sep 17 00:00:00 2001
From: Doug Gregor <dgregor@apple.com>
Date: Mon, 10 Feb 2025 14:51:52 -0800
Subject: [PATCH 1/3] [Unsafe] Teach for..in loops to let the sequence's
 'unsafe' cover next()

Warnings about unsafe uses due to an @unsafe IteratorProtocol conformance
(for the implicit call to next()) could not be silenced. Follow the same
path we did for the Sequence conformance (and makeIterator() call) by
associating it with the `unsafe` on the sequence argument.

This isn't the only solution here, but it's a reasonable one.
---
 lib/Sema/CSGen.cpp                   |  6 ++++++
 lib/Sema/TypeCheckEffects.cpp        |  6 ++++++
 test/Unsafe/unsafe-suppression.swift | 12 ++++++++++++
 3 files changed, 24 insertions(+)

diff --git a/lib/Sema/CSGen.cpp b/lib/Sema/CSGen.cpp
index b94d96491b54c..ed935427b53f6 100644
--- a/lib/Sema/CSGen.cpp
+++ b/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(
diff --git a/lib/Sema/TypeCheckEffects.cpp b/lib/Sema/TypeCheckEffects.cpp
index 88b1c55d6770e..ff1f670635b86 100644
--- a/lib/Sema/TypeCheckEffects.cpp
+++ b/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;
+        }
       }
     }
 
diff --git a/test/Unsafe/unsafe-suppression.swift b/test/Unsafe/unsafe-suppression.swift
index fc0983b0a0565..57e049aaed36c 100644
--- a/test/Unsafe/unsafe-suppression.swift
+++ b/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 { }
+}

From 05113a69fbef95ab5a0407c99ea04ac903a5a233 Mon Sep 17 00:00:00 2001
From: Doug Gregor <dgregor@apple.com>
Date: Mon, 10 Feb 2025 16:08:37 -0800
Subject: [PATCH 2/3] Fixup test case based on limitations of the existing
 implementation

---
 test/Unsafe/safe.swift | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test/Unsafe/safe.swift b/test/Unsafe/safe.swift
index 26ff07135358b..f7ee17d1c955e 100644
--- a/test/Unsafe/safe.swift
+++ b/test/Unsafe/safe.swift
@@ -87,7 +87,8 @@ func returnsExistentialP() -> any P {
   // expected-note@-1{{@unsafe conformance of 'Int' to protocol 'P' involves unsafe code}}
 }
 
-struct UnsafeAsSequence: @unsafe Sequence, IteratorProtocol {
+// FIXME: Should work even if the IteratorProtocol conformance is safe
+struct UnsafeAsSequence: @unsafe Sequence, @unsafe IteratorProtocol {
   mutating func next() -> Int? { nil }
 }
 
@@ -96,6 +97,7 @@ func testUnsafeAsSequenceForEach() {
 
   // 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
 }

From 02ada804f1c86d677a65ce2cb6e797842bb77fc9 Mon Sep 17 00:00:00 2001
From: Doug Gregor <dgregor@apple.com>
Date: Mon, 10 Feb 2025 17:10:52 -0800
Subject: [PATCH 3/3] More temporary test fixes

---
 test/Unsafe/safe.swift | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/Unsafe/safe.swift b/test/Unsafe/safe.swift
index f7ee17d1c955e..31bca73c797ca 100644
--- a/test/Unsafe/safe.swift
+++ b/test/Unsafe/safe.swift
@@ -89,7 +89,7 @@ func returnsExistentialP() -> any P {
 
 // FIXME: Should work even if the IteratorProtocol conformance is safe
 struct UnsafeAsSequence: @unsafe Sequence, @unsafe IteratorProtocol {
-  mutating func next() -> Int? { nil }
+  @unsafe mutating func next() -> Int? { nil }
 }
 
 func testUnsafeAsSequenceForEach() {