[Constraint system] Implement switch support for function builders.

Implement support for switch statements within function builders. Cases can
perform arbitrary pattern matches, e.g.,

    tuplify(true) { c in
      "testSwitchCombined"
      switch e {
      case .a:
        "a"
      case .b(let i, _?), .b(let i, nil):
        i + 17
      }
    }

subject to the normal rules of switch statements. Cases within function
builders cannot, however, include “fallthrough” statements, because those
(like “break” and “continue”) are control flow.

The translation of performed for `switch` statements is similar to that of
`if` statements, using `buildEither(first:)` and `buildEither(second:)` on
the function builder type.

This is the bulk of switch support, tracked by rdar://problem/50426203.

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -622,12 +622,104 @@ class BuilderClosureVisitor
         DeclNameLoc(endLoc), /*implicit=*/true);
   }
 
+  VarDecl *visitSwitchStmt(SwitchStmt *switchStmt) {
+    // Generate constraints for the subject expression, and capture its
+    // type for use in matching the various patterns.
+    Expr *subjectExpr = switchStmt->getSubjectExpr();
+    if (cs) {
+      // FIXME: Add contextual type purpose for switch subjects?
+      SolutionApplicationTarget target(subjectExpr, dc, CTP_Unused, Type(),
+                                       /*isDiscarded=*/false);
+      if (cs->generateConstraints(target, FreeTypeVariableBinding::Disallow)) {
+        hadError = true;
+        return nullptr;
+      }
+
+      cs->setSolutionApplicationTarget(switchStmt, target);
+      subjectExpr = target.getAsExpr();
+      assert(subjectExpr && "Must have a subject expression here");
+    }
+
+    // Generate constraints and capture variables for all of the cases.
+    SmallVector<std::pair<CaseStmt *, VarDecl *>, 4> capturedCaseVars;
+    for (auto *caseStmt : switchStmt->getCases()) {
+      if (auto capturedCaseVar = visitCaseStmt(caseStmt, subjectExpr)) {
+        capturedCaseVars.push_back({caseStmt, capturedCaseVar});
+      }
+    }
+
+    if (!cs)
+      return nullptr;
+
+    // Form the expressions that inject the result of each case into the
+    // appropriate
+    llvm::TinyPtrVector<Expr *> injectedCaseExprs;
+    SmallVector<std::pair<Type, ConstraintLocator *>, 4> injectedCaseTerms;
+    for (unsigned idx : indices(capturedCaseVars)) {
+      auto caseStmt = capturedCaseVars[idx].first;
+      auto caseVar = capturedCaseVars[idx].second;
+
+      // Build the expression that injects the case variable into appropriate
+      // buildEither(first:)/buildEither(second:) chain.
+      Expr *caseVarRef = buildVarRef(caseVar, caseStmt->getEndLoc());
+      Expr *injectedCaseExpr = buildWrappedChainPayload(
+          caseVarRef, idx, capturedCaseVars.size(), /*isOptional=*/false);
+
+      // Generate constraints for this injected case result.
+      injectedCaseExpr = cs->generateConstraints(injectedCaseExpr, dc);
+      if (!injectedCaseExpr) {
+        hadError = true;
+        return nullptr;
+      }
+
+      // Record this injected case expression.
+      injectedCaseExprs.push_back(injectedCaseExpr);
+
+      // Record the type and locator for this injected case expression, to be
+      // used in the "join" constraint later.
+      injectedCaseTerms.push_back(
+        { cs->getType(injectedCaseExpr)->getRValueType(),
+          cs->getConstraintLocator(injectedCaseExpr) });
+    }
+
+    // Form the type of the switch itself.
+    // FIXME: Need a locator for the "switch" statement.
+    Type resultType = cs->addJoinConstraint(nullptr, injectedCaseTerms);
+    if (!resultType) {
+      hadError = true;
+      return nullptr;
+    }
+
+    // Create a variable to capture the result of evaluating the switch.
+    auto switchVar = buildVar(switchStmt->getStartLoc());
+    cs->setType(switchVar, resultType);
+    applied.capturedStmts.insert(
+        {switchStmt, { switchVar, std::move(injectedCaseExprs) } });
+    return switchVar;
+  }
+
+  VarDecl *visitCaseStmt(CaseStmt *caseStmt, Expr *subjectExpr) {
+    // If needed, generate constraints for everything in the case statement.
+    if (cs) {
+      auto locator = cs->getConstraintLocator(
+          subjectExpr, LocatorPathElt::ContextualType());
+      Type subjectType = cs->getType(subjectExpr);
+
+      if (cs->generateConstraints(caseStmt, dc, subjectType, locator)) {
+        hadError = true;
+        return nullptr;
+      }
+    }
+
+    // Translate the body.
+    return visit(caseStmt->getBody());
+  }
+
   CONTROL_FLOW_STMT(Guard)
   CONTROL_FLOW_STMT(While)
   CONTROL_FLOW_STMT(DoCatch)
   CONTROL_FLOW_STMT(RepeatWhile)
   CONTROL_FLOW_STMT(ForEach)
