Move out-of-place ThenStmt checking to constraint generation

Author: hamishknight

include/swift/Sema/CSFix.h

@@ -458,6 +458,9 @@ enum class FixKind : uint8_t {
   /// Ignore an attempt to specialize non-generic type.
   AllowConcreteTypeSpecialization,
 
+  /// Ignore an out-of-place \c then statement.
+  IgnoreOutOfPlaceThenStmt,
+
   /// Ignore situations when provided number of generic arguments didn't match
   /// expected number of parameters.
   IgnoreGenericSpecializationArityMismatch,
@@ -3684,6 +3687,29 @@ class AllowConcreteTypeSpecialization final : public ConstraintFix {
   }
 };
 
+class IgnoreOutOfPlaceThenStmt final : public ConstraintFix {
+  IgnoreOutOfPlaceThenStmt(ConstraintSystem &cs, ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::IgnoreOutOfPlaceThenStmt, locator) {}
+
+public:
+  std::string getName() const override {
+    return "ignore out-of-place ThenStmt";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  bool diagnoseForAmbiguity(CommonFixesArray commonFixes) const override {
+    return diagnose(*commonFixes.front().first);
+  }
+
+  static IgnoreOutOfPlaceThenStmt *create(ConstraintSystem &cs,
+                                          ConstraintLocator *locator);
+
+  static bool classof(const ConstraintFix *fix) {
+    return fix->getKind() == FixKind::IgnoreOutOfPlaceThenStmt;
+  }
+};
+
 class IgnoreGenericSpecializationArityMismatch final : public ConstraintFix {
   ValueDecl *D;
   unsigned NumParams;

lib/Sema/CSDiagnostics.cpp

@@ -9158,6 +9158,11 @@ bool ConcreteTypeSpecialization::diagnoseAsError() {
   return true;
 }
 
+bool OutOfPlaceThenStmtFailure::diagnoseAsError() {
+  emitDiagnostic(diag::out_of_place_then_stmt);
+  return true;
+}
+
 bool InvalidTypeSpecializationArity::diagnoseAsError() {
   emitDiagnostic(diag::type_parameter_count_mismatch, D->getBaseIdentifier(),
                  NumParams, NumArgs, NumArgs < NumParams, HasParameterPack);

lib/Sema/CSDiagnostics.h

@@ -3073,6 +3073,16 @@ class ConcreteTypeSpecialization final : public FailureDiagnostic {
   bool diagnoseAsError() override;
 };
 
+/// Diagnose an out-of-place ThenStmt.
+class OutOfPlaceThenStmtFailure final : public FailureDiagnostic {
+public:
+  OutOfPlaceThenStmtFailure(const Solution &solution,
+                            ConstraintLocator *locator)
+      : FailureDiagnostic(solution, locator) {}
+
+  bool diagnoseAsError() override;
+};
+
 /// Diagnose attempts to specialize with invalid number of generic arguments:
 ///
 /// \code

lib/Sema/CSFix.cpp

@@ -2843,6 +2843,18 @@ AllowConcreteTypeSpecialization::create(ConstraintSystem &cs, Type concreteTy,
       AllowConcreteTypeSpecialization(cs, concreteTy, locator);
 }
 
+bool IgnoreOutOfPlaceThenStmt::diagnose(const Solution &solution,
+                                        bool asNote) const {
+  OutOfPlaceThenStmtFailure failure(solution, getLocator());
+  return failure.diagnose(asNote);
+}
+
+IgnoreOutOfPlaceThenStmt *
+IgnoreOutOfPlaceThenStmt::create(ConstraintSystem &cs,
+                                 ConstraintLocator *locator) {
+  return new (cs.getAllocator()) IgnoreOutOfPlaceThenStmt(cs, locator);
+}
+
 bool IgnoreGenericSpecializationArityMismatch::diagnose(
     const Solution &solution, bool asNote) const {
   InvalidTypeSpecializationArity failure(solution, D, NumParams, NumArgs,

lib/Sema/CSSimplify.cpp

@@ -15291,6 +15291,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::DestructureTupleToMatchPackExpansionParameter:
   case FixKind::AllowValueExpansionWithoutPackReferences:
   case FixKind::IgnoreInvalidPatternInExpr:
+  case FixKind::IgnoreOutOfPlaceThenStmt:
   case FixKind::IgnoreMissingEachKeyword:
     llvm_unreachable("handled elsewhere");
   }

lib/Sema/CSSyntacticElement.cpp

@@ -1255,11 +1255,21 @@ class SyntacticElementConstraintGenerator
     auto *resultExpr = thenStmt->getResult();
     auto contextInfo = cs.getContextualTypeInfo(resultExpr);
 
+    // First check to make sure the ThenStmt is in a valid position.
+    SmallVector<ThenStmt *, 4> validThenStmts;
+    if (auto SVE = context.getAsSingleValueStmtExpr())
+      (void)SVE.get()->getThenStmts(validThenStmts);
+
+    if (!llvm::is_contained(validThenStmts, thenStmt)) {
+      auto *thenLoc = cs.getConstraintLocator(thenStmt);
+      (void)cs.recordFix(IgnoreOutOfPlaceThenStmt::create(cs, thenLoc));
+    }
+
     // For an if/switch expression, if the contextual type for the branch is
     // still a type variable, we can drop it. This avoids needlessly
     // propagating the type of the branch to subsequent branches, instead
     // we'll let the join handle the conversion.
-    if (contextInfo && isExpr<SingleValueStmtExpr>(locator->getAnchor())) {
+    if (contextInfo) {
       auto contextualFixedTy =
           cs.getFixedTypeRecursive(contextInfo->getType(), /*wantRValue*/ true);
       if (contextualFixedTy->isTypeVariableOrMember())
@@ -2176,29 +2186,22 @@ class SyntacticElementSolutionApplication
   }
 
   ASTNode visitThenStmt(ThenStmt *thenStmt) {
-    // Note we're defensive here over whether we're in a SingleValueStmtExpr,
-    // as we don't diagnose out-of-place ThenStmts until MiscDiagnostics. If we
-    // have an out-of-place ThenStmt, just rewrite the expr without a contextual
-    // type.
     auto SVE = context.getAsSingleValueStmtExpr();
-    auto ty = SVE ? solution.getResolvedType(SVE.get()) : Type();
+    assert(SVE && "Should have diagnosed an out-of-place ThenStmt");
+    auto ty = solution.getResolvedType(SVE.get());
 
     // We need to fixup the conversion type to the full result type,
     // not the branch result type. This is necessary as there may be
     // an additional conversion required for the branch.
     auto target = solution.getTargetFor(thenStmt->getResult());
-    if (ty)
-      target->setExprConversionType(ty);
+    target->setExprConversionType(ty);
 
     auto *resultExpr = thenStmt->getResult();
     if (auto newResultTarget = rewriteTarget(*target))
       resultExpr = newResultTarget->getAsExpr();
 
     thenStmt->setResult(resultExpr);
 
-    if (!SVE)
-      return thenStmt;
-
     // If the expression was typed as Void, its branches are effectively
     // discarded, so treat them as ignored expressions.
     if (ty->lookThroughAllOptionalTypes()->isVoid()) {

lib/Sema/MiscDiagnostics.cpp

@@ -3846,7 +3846,6 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
   ASTContext &Ctx;
   DiagnosticEngine &Diags;
   llvm::DenseSet<SingleValueStmtExpr *> ValidSingleValueStmtExprs;
-  llvm::DenseSet<ThenStmt *> ValidThenStmts;
 
 public:
   SingleValueStmtUsageChecker(
@@ -3919,11 +3918,6 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
                        SVE->getStmt()->getKind());
       }
 
-      // Mark the valid 'then' statements.
-      SmallVector<ThenStmt *, 4> scratch;
-      for (auto *thenStmt : SVE->getThenStmts(scratch))
-        ValidThenStmts.insert(thenStmt);
-
       // Diagnose invalid SingleValueStmtExprs. This should only happen for
       // expressions in positions that we didn't support before
       // (e.g assignment or *explicit* return).
@@ -4008,11 +4002,8 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
     if (auto *TS = dyn_cast<ThrowStmt>(S))
       markValidSingleValueStmt(TS->getSubExpr());
 
-    if (auto *TS = dyn_cast<ThenStmt>(S)) {
+    if (auto *TS = dyn_cast<ThenStmt>(S))
       markValidSingleValueStmt(TS->getResult());
-      if (!ValidThenStmts.contains(TS))
-        Diags.diagnose(TS->getThenLoc(), diag::out_of_place_then_stmt);
-    }
 
     return Action::Continue(S);
   }