[Sema] Catch invalid if/switch exprs in more places

Move out-of-place SingleValueStmtExpr checking into
`performSyntacticExprDiagnostics`, to ensure we
catch all expressions. Previously we did the walk
as a part of Decl-based MiscDiagnostics, but it
turns out that can miss expressions in property
initializers, subscript default arguments, and
custom attrs.

This does mean that we'll now no longer diagnose
out-of-place if/switch exprs if the expression
didn't type-check, but that's consistent with the
rest of MiscDiagnostics, and I don't think it will
be a major issue in practice. We probably ought to
consider moving this checking into PreCheckExpr,
but that would require first separating out
SequenceExpr folding, which has other consequences,
and so I'm leaving as future work for now.

Author: hamishknight

lib/Sema/MiscDiagnostics.cpp

@@ -3850,7 +3850,20 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
   llvm::DenseSet<SingleValueStmtExpr *> ValidSingleValueStmtExprs;
 
 public:
-  SingleValueStmtUsageChecker(ASTContext &ctx) : Ctx(ctx), Diags(ctx.Diags) {}
+  SingleValueStmtUsageChecker(
+      ASTContext &ctx, ASTNode root,
+      llvm::Optional<ContextualTypePurpose> contextualPurpose)
+      : Ctx(ctx), Diags(ctx.Diags) {
+    assert(!root.is<Expr *>() || contextualPurpose &&
+           "Must provide contextual purpose for expr");
+
+    // If we have a contextual purpose, this is for an expression. Check if it's
+    // an expression in a valid position.
+    if (contextualPurpose) {
+      markAnyValidTopLevelSingleValueStmt(root.get<Expr *>(),
+                                          *contextualPurpose);
+    }
+  }
 
 private:
   /// Mark a given expression as a valid position for a SingleValueStmtExpr.
@@ -3862,8 +3875,23 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
       ValidSingleValueStmtExprs.insert(SVE);
   }
 
+  /// Mark a valid top-level expression with a given contextual purpose.
+  void markAnyValidTopLevelSingleValueStmt(Expr *E, ContextualTypePurpose ctp) {
+    // Allowed in returns, throws, and bindings.
+    switch (ctp) {
+    case CTP_ReturnStmt:
+    case CTP_ReturnSingleExpr:
+    case CTP_ThrowStmt:
+    case CTP_Initialization:
+      markValidSingleValueStmt(E);
+      break;
+    default:
+      break;
+    }
+  }
+
   MacroWalking getMacroWalkingBehavior() const override {
-    return MacroWalking::Expansion;
+    return MacroWalking::ArgumentsAndExpansion;
   }
 
   AssignExpr *findAssignment(Expr *E) const {
@@ -3989,28 +4017,33 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
     if (auto *PBD = dyn_cast<PatternBindingDecl>(D)) {
       for (auto idx : range(PBD->getNumPatternEntries()))
         markValidSingleValueStmt(PBD->getInit(idx));
+
+      return Action::Continue();
     }
-    // Valid as a single expression body of a function. This is needed in
-    // addition to ReturnStmt checking, as we will remove the return if the
-    // expression is inferred to be Never.
-    if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
-      if (AFD->hasSingleExpressionBody())
-        markValidSingleValueStmt(AFD->getSingleExpressionBody());
-    }
-    return Action::Continue();
+    // We don't want to walk into any other decl, we will visit them as part of
+    // typeCheckDecl.
+    return Action::SkipChildren();
   }
 };
 } // end anonymous namespace
 
+void swift::diagnoseOutOfPlaceExprs(
+    ASTContext &ctx, ASTNode root,
+    llvm::Optional<ContextualTypePurpose> contextualPurpose) {
+  // TODO: We ought to consider moving this into pre-checking such that we can
+  // still diagnose on invalid code, and don't have to traverse over implicit
+  // exprs. We need to first separate out SequenceExpr folding though.
+  SingleValueStmtUsageChecker sveChecker(ctx, root, contextualPurpose);
+  root.walk(sveChecker);
+}
+
 /// Apply the warnings managed by VarDeclUsageChecker to the top level
 /// code declarations that haven't been checked yet.
 void swift::
 performTopLevelDeclDiagnostics(TopLevelCodeDecl *TLCD) {
   auto &ctx = TLCD->getDeclContext()->getASTContext();
   VarDeclUsageChecker checker(TLCD, ctx.Diags);
   TLCD->walk(checker);
-  SingleValueStmtUsageChecker sveChecker(ctx);
-  TLCD->walk(sveChecker);
 }
 
 /// Perform diagnostics for func/init/deinit declarations.