-  CONTROL_FLOW_STMT(Switch)
   CONTROL_FLOW_STMT(Case)
   CONTROL_FLOW_STMT(Catch)
   CONTROL_FLOW_STMT(Break)
@@ -996,6 +1088,63 @@ class BuilderClosureRewriter
     return doStmt;
   }
 
+  Stmt *visitSwitchStmt(SwitchStmt *switchStmt, FunctionBuilderTarget target) {
+    // Translate the subject expression.
+    ConstraintSystem &cs = solution.getConstraintSystem();
+    auto subjectTarget =
+        rewriteTarget(*cs.getSolutionApplicationTarget(switchStmt));
+    if (!subjectTarget)
+      return nullptr;
+
+    switchStmt->setSubjectExpr(subjectTarget->getAsExpr());
+
+    // Handle any declaration nodes within the case list first; we'll
+    // handle the cases in a second pass.
+    for (auto child : switchStmt->getRawCases()) {
+      if (auto decl = child.dyn_cast<Decl *>()) {
+        TypeChecker::typeCheckDecl(decl);
+      }
+    }
+
+    // Translate all of the cases.
+    assert(target.kind == FunctionBuilderTarget::TemporaryVar);
+    auto temporaryVar = target.captured.first;
+    unsigned caseIndex = 0;
+    for (auto caseStmt : switchStmt->getCases()) {
+      if (!visitCaseStmt(
+            caseStmt,
+            FunctionBuilderTarget::forAssign(
+              temporaryVar, {target.captured.second[caseIndex]})))
+        return nullptr;
+
+      ++caseIndex;
+    }
+
+    return switchStmt;
+  }
+
+  Stmt *visitCaseStmt(CaseStmt *caseStmt, FunctionBuilderTarget target) {
+    // Translate the patterns and guard expressions for each case label item.
+    for (auto &caseLabelItem : caseStmt->getMutableCaseLabelItems()) {
+      SolutionApplicationTarget caseLabelTarget(&caseLabelItem, dc);
+      if (!rewriteTarget(caseLabelTarget))
+        return nullptr;
+    }
+
+    // Transform the body of the case.
+    auto body = cast<BraceStmt>(caseStmt->getBody());
+    auto captured = takeCapturedStmt(body);
+    auto newInnerBody = cast<BraceStmt>(
+        visitBraceStmt(
+          body,
+          target,
+          FunctionBuilderTarget::forAssign(
+            captured.first, {captured.second.front()})));
+    caseStmt->setBody(newInnerBody);
+
+    return caseStmt;
+  }
+
 #define UNHANDLED_FUNCTION_BUILDER_STMT(STMT) \
   Stmt *visit##STMT##Stmt(STMT##Stmt *stmt, FunctionBuilderTarget target) { \
     llvm_unreachable("Function builders do not allow statement of kind " \
@@ -1010,8 +1159,6 @@ class BuilderClosureRewriter
   UNHANDLED_FUNCTION_BUILDER_STMT(DoCatch)
   UNHANDLED_FUNCTION_BUILDER_STMT(RepeatWhile)
   UNHANDLED_FUNCTION_BUILDER_STMT(ForEach)
-  UNHANDLED_FUNCTION_BUILDER_STMT(Switch)
-  UNHANDLED_FUNCTION_BUILDER_STMT(Case)
   UNHANDLED_FUNCTION_BUILDER_STMT(Catch)
   UNHANDLED_FUNCTION_BUILDER_STMT(Break)
   UNHANDLED_FUNCTION_BUILDER_STMT(Continue)

lib/Sema/CSApply.cpp

@@ -7474,6 +7474,45 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
       }
     }
 
