Allow implicit last expression results for if/switch expressions

Allow implicitly treating the last expression of
the branch as the result, behind the experimental
feature `ImplicitLastExprResults`.

Author: hamishknight

include/swift/AST/Expr.h

@@ -6204,12 +6204,19 @@ class SingleValueStmtExpr : public Expr {
   /// SingleValueStmtExpr.
   static SingleValueStmtExpr *tryDigOutSingleValueStmtExpr(Expr *E);
 
+  /// Whether the last ASTNode in the given BraceStmt can potentially be used as
+  /// the implicit result for a SingleValueStmtExpr. If \p mustBeSingleValueStmt
+  /// is \c true, a result will be considered even if it may not be valid.
+  static bool isLastElementImplicitResult(BraceStmt *BS, ASTContext &ctx,
+                                          bool mustBeSingleValueStmt);
+
   /// Retrieves a resulting ThenStmt from the given BraceStmt, or \c nullptr if
   /// the brace does not have a resulting ThenStmt.
   static ThenStmt *getThenStmtFrom(BraceStmt *BS);
 
   /// Whether the given BraceStmt has a result to be produced from a parent
-  /// SingleValueStmtExpr.
+  /// SingleValueStmtExpr. Note this does not consider elements that may
+  /// implicitly become results, check \c isLastElementImplicitResult for that.
   static bool hasResult(BraceStmt *BS) {
     return getThenStmtFrom(BS);
   }

include/swift/AST/TypeCheckRequests.h

@@ -4258,9 +4258,13 @@ class IsSingleValueStmtResult {
     return InvalidJumps;
   }
 
-  explicit operator bool() const {
+  bool isValid() const {
     return TheKind == Kind::Valid;
   }
+
+  explicit operator bool() const {
+    return isValid();
+  }
 };
 
 /// Computes whether a given statement can be treated as a SingleValueStmtExpr.

include/swift/Basic/Features.def

@@ -231,6 +231,10 @@ EXPERIMENTAL_FEATURE(ThenStatements, false)
 /// Enable 'do' expressions.
 EXPERIMENTAL_FEATURE(DoExpressions, false)
 
+/// Enable implicitly treating the last expression in an 'if'/'switch'
+/// expression as the result
+EXPERIMENTAL_FEATURE(ImplicitLastExprResults, false)
+
 /// Enable the `@_rawLayout` attribute.
 EXPERIMENTAL_FEATURE(RawLayout, true)
 

lib/AST/ASTPrinter.cpp

@@ -3846,6 +3846,10 @@ static bool usesFeatureDoExpressions(Decl *decl) {
   return false;
 }
 
+static bool usesFeatureImplicitLastExprResults(Decl *decl) {
+  return false;
+}
+
 static bool usesFeatureNewCxxMethodSafetyHeuristics(Decl *decl) {
   return decl->hasClangNode();
 }

lib/AST/Expr.cpp

@@ -2550,37 +2550,21 @@ SingleValueStmtExpr *SingleValueStmtExpr::createWithWrappedBranches(
     if (!BS)
       continue;
 
-    if (auto *S = BS->getSingleActiveStatement()) {
-      if (mustBeExpr) {
-        // If this must be an expression, we can eagerly wrap any exhaustive if,
-        // switch, and do statement.
-        if (auto *IS = dyn_cast<IfStmt>(S)) {
-          if (!IS->isSyntacticallyExhaustive())
-            continue;
-        } else if (auto *DCS = dyn_cast<DoCatchStmt>(S)) {
-          if (!ctx.LangOpts.hasFeature(Feature::DoExpressions))
-            continue;
-          if (!DCS->isSyntacticallyExhaustive())
-            continue;
-        } else if (isa<DoStmt>(S)) {
-          if (!ctx.LangOpts.hasFeature(Feature::DoExpressions))
-            continue;
-        } else if (!isa<SwitchStmt>(S)) {
-          continue;
-        }
-      } else {
-        // Otherwise do the semantic checking to verify that we can wrap the
-        // branch.
-        if (!S->mayProduceSingleValue(ctx))
-          continue;
-      }
-      BS->setLastElement(SingleValueStmtExpr::createWithWrappedBranches(
-          ctx, S, DC, mustBeExpr));
-    }
+    // Check to see if we can wrap an implicit last element of the brace.
+    if (!isLastElementImplicitResult(BS, ctx, mustBeExpr))
+      continue;
+
+    auto &result = BS->getElements().back();
+    assert(result.is<Expr *>() || result.is<Stmt *>());
 
-    // Wrap single expression branches in an implicit 'then <expr>'.
-    if (auto *E = BS->getSingleActiveExpression())
-      BS->setLastElement(ThenStmt::createImplicit(ctx, E));
+    // Wrap a statement in a SingleValueStmtExpr.
+    if (auto *S = result.dyn_cast<Stmt *>()) {
+      result = SingleValueStmtExpr::createWithWrappedBranches(ctx, S, DC,
+                                                              mustBeExpr);
+    }
+    // Wrap an expression in an implicit 'then <expr>'.
+    if (auto *E = result.dyn_cast<Expr *>())
+      result = ThenStmt::createImplicit(ctx, E);
   }
   return SVE;
 }
