[test] Add a couple of locator tests for if/switch exprs

Author: hamishknight

unittests/Sema/ConstraintGenerationTests.cpp

@@ -220,3 +220,125 @@ TEST_F(SemaTest, TestMultiStmtClosureBodyParentAndDepth) {
   ASSERT_EQ(cs.getParentExpr(result), closure);
   ASSERT_EQ(cs.getExprDepth(result), 1);
 }
+
+TEST_F(SemaTest, TestIfExprLocator) {
+  // Test to make sure we get the expected conjunction and locators.
+
+  // if true { 1 } else { 2 }
+  auto *ifCond = new (Context) BooleanLiteralExpr(true, SourceLoc());
+  auto *thenBrace = BraceStmt::createImplicit(
+      Context,
+      {IntegerLiteralExpr::createFromUnsigned(Context, 1, SourceLoc())});
+  auto *elseBrace = BraceStmt::createImplicit(
+      Context,
+      {IntegerLiteralExpr::createFromUnsigned(Context, 2, SourceLoc())});
+  auto *ifStmt =
+      new (Context) IfStmt(SourceLoc(), ifCond, thenBrace, SourceLoc(),
+                           elseBrace, /*implicit*/ true, Context);
+  auto *ifExpr = SingleValueStmtExpr::createWithWrappedBranches(
+      Context, ifStmt, DC, /*mustBeExpr*/ true);
+
+  SyntacticElementTarget target(ifExpr, DC, ContextualTypeInfo(),
+                                /*isDiscarded*/ true);
+
+  ConstraintSystem cs(DC, ConstraintSystemOptions());
+  auto hadError = cs.generateConstraints(target);
+  ASSERT_FALSE(hadError);
+
+  auto *conjunction = &cs.getConstraints().front();
+  ASSERT_EQ(conjunction->getKind(), ConstraintKind::Conjunction);
+
+  auto nested = conjunction->getNestedConstraints();
+  ASSERT_EQ(nested.size(), 4);
+
+  auto *condConstraint = nested[0];
+  auto *thenConstraint = nested[1];
+  auto *elseConstraint = nested[2];
+  auto *joinConstraint = nested[3];
+
+  auto *ifStmtLoc =
+      cs.getConstraintLocator(ifExpr, LocatorPathElt::SyntacticElement(ifStmt));
+
+  auto *condLoc =
+      cs.getConstraintLocator(ifStmtLoc, LocatorPathElt::Condition());
+  ASSERT_EQ(condConstraint->getLocator(), condLoc);
+
+  auto *thenLoc =
+      cs.getConstraintLocator(ifStmtLoc, LocatorPathElt::TernaryBranch(true));
+  ASSERT_EQ(thenConstraint->getLocator(), thenLoc);
+
+  auto *elseLoc =
+      cs.getConstraintLocator(ifStmtLoc, LocatorPathElt::TernaryBranch(false));
+  ASSERT_EQ(elseConstraint->getLocator(), elseLoc);
+
+  auto *joinLoc = cs.getConstraintLocator(
+      ifStmtLoc,
+      LocatorPathElt::SyntacticElement(joinConstraint->getSyntacticElement()));
+  ASSERT_EQ(joinConstraint->getLocator(), joinLoc);
+}
+
+TEST_F(SemaTest, TestSwitchExprLocator) {
+  // Test to make sure we get the expected conjunction and locators.
+
+  // case true: 1
+  auto *trueBrace = BraceStmt::createImplicit(
+      Context,
+      {IntegerLiteralExpr::createFromUnsigned(Context, 1, SourceLoc())});
+  auto *truePattern = ExprPattern::createImplicit(
+      Context, new (Context) BooleanLiteralExpr(true, SourceLoc()), DC);
+  auto *trueCase =
+      CaseStmt::create(Context, CaseParentKind::Switch, SourceLoc(),
+                       {CaseLabelItem(truePattern)}, SourceLoc(), SourceLoc(),
+                       trueBrace, /*caseBodyVars*/ llvm::None);
+
+  // case false: 2
+  auto *falseBrace = BraceStmt::createImplicit(
+      Context,
+      {IntegerLiteralExpr::createFromUnsigned(Context, 2, SourceLoc())});
+  auto *falsePattern = ExprPattern::createImplicit(
+      Context, new (Context) BooleanLiteralExpr(false, SourceLoc()), DC);
+  auto *falseCase =
+      CaseStmt::create(Context, CaseParentKind::Switch, SourceLoc(),
+                       {CaseLabelItem(falsePattern)}, SourceLoc(), SourceLoc(),
+                       falseBrace, /*caseBodyVars*/ llvm::None);
+
+  auto *subject = new (Context) BooleanLiteralExpr(true, SourceLoc());
+
+  // switch true { case true: 1 case false: 2 }
+  auto *switchStmt = SwitchStmt::create(LabeledStmtInfo(), SourceLoc(), subject,
+                                        SourceLoc(), {trueCase, falseCase},
+                                        SourceLoc(), SourceLoc(), Context);
+  auto *switchExpr = SingleValueStmtExpr::createWithWrappedBranches(
+      Context, switchStmt, DC, /*mustBeExpr*/ true);
+
+  SyntacticElementTarget target(switchExpr, DC, ContextualTypeInfo(),
+                                /*isDiscarded*/ true);
+
+  ConstraintSystem cs(DC, ConstraintSystemOptions());
+  auto hadError = cs.generateConstraints(target);
+  ASSERT_FALSE(hadError);
+
+  auto *conjunction = &cs.getConstraints().front();
+  ASSERT_EQ(conjunction->getKind(), ConstraintKind::Conjunction);
+
+  auto nested = conjunction->getNestedConstraints();
+  ASSERT_EQ(nested.size(), 4);
+
+  auto *subjectConstraint = nested[0];
+  auto *trueConstraint = nested[1];
+  auto *falseConstraint = nested[2];
+  auto *joinConstraint = nested[3];
+
+  auto *switchStmtLoc = cs.getConstraintLocator(
+      switchExpr, LocatorPathElt::SyntacticElement(switchStmt));
+
+  // These currently all get the switch statement locator.
+  ASSERT_EQ(subjectConstraint->getLocator(), switchStmtLoc);
+  ASSERT_EQ(trueConstraint->getLocator(), switchStmtLoc);
+  ASSERT_EQ(falseConstraint->getLocator(), switchStmtLoc);
+
+  auto *joinLoc = cs.getConstraintLocator(
+      switchStmtLoc,
+      LocatorPathElt::SyntacticElement(joinConstraint->getSyntacticElement()));
+  ASSERT_EQ(joinConstraint->getLocator(), joinLoc);
+}