Merge pull request #67454 from hamishknight/then

Author: hamishknight

include/swift/AST/DiagnosticsParse.def

@@ -1244,8 +1244,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", ())
@@ -1399,6 +1400,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

@@ -1231,14 +1231,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 *getThenStmtFrom(BraceStmt *BS);
+
+  /// Whether the given BraceStmt has a result to be produced from a parent
+  /// SingleValueStmtExpr.
+  static bool hasResult(BraceStmt *BS) {
+    return getThenStmtFrom(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 resulting ThenStmts from each branch of the
+  /// SingleValueStmtExpr.
+  ArrayRef<ThenStmt *>
+  getThenStmts(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/TokenKinds.def

@@ -258,6 +258,7 @@ MISC(string_segment)
 MISC(string_interpolation_anchor)
 MISC(kw_yield)
 MISC(kw_discard)
+MISC(kw_then)
 
 // The following tokens are irrelevant for swiftSyntax and thus only declared 
 // in this .def file

include/swift/AST/TypeCheckRequests.h

@@ -3901,8 +3901,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.
@@ -3957,8 +3957,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

@@ -226,6 +226,9 @@ EXPERIMENTAL_FEATURE(GlobalConcurrency, 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,