[ASTScopes] Fix scope expansion for if/switch expressions

lib/AST/ASTScopeCreation.cpp

@@ -100,13 +100,14 @@ class ScopeCreator final : public ASTAllocated<ScopeCreator> {
     ASTScopeAssert(expr,
                  "If looking for closures, must have an expression to search.");
 
-    /// AST walker that finds top-level closures in an expression.
-    class ClosureFinder : public ASTWalker {
+    /// AST walker that finds nested scopes in expressions. This handles both
+    /// closures and if/switch expressions.
+    class NestedExprScopeFinder : public ASTWalker {
       ScopeCreator &scopeCreator;
       ASTScopeImpl *parent;
 
     public:
-      ClosureFinder(ScopeCreator &scopeCreator, ASTScopeImpl *parent)
+      NestedExprScopeFinder(ScopeCreator &scopeCreator, ASTScopeImpl *parent)
           : scopeCreator(scopeCreator), parent(parent) {}
 
       PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
@@ -122,6 +123,13 @@ class ScopeCreator final : public ASTAllocated<ScopeCreator> {
                   parent, capture);
           return Action::SkipChildren(E);
         }
+
+        // If we have a single value statement expression, we need to add any
+        // scopes in the underlying statement.
+        if (auto *SVE = dyn_cast<SingleValueStmtExpr>(E)) {
+          scopeCreator.addToScopeTree(SVE->getStmt(), parent);
+          return Action::SkipChildren(E);
+        }
         return Action::Continue(E);
       }
       PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
@@ -148,7 +156,7 @@ class ScopeCreator final : public ASTAllocated<ScopeCreator> {
       }
     };
 
-    expr->walk(ClosureFinder(*this, parent));
+    expr->walk(NestedExprScopeFinder(*this, parent));
   }
 
 public:
@@ -518,11 +526,6 @@ class NodeAdder
     if (!expr)
       return p;
 
-    // If we have a single value statement expression, we expand scopes based
-    // on the underlying statement.
-    if (auto *SVE = dyn_cast<SingleValueStmtExpr>(expr))
-      return visit(SVE->getStmt(), p, scopeCreator);
-
     scopeCreator.addExprToScopeTree(expr, p);
     return p;
   }

test/SILGen/if_expr.swift

@@ -263,3 +263,98 @@ func nestedType() throws -> Int {
     0
   }
 }
+
+// MARK: Bindings
+
+enum E {
+  case e(Int)
+}
+
+struct S {
+  var i: Int
+
+  mutating func testAssign1(_ x: E) {
+    i = if case .e(let y) = x { y } else { 0 }
+  }
+
+
+  mutating func testAssign2(_ x: E) {
+    i = if case .e(let y) = x { Int(y) } else { 0 }
+  }
+
+  func testAssign3(_ x: E) {
+    var i = 0
+    i = if case .e(let y) = x { y } else { 0 }
+    _ = i
+  }
+
+  func testAssign4(_ x: E) {
+    var i = 0
+    let _ = {
+      i = if case .e(let y) = x { y } else { 0 }
+    }
+    _ = i
+  }
+
+  mutating func testAssign5(_ x: E) {
+    i = switch Bool.random() {
+    case true:
+      if case .e(let y) = x { y } else { 0 }
+    case let z:
+      z ? 0 : 1
+    }
+  }
+
+  mutating func testAssign6(_ x: E) {
+    i = if case .e(let y) = x {
+      switch Bool.random() {
+      case true: y
+      case false: y
+      }
+    } else {
+      0
+    }
+  }
+
+  mutating func testAssign7(_ x: E?) {
+    i = if let x = x {
+      switch x {
+      case .e(let y): y
+      }
+    } else {
+      0
+    }
+  }
+
+  func testReturn1(_ x: E) -> Int {
+    if case .e(let y) = x { y } else { 0 }
+  }
+
+  func testReturn2(_ x: E) -> Int {
+    return if case .e(let y) = x { y } else { 0 }
+  }
+
+  func testReturn3(_ x: E) -> Int {
+    {
+      if case .e(let y) = x { y } else { 0 }
+    }()
+  }
+
+  func testReturn4(_ x: E) -> Int {
+    return {
+      if case .e(let y) = x { y } else { 0 }
+    }()
+  }
+
+  func testBinding1(_ x: E) -> Int {
+    let i = if case .e(let y) = x { y } else { 0 }
+    return i
+  }
+
+  func testBinding2(_ x: E) -> Int {
+    let i = {
+      if case .e(let y) = x { y } else { 0 }
+    }()
+    return i
+  }
+}

test/SILGen/switch_expr.swift

@@ -358,3 +358,108 @@ func nestedType() throws -> Int {
     0
   }
 }
+
+// MARK: Bindings
+
+enum F {
+  case e(Int)
+}
+
+struct S {
+  var i: Int
+
+  mutating func testAssign1(_ x: F) {
+    i = switch x {
+    case .e(let y): y
+    }
+  }
+
+  mutating func testAssign2(_ x: F) {
+    i = switch x {
+    case .e(let y): Int(y)
+    }
+  }
+
+  func testAssign3(_ x: F) {
+    var i = 0
+    i = switch x {
+    case .e(let y): y
+    }
+    _ = i
+  }
+
+  func testAssign4(_ x: F) {
+    var i = 0
+    let _ = {
+      i = switch x {
+      case .e(let y): y
+      }
+    }
+    _ = i
+  }
+
+  mutating func testAssign5(_ x: F) {
+    i = switch Bool.random() {
+    case true:
+      switch x {
+      case .e(let y): y
+      }
+    case let z:
+      z ? 0 : 1
+    }
+  }
+
+  mutating func testAssign6(_ x: F) {
+    i = switch x {
+    case .e(let y):
+      switch Bool.random() {
+      case true: y
+      case false: y
+      }
+    }
+  }
+
+  func testReturn1(_ x: F) -> Int {
+    switch x {
+    case .e(let y): y
+    }
+  }
+
+  func testReturn2(_ x: F) -> Int {
+    return switch x {
+    case .e(let y): y
+    }
+  }
+
+  func testReturn3(_ x: F) -> Int {
+    {
+      switch x {
+      case .e(let y): y
+      }
+    }()
+  }
+
+  func testReturn4(_ x: F) -> Int {
+    return {
+      switch x {
+      case .e(let y): y
+      }
+    }()
+  }
+
+  func testBinding1(_ x: F) -> Int {
+    let i = switch x {
+    case .e(let y): y
+    }
+    return i
+  }
+
+  func testBinding2(_ x: F) -> Int {
+    let i = {
+      switch x {
+      case .e(let y): y
+      }
+    }()
+    return i
+  }
+}