[SE-0458] Implement "unsafe" effect for the for-in loop

Memory unsafety in the iteration part of the for-in loop (i.e., the part
that works on the iterator) can be covered by the "unsafe" effect on
the for..in loop, before the pattern.

Author: DougGregor

include/swift/AST/ASTBridging.h

@@ -2289,11 +2289,12 @@ BridgedFallthroughStmt_createParsed(BridgedSourceLoc cLoc,
                                     BridgedDeclContext cDC);
 
 SWIFT_NAME("BridgedForEachStmt.createParsed(_:labelInfo:forLoc:tryLoc:awaitLoc:"
-           "pattern:inLoc:sequence:whereLoc:whereExpr:body:)")
+           "unsafeLoc:pattern:inLoc:sequence:whereLoc:whereExpr:body:)")
 BridgedForEachStmt BridgedForEachStmt_createParsed(
     BridgedASTContext cContext, BridgedLabeledStmtInfo cLabelInfo,
     BridgedSourceLoc cForLoc, BridgedSourceLoc cTryLoc,
-    BridgedSourceLoc cAwaitLoc, BridgedPattern cPat, BridgedSourceLoc cInLoc,
+    BridgedSourceLoc cAwaitLoc, BridgedSourceLoc cUnsafeLoc,
+    BridgedPattern cPat, BridgedSourceLoc cInLoc,
     BridgedExpr cSequence, BridgedSourceLoc cWhereLoc,
     BridgedNullableExpr cWhereExpr, BridgedBraceStmt cBody);
 

include/swift/AST/DiagnosticsSema.def

@@ -8212,8 +8212,12 @@ GROUPED_WARNING(preconcurrency_import_unsafe,Unsafe,none,
       "introduce data races", ())
 GROUPED_WARNING(unsafe_without_unsafe,Unsafe,none,
       "expression uses unsafe constructs but is not marked with 'unsafe'", ())
+GROUPED_WARNING(for_unsafe_without_unsafe,Unsafe,none,
+      "for-in loop uses unsafe constructs but is not marked with 'unsafe'", ())
 WARNING(no_unsafe_in_unsafe,none,
         "no unsafe operations occur within 'unsafe' expression", ())
+WARNING(no_unsafe_in_unsafe_for,none,
+        "no unsafe operations occur within 'unsafe' for-in loop", ())
 NOTE(make_subclass_unsafe,none,
      "make class %0 @unsafe to allow unsafe overrides of safe superclass methods",
      (DeclName))

include/swift/AST/Stmt.h