+    return target;
+  } else if (auto caseLabelItem = target.getAsCaseLabelItem()) {
+    ConstraintSystem &cs = solution.getConstraintSystem();
+    auto info = *cs.getCaseLabelItemInfo(*caseLabelItem);
+
+    // Figure out the pattern type.
+    Type patternType = solution.simplifyType(solution.getType(info.pattern));
+    patternType = patternType->reconstituteSugar(/*recursive=*/false);
+
+    // Coerce the pattern to its appropriate type.
+    TypeResolutionOptions patternOptions(TypeResolverContext::InExpression);
+    patternOptions |= TypeResolutionFlags::OverrideType;
+    auto contextualPattern =
+        ContextualPattern::forRawPattern(info.pattern,
+                                         target.getDeclContext());
+    if (auto coercedPattern = TypeChecker::coercePatternToType(
+            contextualPattern, patternType, patternOptions)) {
+      (*caseLabelItem)->setPattern(coercedPattern);
+    } else {
+      return None;
+    }
+
+    // If there is a guard expression, coerce that.
+    if (auto guardExpr = info.guardExpr) {
+      guardExpr = guardExpr->walk(*this);
+      if (!guardExpr)
+        return None;
+
+      // FIXME: Feels like we could leverage existing code more.
+      Type boolType = cs.getASTContext().getBoolDecl()->getDeclaredType();
+      guardExpr = solution.coerceToType(
+          guardExpr, boolType, cs.getConstraintLocator(info.guardExpr));
+      if (!guardExpr)
+        return None;
+
+      (*caseLabelItem)->setGuardExpr(guardExpr);
+      solution.setExprTypes(guardExpr);
+    }
+
     return target;
   } else {
     auto fn = *target.getAsFunction();
@@ -7747,5 +7786,17 @@ SolutionApplicationTarget SolutionApplicationTarget::walk(ASTWalker &walker) {
       condElement = *condElement.walk(walker);
     }
     return *this;
+
+  case Kind::caseLabelItem:
+    if (auto newPattern =
+            caseLabelItem.caseLabelItem->getPattern()->walk(walker)) {
+      caseLabelItem.caseLabelItem->setPattern(newPattern);
+    }
+    if (auto guardExpr = caseLabelItem.caseLabelItem->getGuardExpr()) {
+      if (auto newGuardExpr = guardExpr->walk(walker))
+        caseLabelItem.caseLabelItem->setGuardExpr(newGuardExpr);
+    }
+
+    return *this;
   }
 }

lib/Sema/CSGen.cpp

@@ -4223,6 +4223,62 @@ bool ConstraintSystem::generateConstraints(StmtCondition condition,
   return false;
 }
 