@@ -4026,10 +4059,6 @@ void swift::performAbstractFuncDeclDiagnostics(AbstractFunctionDecl *AFD) {
     auto &ctx = AFD->getDeclContext()->getASTContext();
     VarDeclUsageChecker checker(AFD, ctx.Diags);
     AFD->walk(checker);
-
-    // Do a similar walk to check for out of place SingleValueStmtExprs.
-    SingleValueStmtUsageChecker sveChecker(ctx);
-    AFD->walk(sveChecker);
   }
 
   auto *body = AFD->getBody();
@@ -5864,10 +5893,10 @@ diagnoseDictionaryLiteralDuplicateKeyEntries(const Expr *E,
 //===----------------------------------------------------------------------===//
 
 /// Emit diagnostics for syntactic restrictions on a given expression.
-void swift::performSyntacticExprDiagnostics(const Expr *E,
-                                            const DeclContext *DC,
-                                            bool isExprStmt,
-                                            bool disableExprAvailabilityChecking) {
+void swift::performSyntacticExprDiagnostics(
+    const Expr *E, const DeclContext *DC,
+    llvm::Optional<ContextualTypePurpose> contextualPurpose, bool isExprStmt,
+    bool disableExprAvailabilityChecking, bool disableOutOfPlaceExprChecking) {
   auto &ctx = DC->getASTContext();
   TypeChecker::diagnoseSelfAssignment(E);
   diagSyntacticUseRestrictions(E, DC, isExprStmt);
@@ -5886,6 +5915,8 @@ void swift::performSyntacticExprDiagnostics(const Expr *E,
   diagnoseConstantArgumentRequirement(E, DC);
   diagUnqualifiedAccessToMethodNamedSelf(E, DC);
   diagnoseDictionaryLiteralDuplicateKeyEntries(E, DC);
+  if (!disableOutOfPlaceExprChecking)
+    diagnoseOutOfPlaceExprs(ctx, const_cast<Expr *>(E), contextualPurpose);
 }
 
 void swift::performStmtDiagnostics(const Stmt *S, DeclContext *DC) {

lib/Sema/MiscDiagnostics.h

@@ -28,6 +28,7 @@ namespace swift {
   class ApplyExpr;
   class CallExpr;
   class ClosureExpr;
+  enum ContextualTypePurpose : uint8_t;
   class DeclContext;
   class Decl;
   class Expr;
@@ -37,10 +38,22 @@ namespace swift {
   class ValueDecl;
   class ForEachStmt;
 
+/// Diagnose any expressions that appear in an unsupported position. If visiting
+/// an expression directly, its \p contextualPurpose should be provided to
+/// evaluate its position.
+void diagnoseOutOfPlaceExprs(
+    ASTContext &ctx, ASTNode root,
+    llvm::Optional<ContextualTypePurpose> contextualPurpose);
+
 /// Emit diagnostics for syntactic restrictions on a given expression.
+///
+/// Note: \p contextualPurpose must be non-nil, unless
+/// \p disableOutOfPlaceExprChecking is set to \c true.
 void performSyntacticExprDiagnostics(
     const Expr *E, const DeclContext *DC,
-    bool isExprStmt, bool disableExprAvailabilityChecking = false);
+    llvm::Optional<ContextualTypePurpose> contextualPurpose,
+    bool isExprStmt, bool disableExprAvailabilityChecking = false,
+    bool disableOutOfPlaceExprChecking = false);
 
 /// Emit diagnostics for a given statement.
 void performStmtDiagnostics(const Stmt *S, DeclContext *DC);

lib/Sema/TypeCheckConstraints.cpp

@@ -299,7 +299,11 @@ class FunctionSyntacticDiagnosticWalker : public ASTWalker {
   }
 
   PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
-    performSyntacticExprDiagnostics(expr, dcStack.back(), /*isExprStmt=*/false);
+    // We skip out-of-place expr checking here since we've already performed it.
+    performSyntacticExprDiagnostics(expr, dcStack.back(), /*ctp*/ llvm::None,
+                                    /*isExprStmt=*/false,
+                                    /*disableAvailabilityChecking*/ false,
+                                    /*disableOutOfPlaceExprChecking*/ true);
 
     if (auto closure = dyn_cast<ClosureExpr>(expr)) {
       if (closure->isSeparatelyTypeChecked()) {
@@ -346,8 +350,9 @@ void constraints::performSyntacticDiagnosticsForTarget(
   switch (target.kind) {
   case SyntacticElementTarget::Kind::expression: {
     // First emit diagnostics for the main expression.
-    performSyntacticExprDiagnostics(target.getAsExpr(), dc,
-                                    isExprStmt, disableExprAvailabilityChecking);
+    performSyntacticExprDiagnostics(
+        target.getAsExpr(), dc, target.getExprContextualTypePurpose(),
+        isExprStmt, disableExprAvailabilityChecking);
     return;
   }
 
@@ -356,17 +361,25 @@ void constraints::performSyntacticDiagnosticsForTarget(
 
     // First emit diagnostics for the main expression.
     performSyntacticExprDiagnostics(stmt->getTypeCheckedSequence(), dc,
-                                    isExprStmt,
+                                    CTP_ForEachSequence, isExprStmt,
                                     disableExprAvailabilityChecking);
 
     if (auto *whereExpr = stmt->getWhere())
-      performSyntacticExprDiagnostics(whereExpr, dc, /*isExprStmt*/ false);
+      performSyntacticExprDiagnostics(whereExpr, dc, CTP_Condition,
+                                      /*isExprStmt*/ false);
     return;
   }
 
   case SyntacticElementTarget::Kind::function: {
+    // Check for out of place expressions. This needs to be done on the entire
+    // function body rather than on individual expressions since we need the
+    // context of the parent nodes.
+    auto *body = target.getFunctionBody();
+    diagnoseOutOfPlaceExprs(dc->getASTContext(), body,
+                            /*contextualPurpose*/ llvm::None);
+
     FunctionSyntacticDiagnosticWalker walker(dc);
-    target.getFunctionBody()->walk(walker);
+    body->walk(walker);
     return;
   }
   case SyntacticElementTarget::Kind::closure:

test/Constraints/closures.swift

@@ -1149,11 +1149,9 @@ struct R_76250381<Result, Failure: Error> {
 // rdar://77022842 - crash due to a missing argument to a ternary operator
 func rdar77022842(argA: Bool? = nil, argB: Bool? = nil) {
   if let a = argA ?? false, if let b = argB ?? {
-    // expected-error@-1 {{'if' may only be used as expression in return, throw, or as the source of an assignment}}
-    // expected-error@-2 {{initializer for conditional binding must have Optional type, not 'Bool'}}
-    // expected-error@-3 {{cannot convert value of type '() -> ()' to expected argument type 'Bool?'}}
-    // expected-error@-4 {{cannot convert value of type 'Void' to expected condition type 'Bool'}}
-    // expected-error@-5 {{'if' must have an unconditional 'else' to be used as expression}}
+    // expected-error@-1 {{initializer for conditional binding must have Optional type, not 'Bool'}}
+    // expected-error@-2 {{cannot convert value of type '() -> ()' to expected argument type 'Bool?'}}
+    // expected-error@-3 {{cannot convert value of type 'Void' to expected condition type 'Bool'}}
   } // expected-error {{expected '{' after 'if' condition}}
 }
 

test/Constraints/if_expr.swift

@@ -387,7 +387,6 @@ func testReturnMismatch() {
   let _ = if .random() {
     return 1 // expected-error {{unexpected non-void return value in void function}}
     // expected-note@-1 {{did you mean to add a return type?}}
-    // expected-error@-2 {{cannot 'return' in 'if' when used as expression}}
   } else {
     0
   }
@@ -651,9 +650,46 @@ func builderWithBinding() -> Either<String, Int> {
   }
 }
 
+@Builder
+func builderWithInvalidBinding() -> Either<String, Int> {
+  let str = (if .random() { "a" } else { "b" })
+  // expected-error@-1 {{'if' may only be used as expression in return, throw, or as the source of an assignment}}
+  if .random() {
+    str
+  } else {
+    1
+  }
+}
+
+func takesBuilder(@Builder _ fn: () -> Either<String, Int>) {}
+
+func builderClosureWithBinding() {
+  takesBuilder {
+    // Make sure the binding gets type-checked as an if expression, but the
+    // other if block gets type-checked as a stmt.
+    let str = if .random() { "a" } else { "b" }
+    if .random() {
+      str
+    } else {
+      1
+    }
+  }
+}
+
+func builderClosureWithInvalidBinding()  {
+  takesBuilder {
+    let str = (if .random() { "a" } else { "b" })
+    // expected-error@-1 {{'if' may only be used as expression in return, throw, or as the source of an assignment}}
+    if .random() {
+      str
+    } else {
+      1
+    }
+  }
+}
+
 func builderInClosure() {
-  func build(@Builder _ fn: () -> Either<String, Int>) {}
-  build {
+  takesBuilder {
     if .random() {
       ""
     } else {

test/Constraints/switch_expr.swift

@@ -754,7 +754,7 @@ func builderNotPostfix() -> (Either<String, Int>, Bool) {
 
 @Builder
 func builderWithBinding() -> Either<String, Int> {
-  // Make sure the binding gets type-checked as an if expression, but the
+  // Make sure the binding gets type-checked as a switch expression, but the
   // other if block gets type-checked as a stmt.
   let str = switch Bool.random() {
     case true: "a"
@@ -767,9 +767,49 @@ func builderWithBinding() -> Either<String, Int> {
   }
 }
 
+
+@Builder
+func builderWithInvalidBinding() -> Either<String, Int> {
+  let str = (switch Bool.random() { default: "a" })
+  // expected-error@-1 {{'switch' may only be used as expression in return, throw, or as the source of an assignment}}
+  if .random() {
+    str
+  } else {
+    1
+  }
+}
+
+func takesBuilder(@Builder _ fn: () -> Either<String, Int>) {}
+
+func builderClosureWithBinding() {
+  takesBuilder {
+    // Make sure the binding gets type-checked as a switch expression, but the
+    // other if block gets type-checked as a stmt.
+    let str = switch Bool.random() { case true: "a" case false: "b" }
+    switch Bool.random() {
+    case true:
+      str
+    case false:
+      1
+    }
+  }
+}
+
+func builderClosureWithInvalidBinding()  {
+  takesBuilder {
+    let str = (switch Bool.random() { case true: "a" case false: "b" })
+    // expected-error@-1 {{'switch' may only be used as expression in return, throw, or as the source of an assignment}}
+    switch Bool.random() {
+    case true:
+      str
+    case false:
+      1
+    }
+  }
+}
+
 func builderInClosure() {
-  func build(@Builder _ fn: () -> Either<String, Int>) {}
-  build {
+  takesBuilder {
     switch Bool.random() {
     case true:
       ""

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -1998,6 +1998,20 @@ public struct NestedMagicLiteralMacro: ExpressionMacro {
   }
 }
 
+public struct InvalidIfExprMacro: MemberMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingMembersOf decl: some DeclGroupSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    return ["""
+      func bar() {
+        let _ = (if .random() { 0 } else { 1 })
+      }
+      """]
+  }
+}
+
 public struct InitWithProjectedValueWrapperMacro: PeerMacro {
   public static func expansion(
     of node: AttributeSyntax,

test/Macros/if_expr.swift

@@ -0,0 +1,18 @@
+// RUN: %empty-directory(%t)
+// RUN: %host-build-swift -swift-version 5 -emit-library -o %t/%target-library-name(MacroDefinition) -module-name=MacroDefinition %S/Inputs/syntax_macro_definitions.swift -g -no-toolchain-stdlib-rpath
+
+// RUN: not %target-swift-frontend -typecheck %s -swift-version 5 -load-plugin-library %t/%target-library-name(MacroDefinition) -serialize-diagnostics-path %t/diags.dia
+
+// RUN: c-index-test -read-diagnostics %t/diags.dia 2>&1 | %FileCheck -check-prefix DIAGS %s
+
+// REQUIRES: swift_swift_parser
+
+@attached(member, names: named(bar))
+macro TestMacro<T>(_ x: T) = #externalMacro(module: "MacroDefinition", type: "InvalidIfExprMacro")
+
+// Make sure we diagnose both the use in the custom attribute and in the expansion.
+// DIAGS-DAG: if_expr.swift:[[@LINE+1]]:12: error: 'if' may only be used as expression in return, throw, or as the source of an assignment
+@TestMacro(if .random() { 0 } else { 0 })
+struct S {}
+
+// DIAGS-DAG: @__swiftmacro_7if_expr1S9TestMacrofMm_.swift:2:12: error: 'if' may only be used as expression in return, throw, or as the source of an assignment