@@ -995,6 +995,7 @@ class ForEachStmt : public LabeledStmt {
   SourceLoc ForLoc;
   SourceLoc TryLoc;
   SourceLoc AwaitLoc;
+  SourceLoc UnsafeLoc;
   Pattern *Pat;
   SourceLoc InLoc;
   Expr *Sequence;
@@ -1012,13 +1013,14 @@ class ForEachStmt : public LabeledStmt {
 
 public:
   ForEachStmt(LabeledStmtInfo LabelInfo, SourceLoc ForLoc, SourceLoc TryLoc,
-              SourceLoc AwaitLoc, Pattern *Pat, SourceLoc InLoc, Expr *Sequence,
+              SourceLoc AwaitLoc, SourceLoc UnsafeLoc, Pattern *Pat,
+              SourceLoc InLoc, Expr *Sequence,
               SourceLoc WhereLoc, Expr *WhereExpr, BraceStmt *Body,
               std::optional<bool> implicit = std::nullopt)
       : LabeledStmt(StmtKind::ForEach, getDefaultImplicitFlag(implicit, ForLoc),
                     LabelInfo),
-        ForLoc(ForLoc), TryLoc(TryLoc), AwaitLoc(AwaitLoc), Pat(nullptr),
-        InLoc(InLoc), Sequence(Sequence), WhereLoc(WhereLoc),
+        ForLoc(ForLoc), TryLoc(TryLoc), AwaitLoc(AwaitLoc), UnsafeLoc(UnsafeLoc),
+        Pat(nullptr), InLoc(InLoc), Sequence(Sequence), WhereLoc(WhereLoc),
         WhereExpr(WhereExpr), Body(Body) {
     setPattern(Pat);
   }
@@ -1056,6 +1058,7 @@ class ForEachStmt : public LabeledStmt {
 
   SourceLoc getAwaitLoc() const { return AwaitLoc; }
   SourceLoc getTryLoc() const { return TryLoc; }
+  SourceLoc getUnsafeLoc() const { return UnsafeLoc; }
 
   /// getPattern - Retrieve the pattern describing the iteration variables.
   /// These variables will only be visible within the body of the loop.

include/swift/Parse/IDEInspectionCallbacks.h

@@ -287,7 +287,8 @@ class CodeCompletionCallbacks {
   virtual void completeStmtLabel(StmtKind ParentKind) {};
 
   virtual
-  void completeForEachPatternBeginning(bool hasTry, bool hasAwait) {};
+  void completeForEachPatternBeginning(
+      bool hasTry, bool hasAwait, bool hasUnsafe) {};
 
   virtual void completeTypeAttrBeginning() {};
 

lib/AST/Bridging/StmtBridging.cpp

@@ -191,14 +191,15 @@ BridgedFallthroughStmt_createParsed(BridgedSourceLoc cLoc,
 BridgedForEachStmt BridgedForEachStmt_createParsed(
     BridgedASTContext cContext, BridgedLabeledStmtInfo cLabelInfo,
     BridgedSourceLoc cForLoc, BridgedSourceLoc cTryLoc,
-    BridgedSourceLoc cAwaitLoc, BridgedPattern cPat, BridgedSourceLoc cInLoc,
+    BridgedSourceLoc cAwaitLoc, BridgedSourceLoc cUnsafeLoc,
+    BridgedPattern cPat, BridgedSourceLoc cInLoc,
     BridgedExpr cSequence, BridgedSourceLoc cWhereLoc,
     BridgedNullableExpr cWhereExpr, BridgedBraceStmt cBody) {
   return new (cContext.unbridged()) ForEachStmt(
       cLabelInfo.unbridged(), cForLoc.unbridged(), cTryLoc.unbridged(),
-      cAwaitLoc.unbridged(), cPat.unbridged(), cInLoc.unbridged(),
-      cSequence.unbridged(), cWhereLoc.unbridged(), cWhereExpr.unbridged(),
-      cBody.unbridged());
+      cAwaitLoc.unbridged(), cUnsafeLoc.unbridged(), cPat.unbridged(),
+      cInLoc.unbridged(), cSequence.unbridged(), cWhereLoc.unbridged(),
+      cWhereExpr.unbridged(), cBody.unbridged());
 }
 
 BridgedGuardStmt BridgedGuardStmt_createParsed(BridgedASTContext cContext,

lib/ASTGen/Sources/ASTGen/Stmts.swift

@@ -323,6 +323,7 @@ extension ASTGenVisitor {
       forLoc: self.generateSourceLoc(node.forKeyword),
       tryLoc: self.generateSourceLoc(node.tryKeyword),
       awaitLoc: self.generateSourceLoc(node.awaitKeyword),
+      unsafeLoc: self.generateSourceLoc(node.unsafeKeyword),
       // NOTE: The pattern can be either a refutable pattern after `case` or a
       // normal binding pattern. ASTGen doesn't care because it should be handled
       // by the parser.

lib/IDE/CodeCompletion.cpp

@@ -302,7 +302,8 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
   void completeGenericRequirement() override;
   void completeAfterIfStmtElse() override;
   void completeStmtLabel(StmtKind ParentKind) override;
-  void completeForEachPatternBeginning(bool hasTry, bool hasAwait) override;
+  void completeForEachPatternBeginning(
+      bool hasTry, bool hasAwait, bool hasUnsafe) override;
   void completeTypeAttrBeginning() override;
   void completeTypeAttrInheritanceBeginning() override;
   void completeOptionalBinding() override;
@@ -636,14 +637,16 @@ void CodeCompletionCallbacksImpl::completeStmtLabel(StmtKind ParentKind) {
 }
 
 void CodeCompletionCallbacksImpl::completeForEachPatternBeginning(
-    bool hasTry, bool hasAwait) {
+    bool hasTry, bool hasAwait, bool hasUnsafe) {
   CurDeclContext = P.CurDeclContext;
   Kind = CompletionKind::ForEachPatternBeginning;
   ParsedKeywords.clear();
   if (hasTry)
     ParsedKeywords.emplace_back("try");
   if (hasAwait)
     ParsedKeywords.emplace_back("await");
+  if (hasUnsafe)
+    ParsedKeywords.emplace_back("unsafe");
 }
 
 void CodeCompletionCallbacksImpl::completeOptionalBinding() {

lib/Parse/ParseStmt.cpp

@@ -2395,6 +2395,7 @@ ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
   auto StartOfControl = Tok.getLoc();
   SourceLoc AwaitLoc;
   SourceLoc TryLoc;
+  SourceLoc UnsafeLoc;
 
   if (Tok.isContextualKeyword("await")) {
     AwaitLoc = consumeToken();
@@ -2405,10 +2406,15 @@ ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
     }
   }
 
+  if (Context.LangOpts.hasFeature(Feature::WarnUnsafe) &&
+      Tok.isContextualKeyword("unsafe")) {
+    UnsafeLoc = consumeToken();
+  }
+  
   if (Tok.is(tok::code_complete)) {
     if (CodeCompletionCallbacks) {
       CodeCompletionCallbacks->completeForEachPatternBeginning(
-          TryLoc.isValid(), AwaitLoc.isValid());
+          TryLoc.isValid(), AwaitLoc.isValid(), UnsafeLoc.isValid());
     }
     consumeToken(tok::code_complete);
     // Since 'completeForeachPatternBeginning' is a keyword only completion,
@@ -2522,7 +2528,8 @@ ParserResult<Stmt> Parser::parseStmtForEach(LabeledStmtInfo LabelInfo) {
 
   return makeParserResult(
       Status,
-      new (Context) ForEachStmt(LabelInfo, ForLoc, TryLoc, AwaitLoc, pattern.get(), InLoc,
+      new (Context) ForEachStmt(LabelInfo, ForLoc, TryLoc, AwaitLoc, UnsafeLoc,
+                                pattern.get(), InLoc,
                                 Container.get(), WhereLoc, Where.getPtrOrNull(),
                                 Body.get()));
 }

lib/Sema/BuilderTransform.cpp

@@ -726,6 +726,7 @@ class ResultBuilderTransform
     auto *newForEach = new (ctx)
         ForEachStmt(forEachStmt->getLabelInfo(), forEachStmt->getForLoc(),
                     forEachStmt->getTryLoc(), forEachStmt->getAwaitLoc(),
+                    forEachStmt->getUnsafeLoc(),
                     forEachStmt->getPattern(), forEachStmt->getInLoc(),
                     forEachStmt->getParsedSequence(),
                     forEachStmt->getWhereLoc(), forEachStmt->getWhere(),

lib/Sema/CSGen.cpp

@@ -4632,10 +4632,11 @@ 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);
+    // Wrap the 'next' call in 'unsafe', if the for..in loop has that
+    // effect.
+    if (stmt->getUnsafeLoc().isValid()) {
+      nextCall = new (ctx) UnsafeExpr(
+          stmt->getUnsafeLoc(), nextCall, Type(), /*implicit=*/true);
     }
 
     // The iterator type must conform to IteratorProtocol.