Diagnose unsafe sequences in the for..in loop

The `unsafe` goes on the sequence expression, e.g.,

    for x in unsafe mySequence { ... }

Author: DougGregor

include/swift/AST/UnsafeUse.h

@@ -225,6 +225,33 @@ class UnsafeUse {
     }
   }
 
+  /// Replace the location, if possible.
+  void replaceLocation(SourceLoc loc) {
+    switch (getKind()) {
+    case Override:
+    case Witness:
+    case PreconcurrencyImport:
+      // Cannot replace location.
+      return;
+
+    case UnsafeConformance:
+      storage.conformance.location = loc.getOpaquePointerValue();
+      break;
+
+    case TypeWitness:
+      storage.typeWitness.location = loc.getOpaquePointerValue();
+      break;
+
+    case UnownedUnsafe:
+    case ExclusivityUnchecked:
+    case NonisolatedUnsafe:
+    case ReferenceToUnsafe:
+    case ReferenceToUnsafeThroughTypealias:
+    case CallToUnsafe:
+      storage.entity.location = loc.getOpaquePointerValue();
+    }
+  }
+
   /// Get the main declaration, when there is one.
   const Decl *getDecl() const {
     switch (getKind()) {

lib/Sema/CSGen.cpp

@@ -3663,6 +3663,15 @@ generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
   ASTContext &ctx = cs.getASTContext();
   bool isAsync = stmt->getAwaitLoc().isValid();
   auto *sequenceExpr = stmt->getParsedSequence();
+
+  // If we have an unsafe expression for the sequence, lift it out of the
+  // sequence expression. We'll put it back after we've introduced the
+  // various calls.
+  UnsafeExpr *unsafeExpr = dyn_cast<UnsafeExpr>(sequenceExpr);
+  if (unsafeExpr) {
+    sequenceExpr = unsafeExpr->getSubExpr();
+  }
+
   auto contextualLocator = cs.getConstraintLocator(
       sequenceExpr, LocatorPathElt::ContextualType(CTP_ForEachSequence));
   auto elementLocator = cs.getConstraintLocator(
@@ -3712,9 +3721,15 @@ generateForEachStmtConstraints(ConstraintSystem &cs, DeclContext *dc,
         ctx, sequenceExpr, makeIterator->getName());
     makeIteratorRef->setFunctionRefInfo(FunctionRefInfo::singleBaseNameApply());
 
-    auto *makeIteratorCall =
+    Expr *makeIteratorCall =
         CallExpr::createImplicitEmpty(ctx, makeIteratorRef);
 
+    // Swap in the 'unsafe' expression.
+    if (unsafeExpr) {
+      unsafeExpr->setSubExpr(makeIteratorCall);
+      makeIteratorCall = unsafeExpr;
+    }
+
     Pattern *pattern = NamedPattern::createImplicit(ctx, makeIteratorVar);
     auto *PB = PatternBindingDecl::createImplicit(
         ctx, StaticSpellingKind::None, pattern, makeIteratorCall, dc);

lib/Sema/TypeCheckEffects.cpp

@@ -3091,6 +3091,9 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
   llvm::MapVector<Expr *, std::vector<UnsafeUse>> uncoveredUnsafeUses;
 
   static bool isEffectAnchor(Expr *e) {
+    if (e->getLoc().isInvalid())
+      return false;
+
     return isa<AbstractClosureExpr>(e) || isa<DiscardAssignmentExpr>(e) ||
            isa<AssignExpr>(e) || (isa<DeclRefExpr>(e) && e->isImplicit());
   }
@@ -3757,11 +3760,23 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
         Expr *expr = E.dyn_cast<Expr*>();
         Expr *anchor = walkToAnchor(expr, parentMap,
                                     CurContext.isWithinInterpolatedString());
+
+        // Figure out a location to use if the unsafe use didn't have one.
+        SourceLoc replacementLoc;
+        if (anchor)
+          replacementLoc = anchor->getLoc();
+        if (replacementLoc.isInvalid())
+          replacementLoc = E.getStartLoc();
+
         auto &anchorUncovered = uncoveredUnsafeUses[anchor];
-        anchorUncovered.insert(
-            anchorUncovered.end(),
-            classification.getUnsafeUses().begin(),
-            classification.getUnsafeUses().end());
+        for (UnsafeUse unsafeUse : classification.getUnsafeUses()) {
+          // If we don't have a source location for this use, use the
+          // anchor instead.
+          if (unsafeUse.getLocation().isInvalid())
+            unsafeUse.replaceLocation(replacementLoc);
+
+          anchorUncovered.push_back(unsafeUse);
+        }
       }
       break;
     }
@@ -3940,6 +3955,16 @@ class CheckEffectsCoverage : public EffectsHandlingWalker<CheckEffectsCoverage>
   }
 
   ShouldRecurse_t checkForEach(ForEachStmt *S) {
+    // Reparent the type-checked sequence on the parsed sequence, so we can
+    // find an anchor.
+    if (auto typeCheckedExpr = S->getTypeCheckedSequence()) {
+      parentMap = typeCheckedExpr->getParentMap();
+
+      if (auto parsedSequence = S->getParsedSequence()) {
+        parentMap[typeCheckedExpr] = parsedSequence;
+      }
+    }
+
     if (!S->getAwaitLoc().isValid())
       return ShouldRecurse;
 

test/Unsafe/safe.swift

@@ -81,6 +81,19 @@ 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 }
+}
+
+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}}
+
+  for _ in unsafe uas { } // okay
+}
+
 class MyRange {
   @unsafe init(unchecked bounds: Range<Int>) { }