+bool ConstraintSystem::generateConstraints(
+    CaseStmt *caseStmt, DeclContext *dc, Type subjectType,
+    ConstraintLocator *locator) {
+  // Pre-bind all of the pattern variables within the case.
+  bindSwitchCasePatternVars(caseStmt);
+
+  for (auto &caseLabelItem : caseStmt->getMutableCaseLabelItems()) {
+    // Resolve the pattern.
+    auto *pattern = TypeChecker::resolvePattern(
+        caseLabelItem.getPattern(), dc, /*isStmtCondition=*/false);
+    if (!pattern)
+      return true;
+
+    // Generate constraints for the pattern, including one-way bindings for
+    // any variables that show up in this pattern, because those variables
+    // can be referenced in the guard expressions and the body.
+    Type patternType = generateConstraints(
+        pattern, locator, /* bindPatternVarsOneWay=*/true);
+
+    // Convert the subject type to the pattern, which establishes the
+    // bindings.
+    addConstraint(
+        ConstraintKind::Conversion, subjectType, patternType, locator);
+
+    // Generate constraints for the guard expression, if there is one.
+    Expr *guardExpr = caseLabelItem.getGuardExpr();
+    if (guardExpr) {
+      guardExpr = generateConstraints(guardExpr, dc);
+      if (!guardExpr)
+        return true;
+    }
+
+    // Save this info.
+    setCaseLabelItemInfo(&caseLabelItem, {pattern, guardExpr});
+
+    // For any pattern variable that has a parent variable (i.e., another
+    // pattern variable with the same name in the same case), require that
+    // the types be equivalent.
+    pattern->forEachVariable([&](VarDecl *var) {
+      if (auto parentVar = var->getParentVarDecl()) {
+        addConstraint(
+            ConstraintKind::Equal, getType(parentVar), getType(var), locator);
+      }
+    });
+  }
+
+  // Bind the types of the case body variables.
+  for (auto caseBodyVar : caseStmt->getCaseBodyVariablesOrEmptyArray()) {
+    auto parentVar = caseBodyVar->getParentVarDecl();
+    assert(parentVar && "Case body variables always have parents");
+    setType(caseBodyVar, getType(parentVar));
+  }
+
+  return false;
+}
+
 void ConstraintSystem::optimizeConstraints(Expr *e) {
   if (getASTContext().TypeCheckerOpts.DisableConstraintSolverPerformanceHacks)
     return;

lib/Sema/CSSolver.cpp

@@ -173,6 +173,7 @@ Solution ConstraintSystem::finalize() {
       contextualTypes.begin(), contextualTypes.end());
 
   solution.solutionApplicationTargets = solutionApplicationTargets;
+  solution.caseLabelItems = caseLabelItems;
 
   for (auto &e : CheckedConformances)
     solution.Conformances.push_back({e.first, e.second});
@@ -250,6 +251,11 @@ void ConstraintSystem::applySolution(const Solution &solution) {
     if (!getSolutionApplicationTarget(target.first))
       setSolutionApplicationTarget(target.first, target.second);
   }
+
+  // Register the statement condition targets.
+  for (const auto &info : solution.caseLabelItems) {
+    if (!getCaseLabelItemInfo(info.first))
+      setCaseLabelItemInfo(info.first, info.second);
   }
 
   // Register the conformances checked along the way to arrive to solution.
@@ -357,6 +363,13 @@ void truncate(llvm::MapVector<K, V> &map, unsigned newSize) {
     map.pop_back();
 }
 
+template <typename K, typename V, unsigned N>
+void truncate(llvm::SmallMapVector<K, V, N> &map, unsigned newSize) {
+  assert(newSize <= map.size() && "Not a truncation!");
+  for (unsigned i = 0, n = map.size() - newSize; i != n; ++i)
+    map.pop_back();
+}
+
 } // end anonymous namespace
 
 ConstraintSystem::SolverState::SolverState(
@@ -465,6 +478,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numInferredClosureTypes = cs.ClosureTypes.size();
   numContextualTypes = cs.contextualTypes.size();
   numSolutionApplicationTargets = cs.solutionApplicationTargets.size();
+  numCaseLabelItems = cs.caseLabelItems.size();
 
   PreviousScore = cs.CurrentScore;
 
@@ -545,6 +559,9 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any solution application targets.
   truncate(cs.solutionApplicationTargets, numSolutionApplicationTargets);
 
+  // Remove any case label item infos.
+  truncate(cs.caseLabelItems, numCaseLabelItems);
+
   // Reset the previous score.
   cs.CurrentScore = PreviousScore;