Merge pull request swiftlang#30174 from DougGregor/function-builder-switch

[Constraint system] Implement switch support for function builders.

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -626,12 +626,107 @@ 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) {
+      // Form a one-way constraint to prevent backward propagation.
+      subjectExpr = new (ctx) OneWayExpr(subjectExpr);
+
+      // 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)
@@ -1000,6 +1095,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 " \
@@ -1014,8 +1166,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

@@ -7461,7 +7461,7 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
 
       case StmtConditionElement::CK_PatternBinding: {
         ConstraintSystem &cs = solution.getConstraintSystem();
-        auto target = *cs.getStmtConditionTarget(&condElement);
+        auto target = *cs.getSolutionApplicationTarget(&condElement);
         auto resolvedTarget = rewriteTarget(target);
         if (!resolvedTarget)
           return None;
@@ -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();
@@ -7748,5 +7787,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

@@ -4218,7 +4218,7 @@ bool ConstraintSystem::generateConstraints(StmtCondition condition,
       if (generateConstraints(target, FreeTypeVariableBinding::Disallow))
         return true;
 
-      setStmtConditionTarget(&condElement, target);
+      setSolutionApplicationTarget(&condElement, target);
       continue;
     }
     }
@@ -4227,6 +4227,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

@@ -170,7 +170,8 @@ Solution ConstraintSystem::finalize() {
   solution.contextualTypes.assign(
       contextualTypes.begin(), contextualTypes.end());
 
-  solution.stmtConditionTargets = stmtConditionTargets;
+  solution.solutionApplicationTargets = solutionApplicationTargets;
+  solution.caseLabelItems = caseLabelItems;
 
   for (auto &e : CheckedConformances)
     solution.Conformances.push_back({e.first, e.second});
@@ -244,9 +245,15 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   }
 
   // Register the statement condition targets.
-  for (const auto &stmtConditionTarget : solution.stmtConditionTargets) {
-    if (!getStmtConditionTarget(stmtConditionTarget.first))
-      setStmtConditionTarget(stmtConditionTarget.first, stmtConditionTarget.second);
+  for (const auto &target : solution.solutionApplicationTargets) {
+    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.
@@ -354,6 +361,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(
@@ -461,7 +475,8 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numResolvedOverloads = cs.ResolvedOverloads.size();
   numInferredClosureTypes = cs.ClosureTypes.size();
   numContextualTypes = cs.contextualTypes.size();
-  numStmtConditionTargets = cs.stmtConditionTargets.size();
+  numSolutionApplicationTargets = cs.solutionApplicationTargets.size();
+  numCaseLabelItems = cs.caseLabelItems.size();
 
   PreviousScore = cs.CurrentScore;
 
@@ -539,8 +554,11 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any contextual types.
   truncate(cs.contextualTypes, numContextualTypes);
 
-  // Remove any statement condition types.
-  truncate(cs.stmtConditionTargets, numStmtConditionTargets);
+  // 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;