[Sema] Allow if/switch expressions in optional chain assignments

lib/Sema/MiscDiagnostics.cpp

@@ -3834,6 +3834,23 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
     return MacroWalking::Expansion;
   }
 
+  AssignExpr *findAssignment(Expr *E) const {
+    // Don't consider assignments if we have a parent expression (as otherwise
+    // this would be effectively allowing it in an arbitrary expression
+    // position).
+    if (Parent.getAsExpr())
+      return nullptr;
+
+    // Look through optional exprs, which are present for e.g x?.y = z, as
+    // we wrap the entire assign in the optional evaluation of the destination.
+    if (auto *OEE = dyn_cast<OptionalEvaluationExpr>(E)) {
+      E = OEE->getSubExpr();
+      while (auto *IIO = dyn_cast<InjectIntoOptionalExpr>(E))
+        E = IIO->getSubExpr();
+    }
+    return dyn_cast<AssignExpr>(E);
+  }
+
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
     if (auto *SVE = dyn_cast<SingleValueStmtExpr>(E)) {
       // Diagnose a SingleValueStmtExpr in a context that we do not currently
@@ -3909,13 +3926,9 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
       return Action::Continue(E);
     }
 
-    // Valid as the source of an assignment, as long as it's not a nested
-    // expression (as otherwise this would be effectively allowing it in an
-    // arbitrary expression position).
-    if (auto *AE = dyn_cast<AssignExpr>(E)) {
-      if (!Parent.getAsExpr())
-        markValidSingleValueStmt(AE->getSrc());
-    }
+    // Valid as the source of an assignment.
+    if (auto *AE = findAssignment(E))
+      markValidSingleValueStmt(AE->getSrc());
 
     // Valid as a single expression body of a closure. This is needed in
     // addition to ReturnStmt checking, as we will remove the return if the

test/SILGen/if_expr.swift