@@ -2630,6 +2614,50 @@ SingleValueStmtExpr::tryDigOutSingleValueStmtExpr(Expr *E) {
   return finder.FoundSVE;
 }
 
+bool SingleValueStmtExpr::isLastElementImplicitResult(
+    BraceStmt *BS, ASTContext &ctx, bool mustBeSingleValueStmt) {
+  if (BS->empty())
+    return false;
+
+  // We must either be allowing implicit last expressions, or we must have a
+  // single active element, which is guaranteed to be the last element.
+  if (!ctx.LangOpts.hasFeature(Feature::ImplicitLastExprResults) &&
+      !BS->getSingleActiveElement()) {
+    return false;
+  }
+  auto elt = BS->getLastElement();
+
+  // Expressions are always valid.
+  if (elt.is<Expr *>())
+    return true;
+
+  if (auto *S = elt.dyn_cast<Stmt *>()) {
+    if (mustBeSingleValueStmt) {
+      // If this must be a SingleValueStmtExpr, we can eagerly take any
+      // exhaustive if, switch, and do statement.
+      if (auto *IS = dyn_cast<IfStmt>(S))
+        return IS->isSyntacticallyExhaustive();
+
+      // Guaranteed to be exhaustive.
+      if (isa<SwitchStmt>(S))
+        return true;
+
+      if (ctx.LangOpts.hasFeature(Feature::DoExpressions)) {
+        if (auto *DCS = dyn_cast<DoCatchStmt>(S))
+          return DCS->isSyntacticallyExhaustive();
+
+        if (isa<DoStmt>(S))
+          return true;
+      }
+      return false;
+    }
+    // Otherwise do the semantic checking to verify the statement produces
+    // a single value.
+    return S->mayProduceSingleValue(ctx).isValid();
+  }
+  return false;
+}
+
 ThenStmt *SingleValueStmtExpr::getThenStmtFrom(BraceStmt *BS) {
   if (BS->empty())
    return nullptr;

lib/Sema/MiscDiagnostics.cpp

@@ -3922,8 +3922,9 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
               continue;
             }
           }
-          // TODO: If 'then' statements are enabled by default, the wording of
-          // this diagnostic should be tweaked.
+          // TODO: The wording of this diagnostic will need tweaking if either
+          // implicit last expressions or 'then' statements are enabled by
+          // default.
           Diags.diagnose(branch->getEndLoc(),
                          diag::single_value_stmt_branch_must_end_in_result,
                          S->getKind());

lib/Sema/TypeCheckStmt.cpp

@@ -3057,26 +3057,6 @@ static bool doesBraceEndWithThrow(BraceStmt *BS) {
   return isa<ThrowStmt>(S);
 }
 
