Handle unsafe expressions in yield statements

Author: DougGregor

lib/Sema/TypeCheckStmt.cpp

@@ -1123,6 +1123,7 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
       auto exprToCheck = yieldExprs[i];
 
       InOutExpr *inout = nullptr;
+      UnsafeExpr *unsafeExpr = nullptr;
 
       // Classify whether we're yielding by reference or by value.
       ContextualTypePurpose contextTypePurpose;
@@ -1131,6 +1132,11 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
         contextTypePurpose = CTP_YieldByReference;
         contextType = LValueType::get(contextType);
 
+        // If present, remove an enclosing 'unsafe' expression.
+        unsafeExpr = dyn_cast<UnsafeExpr>(exprToCheck);
+        if (unsafeExpr)
+          exprToCheck = unsafeExpr->getSubExpr();
+
         // Check that the yielded expression is a &.
         if ((inout = dyn_cast<InOutExpr>(exprToCheck))) {
           // Strip the & off so that the constraint system doesn't complain
@@ -1158,6 +1164,13 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
         exprToCheck = inout;
       }
 
+      // Propagate the change to the unsafe expression we stripped before.
+      if (unsafeExpr) {
+        unsafeExpr->setSubExpr(exprToCheck);
+        unsafeExpr->setType(exprToCheck->getType());
+        exprToCheck = unsafeExpr;
+      }
+
       // Note that this modifies the statement's expression list in-place.
       yieldExprs[i] = exprToCheck;
     }

test/Unsafe/unsafe-suppression.swift

@@ -116,3 +116,30 @@ extension UnsafeOuter {
 extension UnsafeOuter {
   func i(_: UnsafeType) { }
 }
+
+// -----------------------------------------------------------------------
+// Miscellaneous issues
+// -----------------------------------------------------------------------
+var yieldUnsafe: Int {
+  _read {
+    @unsafe let x = 5
+    yield x // expected-warning{{expression uses unsafe constructs but is not marked with 'unsafe' [Unsafe]}}
+    // expected-note@-1{{reference to unsafe let 'x'}}
+  }
+  _modify {
+    @unsafe var x = 5
+    yield &x // expected-warning{{expression uses unsafe constructs but is not marked with 'unsafe' [Unsafe]}}
+    // expected-note@-1{{reference to unsafe var 'x'}}
+  }
+}
+
+var yieldUnsafeOkay: Int {
+  _read {
+    @unsafe let x = 5
+    yield unsafe x
+  }
+  _modify {
+    @unsafe var x = 5
+    yield unsafe &x
+  }
+}