@@ -272,6 +272,21 @@ enum E {
 
 struct S {
   var i: Int
+  var opt: Int?
+
+  var computed: Int {
+    get { i }
+    set { i = newValue }
+  }
+  var coroutined: Int {
+    _read { yield i }
+    _modify { yield &i }
+  }
+
+  subscript(x: Int) -> Int {
+    get { i }
+    set { i = newValue }
+  }
 
   mutating func testAssign1(_ x: E) {
     i = if case .e(let y) = x { y } else { 0 }
@@ -358,3 +373,130 @@ struct S {
     return i
   }
 }
+
+enum G {
+  case e(Int)
+  case f
+}
+
+struct TestLValues {
+  var s: S
+  var opt: S?
+  var optopt: S??
+
+  mutating func testOptPromote1() {
+    opt = if .random() { s } else { s }
+  }
+
+  mutating func testOptPromote2() {
+    optopt = if .random() { s } else { s }
+  }
+
+  mutating func testStored1() {
+    s.i = if .random() { 1 } else { 0 }
+  }
+
+  mutating func testStored2() throws {
+    s.i = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testComputed1() {
+    s.computed = if .random() { 1 } else { 0 }
+  }
+
+  mutating func testComputed2() throws {
+    s.computed = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testCoroutined1() {
+    s.coroutined = if .random() { 1 } else { 0 }
+  }
+
+  mutating func testCoroutined2() throws {
+    s.coroutined = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testOptionalChain1() {
+    opt?.i = if .random() { 1 } else { 0 }
+  }
+
+  mutating func testOptionalChain2() throws {
+    opt?.i = if .random() { throw Err() } else { 0 }
+  }
+
+  mutating func testOptionalChain3(_ g: G) {
+    opt?.i = if case .e(let i) = g { i } else { 0 }
+  }
+
+  mutating func testOptionalChain4(_ g: G) throws {
+    opt?.i = if case .e(let i) = g { i } else { throw Err() }
+  }
+
+  mutating func testOptionalChain5(_ g: G) throws {
+    opt?.computed = if case .e(let i) = g { i } else { throw Err() }
+  }
+
+  mutating func testOptionalChain6(_ g: G) throws {
+    opt?.coroutined = if case .e(let i) = g { i } else { throw Err() }
+  }
+
+  mutating func testOptionalChain7() throws {
+    optopt??.i = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testOptionalChain8() throws {
+    optopt??.opt = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testOptionalChain9() throws {
+    optopt??.opt? = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testOptionalForce1() throws {
+    opt!.i = if .random() { throw Err() } else { 0 }
+  }
+
+  mutating func testOptionalForce2() throws {
+    opt!.computed = if .random() { throw Err() } else { 0 }
+  }
+
+  mutating func testOptionalForce3(_ g: G) throws {
+    opt!.coroutined = if case .e(let i) = g { i } else { throw Err() }
+  }
+
+  mutating func testOptionalForce4() throws {
+    optopt!!.i = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testOptionalForce5() throws {
+    optopt!!.opt = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testOptionalForce6() throws {
+    optopt!!.opt! = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testSubscript1() throws {
+    s[5] = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testSubscript2() throws {
+    opt?[5] = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testSubscript3() throws {
+    opt![5] = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testKeyPath1(_ kp: WritableKeyPath<S, Int>) throws {
+    s[keyPath: kp] = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testKeyPath2(_ kp: WritableKeyPath<S, Int>) throws {
+    opt?[keyPath: kp] = if .random() { 1 } else { throw Err() }
+  }
+
+  mutating func testKeyPath3(_ kp: WritableKeyPath<S, Int>) throws {
+    opt![keyPath: kp] = if .random() { 1 } else { throw Err() }
+  }
+}

test/SILGen/switch_expr.swift

@@ -367,6 +367,21 @@ enum F {
 
 struct S {
   var i: Int
+  var opt: Int?
+
+  var computed: Int {
+    get { i }
+    set { i = newValue }
+  }
+  var coroutined: Int {
+    _read { yield i }
+    _modify { yield &i }
+  }
+
+  subscript(x: Int) -> Int {
+    get { i }
+    set { i = newValue }
+  }
 
   mutating func testAssign1(_ x: F) {
     i = switch x {
@@ -463,3 +478,130 @@ struct S {
     return i
   }
 }
+
+enum G {
+  case e(Int)
+  case f
+}
+
+struct TestLValues {
+  var s: S
+  var opt: S?
+  var optopt: S??
+
+  mutating func testOptPromote1() {
+    opt = switch Bool.random() { case true: s case false: s }
+  }
+
+  mutating func testOptPromote2() {
+    optopt = switch Bool.random() { case true: s case false: s }
+  }
+
+  mutating func testStored1() {
+    s.i = switch Bool.random() { case true: 1 case false: 0 }
+  }
+
+  mutating func testStored2() throws {
+    s.i = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testComputed1() {
+    s.computed = switch Bool.random() { case true: 1 case false: 0 }
+  }
+
+  mutating func testComputed2() throws {
+    s.computed = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testCoroutined1() {
+    s.coroutined = switch Bool.random() { case true: 1 case false: 0 }
+  }
+
+  mutating func testCoroutined2() throws {
+    s.coroutined = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testOptionalChain1() {
+    opt?.i = switch Bool.random() { case true: 1 case false: 0 }
+  }
+
+  mutating func testOptionalChain2() throws {
+    opt?.i = switch Bool.random() { case true: throw Err() case false: 0 }
+  }
+
+  mutating func testOptionalChain3(_ g: G) {
+    opt?.i = switch g { case .e(let i): i default: 0 }
+  }
+
+  mutating func testOptionalChain4(_ g: G) throws {
+    opt?.i = switch g { case .e(let i): i default: throw Err() }
+  }
+
+  mutating func testOptionalChain5(_ g: G) throws {
+    opt?.computed = switch g { case .e(let i): i default: throw Err() }
+  }
+
+  mutating func testOptionalChain6(_ g: G) throws {
+    opt?.coroutined = switch g { case .e(let i): i default: throw Err() }
+  }
+
+  mutating func testOptionalChain7() throws {
+    optopt??.i = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testOptionalChain8() throws {
+    optopt??.opt = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testOptionalChain9() throws {
+    optopt??.opt? = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testOptionalForce1() throws {
+    opt!.i = switch Bool.random() { case true: throw Err() case false: 0 }
+  }
+
+  mutating func testOptionalForce2() throws {
+    opt!.computed = switch Bool.random() { case true: throw Err() case false: 0 }
+  }
+
+  mutating func testOptionalForce3(_ g: G) throws {
+    opt!.coroutined = switch g { case .e(let i): i default: throw Err() }
+  }
+
+  mutating func testOptionalForce4() throws {
+    optopt!!.i = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testOptionalForce5() throws {
+    optopt!!.opt = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testOptionalForce6() throws {
+    optopt!!.opt! = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testSubscript1() throws {
+    s[5] = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testSubscript2() throws {
+    opt?[5] = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testSubscript3() throws {
+    opt![5] = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testKeyPath1(_ kp: WritableKeyPath<S, Int>) throws {
+    s[keyPath: kp] = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testKeyPath2(_ kp: WritableKeyPath<S, Int>) throws {
+    opt?[keyPath: kp] = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+
+  mutating func testKeyPath3(_ kp: WritableKeyPath<S, Int>) throws {
+    opt![keyPath: kp] = switch Bool.random() { case true: 1 case false: throw Err() }
+  }
+}