[CS] Ensure DeclContext is switched for if/switch expressions

Similar to multi-statement closures, we need to
make sure we change the DeclContext of the
ConstraintSystem when solving a conjunction for an
if/switch expression. This is unfortunately needed
for cases within single-expression closures since
they can't switch the system's DeclContext (as
they're solved together with the rest of the system).

Author: hamishknight

lib/Sema/CSStep.h

@@ -802,6 +802,27 @@ class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
   }
 };
 
+/// Retrieves the DeclContext that a conjunction should be solved within.
+static DeclContext *getDeclContextForConjunction(ConstraintLocator *loc) {  
+  // Closures introduce a new DeclContext that needs switching into.
+  auto anchor = loc->getAnchor();
+  if (loc->directlyAt<ClosureExpr>())
+    return castToExpr<ClosureExpr>(anchor);
+
+  // SingleValueStmtExprs need to switch to their enclosing context. This
+  // is unfortunately necessary since they can be present in single-expression
+  // closures, which don't have their DeclContext established since they're
+  // solved together with the rest of the system.
+  if (loc->isForSingleValueStmtConjunction())
+    return castToExpr<SingleValueStmtExpr>(anchor)->getDeclContext();
+  
+  // Do the same for TapExprs.
+  if (loc->directlyAt<TapExpr>())
+    return castToExpr<TapExpr>(anchor)->getVar()->getDeclContext();
+
+  return nullptr;
+}
+
 class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
   /// Snapshot of the constraint system before conjunction.
   class SolverSnapshot {
@@ -826,11 +847,9 @@ class ConjunctionStep : public BindingStep<ConjunctionElementProducer> {
         : CS(cs), Conjunction(conjunction),
           TypeVars(std::move(cs.TypeVariables)) {
       auto *locator = Conjunction->getLocator();
-      // If this conjunction represents a closure, we need to
-      // switch declaration context over to it.
-      if (locator->directlyAt<ClosureExpr>()) {
-        DC.emplace(CS.DC, castToExpr<ClosureExpr>(locator->getAnchor()));
-      }
+      // If we need to switch into a new DeclContext for the conjunction, do so.
+      if (auto *newDC = getDeclContextForConjunction(locator))
+        DC.emplace(CS.DC, newDC);
 
       auto &CG = CS.getConstraintGraph();
       // Remove all of the current inactive constraints.

test/Constraints/if_expr.swift

@@ -706,3 +706,25 @@ func testInvalidOptionalChainingInIfContext() {
   let v63796 = 1
   if v63796? {} // expected-error{{cannot use optional chaining on non-optional value of type 'Int'}}
 }
+
+// https://github.com/swiftlang/swift/issues/79395
+_ = {
+  if .random() {
+    struct S: Error {}
+    throw S()
+  } else {
+    1
+  }
+}
+_ = {
+  if .random() {
+    if .random() {
+      struct S: Error {}
+      throw S()
+    } else {
+      0
+    }
+  } else {
+    1
+  }
+}

test/Constraints/switch_expr.swift

@@ -819,3 +819,28 @@ func builderInClosure() {
     }
   }
 }
+
+// https://github.com/swiftlang/swift/issues/79395
+_ = {
+  switch Bool.random() {
+  case true:
+    struct S: Error {}
+    throw S()
+  case false:
+    1
+  }
+}
+_ = {
+  switch Bool.random() {
+  case true:
+    switch Bool.random() {
+    case true:
+      struct S: Error {}
+      throw S()
+    case false:
+      0
+    }
+  case false:
+    1
+  }
+}