Introduce then statements

These allow multi-statement `if`/`switch` expression
branches that can produce a value at the end by
saying `then <expr>`. This is gated behind
`-enable-experimental-feature ThenStatements`
pending evolution discussion.

Author: hamishknight

include/swift/AST/DiagnosticsParse.def

@@ -1248,8 +1248,9 @@ ERROR(case_stmt_without_body,none,
 // 'try' on statements
 ERROR(try_on_stmt,none,
       "'try' cannot be used with '%0'", (StringRef))
-ERROR(try_on_return_throw_yield,none,
-      "'try' must be placed on the %select{returned|thrown|yielded}0 expression",
+ERROR(try_must_come_after_stmt,none,
+      "'try' must be placed on the %select{returned|thrown|yielded|produced}0 "
+      "expression",
       (unsigned))
 ERROR(try_on_var_let,none,
       "'try' must be placed on the initial value expression", ())
@@ -1403,6 +1404,12 @@ ERROR(expected_expr_after_each, none,
       "expected expression after 'each'", ())
 ERROR(expected_expr_after_repeat, none,
       "expected expression after 'repeat'", ())
+ERROR(expected_expr_after_then,none,
+      "expected expression after 'then'", ())
+
+WARNING(unindented_code_after_then,none,
+        "expression following 'then' is treated as an argument of "
+        "the 'then'", ())
 
 WARNING(move_consume_final_spelling, none,
         "'_move' has been renamed to 'consume', and the '_move' spelling will be removed shortly", ())

include/swift/AST/DiagnosticsSema.def

@@ -1225,14 +1225,18 @@ ERROR(if_expr_must_be_syntactically_exhaustive,none,
 ERROR(single_value_stmt_branch_empty,none,
       "expected expression in branch of '%0' expression",
       (StmtKind))
-ERROR(single_value_stmt_branch_must_end_in_throw,none,
-      "non-expression branch of '%0' expression may only end with a 'throw'",
+ERROR(single_value_stmt_branch_must_end_in_result,none,
+      "non-expression branch of '%0' expression may only end with a 'throw' "
+      "or 'then'",
       (StmtKind))
 ERROR(cannot_jump_in_single_value_stmt,none,
       "cannot '%0' in '%1' when used as expression",
       (StmtKind, StmtKind))
 ERROR(effect_marker_on_single_value_stmt,none,
       "'%0' may not be used on '%1' expression", (StringRef, StmtKind))
+ERROR(out_of_place_then_stmt,none,
+      "'then' may only appear as the last statement in an 'if' or 'switch' "
+      "expression", ())
 
 ERROR(did_not_call_function_value,none,
       "function value was used as a property; add () to call it",

include/swift/AST/Expr.h

@@ -70,6 +70,7 @@ namespace swift {
   class CallExpr;
   class KeyPathExpr;
   class CaptureListExpr;
+  class ThenStmt;
 
 enum class ExprKind : uint8_t {
 #define EXPR(Id, Parent) Id,
@@ -6125,10 +6126,10 @@ class SingleValueStmtExpr : public Expr {
   SingleValueStmtExpr(Stmt *S, DeclContext *DC)
       : Expr(ExprKind::SingleValueStmt, /*isImplicit*/ true), S(S), DC(DC) {}
 
-public:
   /// Creates a new SingleValueStmtExpr wrapping a statement.
   static SingleValueStmtExpr *create(ASTContext &ctx, Stmt *S, DeclContext *DC);
 
+public:
   /// Creates a new SingleValueStmtExpr wrapping a statement, and recursively
   /// attempts to wrap any branches of that statement that can become single
   /// value statement expressions.
@@ -6144,6 +6145,16 @@ class SingleValueStmtExpr : public Expr {
   /// SingleValueStmtExpr.
   static SingleValueStmtExpr *tryDigOutSingleValueStmtExpr(Expr *E);
 
+  /// Retrieves a resulting ThenStmt from the given BraceStmt, or \c nullptr if
+  /// the brace does not have a resulting ThenStmt.
+  static ThenStmt *getResultStmtFrom(BraceStmt *BS);
+
+  /// Whether the given BraceStmt has a result to be produced from a parent
+  /// SingleValueStmtExpr.
+  static bool hasResult(BraceStmt *BS) {
+    return getResultStmtFrom(BS);
+  }
+
   /// Retrieve the wrapped statement.
   Stmt *getStmt() const { return S; }
   void setStmt(Stmt *newS) { S = newS; }
@@ -6154,10 +6165,15 @@ class SingleValueStmtExpr : public Expr {
   /// Retrieve the complete set of branches for the underlying statement.
   ArrayRef<Stmt *> getBranches(SmallVectorImpl<Stmt *> &scratch) const;
 
-  /// Retrieve the single expression branches of the statement, excluding
-  /// branches that either have multiple expressions, or have statements.
+  /// Retrieve the result statements from each branch of the
+  /// SingleValueStmtExpr.
+  ArrayRef<ThenStmt *>
+  getResultStmts(SmallVectorImpl<ThenStmt *> &scratch) const;
+
+  /// Retrieve the result expressions from each branch of the
+  /// SingleValueStmtExpr.
   ArrayRef<Expr *>
-  getSingleExprBranches(SmallVectorImpl<Expr *> &scratch) const;
+  getResultExprs(SmallVectorImpl<Expr *> &scratch) const;
 
   DeclContext *getDeclContext() const { return DC; }
 

include/swift/AST/NameLookup.h

@@ -667,6 +667,7 @@ class FindLocalVal : public StmtVisitor<FindLocalVal> {
   void visitFailStmt(FailStmt *) {}
   void visitReturnStmt(ReturnStmt *) {}
   void visitYieldStmt(YieldStmt *) {}
+  void visitThenStmt(ThenStmt *) {}
   void visitThrowStmt(ThrowStmt *) {}
   void visitDiscardStmt(DiscardStmt *) {}
   void visitPoundAssertStmt(PoundAssertStmt *) {}

include/swift/AST/Stmt.h

@@ -304,6 +304,38 @@ class YieldStmt final
   static bool classof(const Stmt *S) { return S->getKind() == StmtKind::Yield; }
 };
 
+/// The statement `then <expr>`. This is used within if/switch expressions to
+/// indicate the value being produced by a given branch.
+class ThenStmt : public Stmt {
+  SourceLoc ThenLoc;
+  Expr *Result;
+
+  ThenStmt(SourceLoc thenLoc, Expr *result, bool isImplicit)
+      : Stmt(StmtKind::Then, isImplicit), ThenLoc(thenLoc), Result(result) {
+    assert(Result && "Must have non-null result");
+  }
+
+public:
+  /// Create a new parsed ThenStmt.
+  static ThenStmt *createParsed(ASTContext &ctx, SourceLoc thenLoc,
+                                Expr *result);
+
+  /// Create an implicit ThenStmt.
+  ///
+  /// Note that such statements will be elided during the result builder
+  /// transform.
+  static ThenStmt *createImplicit(ASTContext &ctx, Expr *result);
+
+  SourceLoc getThenLoc() const { return ThenLoc; }
+
+  SourceRange getSourceRange() const;
+
+  Expr *getResult() const { return Result; }
+  void setResult(Expr *e) { Result = e; }
+
+  static bool classof(const Stmt *S) { return S->getKind() == StmtKind::Then; }
+};
+
 /// DeferStmt - A 'defer' statement.  This runs the substatement it contains
 /// when the enclosing scope is exited.
 ///

include/swift/AST/StmtNodes.def

@@ -46,6 +46,7 @@
 
 STMT(Brace, Stmt)
 STMT(Return, Stmt)
+STMT(Then, Stmt)
 STMT(Yield, Stmt)
 STMT(Defer, Stmt)
 ABSTRACT_STMT(Labeled, Stmt)

include/swift/AST/TypeCheckRequests.h

@@ -3906,8 +3906,8 @@ class IsSingleValueStmtResult {
     /// The statement is an 'if' statement without an unconditional 'else'.
     NonExhaustiveIf,
 
-    /// There are no single-expression branches.
-    NoExpressionBranches,
+    /// There is no branch that produces a resulting value.
+    NoResult,
 
     /// There is an unhandled statement branch. This should only be the case
     /// for invalid AST.
@@ -3962,8 +3962,8 @@ class IsSingleValueStmtResult {
   static IsSingleValueStmtResult nonExhaustiveIf() {
     return IsSingleValueStmtResult(Kind::NonExhaustiveIf);
   }
-  static IsSingleValueStmtResult noExpressionBranches() {
-    return IsSingleValueStmtResult(Kind::NoExpressionBranches);
+  static IsSingleValueStmtResult noResult() {
+    return IsSingleValueStmtResult(Kind::NoResult);
   }
   static IsSingleValueStmtResult unhandledStmt() {
     return IsSingleValueStmtResult(Kind::UnhandledStmt);

include/swift/Basic/Features.def

@@ -224,6 +224,9 @@ EXPERIMENTAL_FEATURE(SendNonSendable, false)
 /// "playground transform" is enabled.
 EXPERIMENTAL_FEATURE(PlaygroundExtendedCallbacks, true)
 
+/// Enable 'then' statements.
+EXPERIMENTAL_FEATURE(ThenStatements, false)
+
 /// Enable the `@_rawLayout` attribute.
 EXPERIMENTAL_FEATURE(RawLayout, true)
 

include/swift/IDE/CodeCompletionResult.h

@@ -213,6 +213,7 @@ enum class CompletionKind : uint8_t {
   CallArg,
   LabeledTrailingClosure,
   ReturnStmtExpr,
+  ThenStmtExpr,
   YieldStmtExpr,
   ForEachSequence,
 

include/swift/Parse/IDEInspectionCallbacks.h

@@ -235,6 +235,8 @@ class CodeCompletionCallbacks {
 
   virtual void completeReturnStmt(CodeCompletionExpr *E) {};
 
+  virtual void completeThenStmt(CodeCompletionExpr *E) {};
+
   /// Complete a yield statement.  A missing yield index means that the
   /// completion immediately follows the 'yield' keyword; it may be either
   /// an expression or a parenthesized expression list.  A present yield

include/swift/Parse/Parser.h

@@ -602,6 +602,10 @@ class Parser {
             cast<AccessorDecl>(CurDeclContext)->isCoroutine());
   }
 
+  /// Whether the current token is the contextual keyword for a \c then
+  /// statement.
+  bool isContextualThenKeyword(bool preferExpr = false);
+
   /// `discard self` is the only valid phrase, but we peek ahead for just any
   /// identifier after `discard` to determine if it's the statement. This helps
   /// us avoid interpreting `discard(self)` as the statement and not a call.
@@ -1924,7 +1928,12 @@ class Parser {
   //===--------------------------------------------------------------------===//
   // Statement Parsing
 
-  bool isStartOfStmt();
+  /// Whether we are at the start of a statement.
+  ///
+  /// \param preferExpr If \c true, if either a statement or expression
+  /// could be parsed, prefer to return \c false such that an expr is parsed.
+  bool isStartOfStmt(bool preferExpr = false);
+
   bool isTerminatorForBraceItemListKind(BraceItemListKind Kind,
                                         ArrayRef<ASTNode> ParsedDecls);
   ParserResult<Stmt> parseStmt();
@@ -1933,6 +1942,7 @@ class Parser {
   ParserResult<Stmt> parseStmtContinue();
   ParserResult<Stmt> parseStmtReturn(SourceLoc tryLoc);
   ParserResult<Stmt> parseStmtYield(SourceLoc tryLoc);
+  ParserResult<Stmt> parseStmtThen(SourceLoc tryLoc);
   ParserResult<Stmt> parseStmtThrow(SourceLoc tryLoc);
   ParserResult<Stmt> parseStmtDiscard();
   ParserResult<Stmt> parseStmtDefer();

include/swift/Parse/Token.h

@@ -108,7 +108,39 @@ class Token {
   bool isBinaryOperator() const {
     return Kind == tok::oper_binary_spaced || Kind == tok::oper_binary_unspaced;
   }
-  
+
+  /// Checks whether the token is either a binary operator, or is a token that
+  /// acts like a binary operator (e.g infix '=', '?', '.').
+  bool isBinaryOperatorLike() const {
+    if (isBinaryOperator())
+      return true;
+
+    switch (Kind) {
+    case tok::period:
+    case tok::equal:
+    case tok::arrow:
+    case tok::question_infix:
+      return true;
+    default:
+      return false;
+    }
+    llvm_unreachable("Unhandled case in switch!");
+  }
+
+  /// Checks whether the token is either a postfix operator, or is a token that
+  /// acts like a postfix operator (e.g postfix '!' and '?').
+  bool isPostfixOperatorLike() const {
+    switch (Kind) {
+    case tok::oper_postfix:
+    case tok::exclaim_postfix:
+    case tok::question_postfix:
+      return true;
+    default:
+      return false;
+    }
+    llvm_unreachable("Unhandled case in switch!");
+  }
+
   bool isAnyOperator() const {
     return isBinaryOperator() || Kind == tok::oper_postfix ||
            Kind == tok::oper_prefix;

include/swift/Sema/ConstraintLocator.h

@@ -103,8 +103,14 @@ enum class ConversionRestrictionKind;
 
 /// The kind of SingleValueStmtExpr branch the locator identifies.
 enum class SingleValueStmtBranchKind {
-  Regular,
-  InSingleExprClosure
+  /// Explicitly written 'then <expr>'.
+  Explicit,
+
+  /// Implicitly written '<expr>'.
+  Implicit,
+
+  /// Implicitly written '<expr>' in a single expr closure body.
+  ImplicitInInSingleExprClosure
 };
 
 /// Locates a given constraint within the expression being
@@ -935,19 +941,19 @@ class LocatorPathElt::TernaryBranch final : public StoredIntegerElement<1> {
   }
 };
 
-/// The branch of a SingleValueStmtExpr. Note the stored index corresponds to
-/// the expression branches, i.e it skips statement branch indices.
-class LocatorPathElt::SingleValueStmtBranch final
+/// A particular result of a ThenStmt at a given index in a SingleValueStmtExpr.
+/// The stored index corresponds to the \c getResultExprs array.
+class LocatorPathElt::SingleValueStmtResult final
     : public StoredIntegerElement<1> {
 public:
-  SingleValueStmtBranch(unsigned exprIdx)
-      : StoredIntegerElement(ConstraintLocator::SingleValueStmtBranch,
+  SingleValueStmtResult(unsigned exprIdx)
+      : StoredIntegerElement(ConstraintLocator::SingleValueStmtResult,
                              exprIdx) {}
 
-  unsigned getExprBranchIndex() const { return getValue(); }
+  unsigned getIndex() const { return getValue(); }
 
   static bool classof(const LocatorPathElt *elt) {
-    return elt->getKind() == ConstraintLocator::SingleValueStmtBranch;
+    return elt->getKind() == ConstraintLocator::SingleValueStmtResult;
   }
 };
 

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -259,8 +259,8 @@ CUSTOM_LOCATOR_PATH_ELT(SyntacticElement)
 /// The element of the pattern binding declaration.
 CUSTOM_LOCATOR_PATH_ELT(PatternBindingElement)
 
-/// A branch of a SingleValueStmtExpr.
-CUSTOM_LOCATOR_PATH_ELT(SingleValueStmtBranch)
+/// A particular result of a ThenStmt at a given index in a SingleValueStmtExpr.
+CUSTOM_LOCATOR_PATH_ELT(SingleValueStmtResult)
 
 /// A declaration introduced by a pattern: name (i.e. `x`) or `_`
 ABSTRACT_LOCATOR_PATH_ELT(PatternDecl)

lib/AST/ASTDumper.cpp

@@ -1640,6 +1640,13 @@ class PrintStmt : public StmtVisitor<PrintStmt> {
     PrintWithColorRAII(OS, ParenthesisColor) << ')';
   }
 
+  void visitThenStmt(ThenStmt *S) {
+    printCommon(S, "then_stmt");
+    OS << '\n';
+    printRec(S->getResult());
+    PrintWithColorRAII(OS, ParenthesisColor) << ')';
+  }
+
   void visitDeferStmt(DeferStmt *S) {
     printCommon(S, "defer_stmt") << '\n';
     printRec(S->getTempDecl());

lib/AST/ASTPrinter.cpp

@@ -3508,6 +3508,10 @@ static bool usesFeaturePlaygroundExtendedCallbacks(Decl *decl) {
   return false;
 }
 
+static bool usesFeatureThenStatements(Decl *decl) {
+  return false;
+}
+
 static bool usesFeatureNewCxxMethodSafetyHeuristics(Decl *decl) {
   return decl->hasClangNode();
 }
@@ -5520,6 +5524,16 @@ void PrintAST::visitYieldStmt(YieldStmt *stmt) {
   if (parens) Printer << ")";
 }
 
+void PrintAST::visitThenStmt(ThenStmt *stmt) {
+  // For now, don't print implicit 'then' statements, since they can be
+  // present when the feature is disabled.
+  // TODO: Once we enable the feature, we can remove this.
+  if (!stmt->isImplicit())
+    Printer.printKeyword("then", Options, " ");
+
+  visit(stmt->getResult());
+}
+
 void PrintAST::visitThrowStmt(ThrowStmt *stmt) {
   Printer << tok::kw_throw << " ";
   visit(stmt->getSubExpr());

lib/AST/ASTScopeCreation.cpp

@@ -437,6 +437,12 @@ class NodeAdder
     return p;
   }
 
+  ASTScopeImpl *visitThenStmt(ThenStmt *ts, ASTScopeImpl *p,
+                              ScopeCreator &scopeCreator) {
+    visitExpr(ts->getResult(), p, scopeCreator);
+    return p;
+  }
+
   ASTScopeImpl *visitDeferStmt(DeferStmt *ds, ASTScopeImpl *p,
                                ScopeCreator &scopeCreator) {
     visitFuncDecl(ds->getTempDecl(), p, scopeCreator);

lib/AST/ASTVerifier.cpp

@@ -1214,7 +1214,7 @@ class Verifier : public ASTWalker {
     void verifyChecked(SingleValueStmtExpr *E) {
       using Kind = IsSingleValueStmtResult::Kind;
       switch (E->getStmt()->mayProduceSingleValue(Ctx).getKind()) {
-      case Kind::NoExpressionBranches:
+      case Kind::NoResult:
         // These are allowed as long as the type is Void.
         checkSameType(
             E->getType(), Ctx.getVoidType(),