[CS] Fix locator for if expressions

Currently we don't insert an intermediate
`SyntacticElement` locator element for the `if`
statement in a SingleValueStmtExpr, which leads
to an assertion failure when attempting to simplify
a following `TernaryBranch` element.

Move the insertion of the `SyntacticElement` locator
element for the statement into the
SingleValueStmtExpr constraint generation function,
and remove it from the statement visitors themselves.

For the other cases, the `SyntacticElement` was
already being inserted for children of a BraceStmt,
so we were actually previously ending up with
duplicated elements in a couple of places.

Author: hamishknight

lib/Sema/CSSyntacticElement.cpp

@@ -999,14 +999,11 @@ class SyntacticElementConstraintGenerator
   }
 
   void visitSwitchStmt(SwitchStmt *switchStmt) {
-    auto *switchLoc = cs.getConstraintLocator(
-        locator, LocatorPathElt::SyntacticElement(switchStmt));
-
     SmallVector<ElementInfo, 4> elements;
     {
       auto *subjectExpr = switchStmt->getSubjectExpr();
       {
-        elements.push_back(makeElement(subjectExpr, switchLoc));
+        elements.push_back(makeElement(subjectExpr, locator));
 
         SyntacticElementTarget target(subjectExpr, context.getAsDeclContext(),
                                       CTP_Unused, Type(),
@@ -1016,32 +1013,29 @@ class SyntacticElementConstraintGenerator
       }
 
       for (auto rawCase : switchStmt->getRawCases())
-        elements.push_back(makeElement(rawCase, switchLoc));
+        elements.push_back(makeElement(rawCase, locator));
     }
 
-    createConjunction(elements, switchLoc);
+    createConjunction(elements, locator);
   }
 
   void visitDoCatchStmt(DoCatchStmt *doStmt) {
-    auto *doLoc = cs.getConstraintLocator(
-        locator, LocatorPathElt::SyntacticElement(doStmt));
-
     SmallVector<ElementInfo, 4> elements;
 
     // First, let's record a body of `do` statement. Note we need to add a
     // SyntaticElement locator path element here to avoid treating the inner
     // brace conjunction as being isolated if 'doLoc' is for an isolated
     // conjunction (as is the case with 'do' expressions).
     auto *doBodyLoc = cs.getConstraintLocator(
-        doLoc, LocatorPathElt::SyntacticElement(doStmt->getBody()));
+        locator, LocatorPathElt::SyntacticElement(doStmt->getBody()));
     elements.push_back(makeElement(doStmt->getBody(), doBodyLoc));
 
     // After that has been type-checked, let's switch to
     // individual `catch` statements.
     for (auto *catchStmt : doStmt->getCatches())
-      elements.push_back(makeElement(catchStmt, doLoc));
+      elements.push_back(makeElement(catchStmt, locator));
 
-    createConjunction(elements, doLoc);
+    createConjunction(elements, locator);
   }
 
   void visitCaseStmt(CaseStmt *caseStmt) {
@@ -1556,7 +1550,9 @@ bool ConstraintSystem::generateConstraints(SingleValueStmtExpr *E) {
 
   // Generate the conjunction for the branches.
   auto context = SyntacticElementContext::forSingleValueStmtExpr(E, join);
-  SyntacticElementConstraintGenerator generator(*this, context, loc);
+  auto *stmtLoc =
+      getConstraintLocator(loc, LocatorPathElt::SyntacticElement(S));
+  SyntacticElementConstraintGenerator generator(*this, context, stmtLoc);
   generator.visit(S);
   return generator.hadError;
 }

test/Constraints/issue-71282.swift

@@ -0,0 +1,20 @@
+// RUN: %target-typecheck-verify-swift
+
+// https://github.com/apple/swift/issues/71282
+
+enum E {
+  case e
+}
+
+// Make sure we don't crash.
+func foo(_ x: E) {
+  return if .random() {
+    ()
+    switch x {
+    case .e:
+      ()
+    }
+  } else { // expected-error {{non-expression branch of 'if' expression may only end with a 'throw'}}
+    ()
+  }
+}