-/// Whether the given brace statement is considered to produce a result for
-/// an if/switch expression.
-static bool doesBraceProduceResult(BraceStmt *BS, ASTContext &ctx) {
-  if (BS->empty())
-    return false;
-
-  // We consider the branch as having a result if there is:
-  // - A single active expression or statement that can be turned into an
-  //   expression.
-  // - 'then <expr>' as the last statement.
-  if (BS->getSingleActiveExpression())
-    return true;
-
-  if (auto *S = BS->getSingleActiveStatement()) {
-    if (S->mayProduceSingleValue(ctx))
-      return true;
-  }
-  return SingleValueStmtExpr::hasResult(BS);
-}
-
 IsSingleValueStmtResult
 areBranchesValidForSingleValueStmt(ASTContext &ctx, ArrayRef<Stmt *> branches) {
   TinyPtrVector<Stmt *> invalidJumps;
@@ -3094,8 +3074,10 @@ areBranchesValidForSingleValueStmt(ASTContext &ctx, ArrayRef<Stmt *> branches) {
     // Check to see if there are any invalid jumps.
     BS->walk(jumpFinder);
 
-    // Check to see if a result is produced from the branch.
-    if (doesBraceProduceResult(BS, ctx)) {
+    // Must either have an explicit or implicit result for the branch.
+    if (SingleValueStmtExpr::hasResult(BS) ||
+        SingleValueStmtExpr::isLastElementImplicitResult(
+            BS, ctx, /*mustBeSingleValueStmt*/ false)) {
       hadResult = true;
       continue;
     }

test/Constraints/if_switch_implicit_last_expr.swift

@@ -0,0 +1,168 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-feature ImplicitLastExprResults
+
+// Required for experimental features
+// REQUIRES: asserts
+
+func testMismatch1() -> Int {
+  if .random() {
+    print("hello")
+    0
+  } else {
+    print("hi")
+    "" // expected-error {{cannot convert value of type 'String' to specified type 'Int'}}
+  }
+}
+
+func testMismatch2() -> Int {
+  let x = if .random() {
+    print("hello")
+    0 // expected-error {{branches have mismatching types 'Int' and 'String'}}
+  } else {
+    print("hi")
+    ""
+  }
+  return x
+}
+
+func testVoidNeverConversion() {
+  // We allow the T -> Void and Never -> T conversions, same as the single
+  // expression case.
+  let _: () -> Void = {
+    if .random() {
+      print("hello")
+      1 // expected-warning {{integer literal is unused}}
+    } else {
+      print("there")
+      2 // expected-warning {{integer literal is unused}}
+    }
+  }
+  let _: () -> Void = {
+    if .random() {
+      print("hello")
+      fatalError()
+    } else {
+      0 // expected-warning {{integer literal is unused}}
+    }
+  }
+  // We fall back to Void if we have a mismatch.
+  let _ = {
+    if .random() {
+      ()
+      "" // expected-warning {{string literal is unused}}
+    } else {
+      ()
+      0 // expected-warning {{integer literal is unused}}
+    }
+  }
+  // Unless there's a contextual type.
+  let _ = { () -> String in
+    if .random() {
+      ()
+      ""
+    } else {
+      ()
+      0 // expected-error {{cannot convert value of type 'Int' to specified type 'String'}}
+    }
+  }
+  let _: () -> String = {
+    if .random() {
+      ()
+      ""
+    } else {
+      ()
+      0 // expected-error {{cannot convert value of type 'Int' to specified type 'String'}}
+    }
+  }
+  let _ = {
+    if .random() {
+      ()
+      ()
+    } else {
+      0 // expected-warning {{integer literal is unused}}
+    }
+  }
+  let _ = {
+    switch Bool.random() {
+    case true:
+      print("hello")
+      ""
+    case false:
+      fatalError()
+    }
+  }
+  func foo() -> Int {
+    switch Bool.random() {
+    case true:
+      fatalError()
+    case false:
+      ()
+      fatalError()
+    }
+  }
+  func bar() -> Int {
+    switch Bool.random() {
+    case true:
+      ()
+      fatalError()
+    case false:
+      ()
+      0
+    }
+  }
+}
+
+enum Either<T, U> {
+  case first(T), second(U)
+}
+
+@resultBuilder
+struct Builder {
+  static func buildBlock<T>(_ x: T) -> T { x }
+  static func buildBlock<T, U>(_ x: T, _ y: U) -> (T, U) { (x, y) }
+
+  static func buildEither<T, U>(first x: T) -> Either<T, U> { .first(x) }
+  static func buildEither<T, U>(second x: U) -> Either<T, U> { .second(x) }
+
+  static func buildExpression(_ x: Double) -> Double { x }
+  static func buildExpression<T>(_ x: T) -> T { x }
+}
+
+// Implicit 'then' statements are transparent to the result builder transform.
+@Builder
+func testBuilder1() -> Either<(Void, Int), (Void, Int)> {
+  if .random() {
+    ()
+    0
+  } else {
+    ()
+    1
+  }
+}
+
+@Builder
+func testBuilder2() -> (Either<(Void, Int), String>, Void) {
+  if .random() {
+    ()
+    0
+  } else {
+    ""
+  }
+  ()
+}
+
+@Builder
+func testBuilder3() -> Either<Int, String> {
+  // Bindings should still work though.
+  let x = if .random() {
+    print("hello")
+    0
+  } else {
+    ()
+    1
+  }
+  if .random() {
+    x
+  } else {
+    ""
+  }
+}