xxx

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,8 @@ 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(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.
@@ -6154,10 +6155,13 @@ 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 stmts from each branch of the SingleValueStmtExpr.
+  ArrayRef<ThenStmt *>
+  getResultStmts(SmallVectorImpl<ThenStmt *> &scratch) const;
+
+  /// Retrieve the result exprs 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

@@ -665,6 +665,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

@@ -3859,8 +3859,8 @@ class IsSingleValueStmtResult {
     /// The statement is an 'if' statement without an unconditional 'else'.
     NonExhaustiveIf,
 
-    /// There are no single-expression branches.
-    NoExpressionBranches,
+    /// There are no branch that produces a resulting value.
+    NoResult,
 
     /// There is an unhandled statement branch. This should only be the case
     /// for invalid AST.
@@ -3915,8 +3915,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/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();
+
   /// `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.
@@ -1933,6 +1937,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/Sema/ConstraintLocator.h

@@ -931,19 +931,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

@@ -5496,6 +5496,11 @@ void PrintAST::visitYieldStmt(YieldStmt *stmt) {
   if (parens) Printer << ")";
 }
 
+void PrintAST::visitThenStmt(ThenStmt *stmt) {
+  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(),

lib/AST/ASTWalker.cpp

@@ -1758,6 +1758,15 @@ Stmt *Traversal::visitYieldStmt(YieldStmt *YS) {
   return YS;
 }
 
+Stmt *Traversal::visitThenStmt(ThenStmt *TS) {
+  auto *E = doIt(TS->getResult());
+  if (!E)
+    return nullptr;
+
+  TS->setResult(E);
+  return TS;
+}
+
 Stmt *Traversal::visitDeferStmt(DeferStmt *DS) {
   if (doIt(DS->getTempDecl()))
     return nullptr;

lib/AST/Expr.cpp

@@ -2492,27 +2492,29 @@ SingleValueStmtExpr *SingleValueStmtExpr::createWithWrappedBranches(
     if (!BS)
       continue;
 
-    auto *S = BS->getSingleActiveStatement();
-    if (!S)
-      continue;
-
-    if (mustBeExpr) {
-      // If this must be an expression, we can eagerly wrap any exhaustive if
-      // and switch branch.
-      if (auto *IS = dyn_cast<IfStmt>(S)) {
-        if (!IS->isSyntacticallyExhaustive())
+    if (auto *S = BS->getSingleActiveStatement()) {
+      if (mustBeExpr) {
+        // If this must be an expression, we can eagerly wrap any exhaustive if
+        // and switch branch.
+        if (auto *IS = dyn_cast<IfStmt>(S)) {
+          if (!IS->isSyntacticallyExhaustive())
+            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;
-      } 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));
     }
-    BS->setLastElement(
-        SingleValueStmtExpr::createWithWrappedBranches(ctx, S, DC, mustBeExpr));
+
+    // Wrap single expression branches in an implicit 'then <expr>'.
+    if (auto *E = BS->getSingleActiveExpression())
+      BS->setLastElement(ThenStmt::createImplicit(ctx, E));
   }
   return SVE;
 }
@@ -2571,20 +2573,31 @@ SingleValueStmtExpr::getBranches(SmallVectorImpl<Stmt *> &scratch) const {
   llvm_unreachable("Unhandled case in switch!");
 }
 
-ArrayRef<Expr *> SingleValueStmtExpr::getSingleExprBranches(
-    SmallVectorImpl<Expr *> &scratch) const {
+ArrayRef<ThenStmt *> SingleValueStmtExpr::getResultStmts(
+    SmallVectorImpl<ThenStmt *> &scratch) const {
   assert(scratch.empty());
   SmallVector<Stmt *, 4> stmtScratch;
   for (auto *branch : getBranches(stmtScratch)) {
     auto *BS = dyn_cast<BraceStmt>(branch);
-    if (!BS)
+    if (!BS || BS->empty())
       continue;
-    if (auto *E = BS->getSingleActiveExpression())
-      scratch.push_back(E);
+    if (auto *S = BS->getLastElement().dyn_cast<Stmt *>()) {
+      if (auto *TS = dyn_cast<ThenStmt>(S))
+        scratch.push_back(TS);
+    }
   }
   return scratch;
 }
 
+ArrayRef<Expr *> SingleValueStmtExpr::getResultExprs(
+    SmallVectorImpl<Expr *> &scratch) const {
+  SmallVector<ThenStmt *, 4> stmtScratch;
+  for (auto *TS : getResultStmts(stmtScratch))
+    scratch.push_back(TS->getResult());
+
+  return scratch;
+}
+
 void InterpolatedStringLiteralExpr::forEachSegment(ASTContext &Ctx, 
     llvm::function_ref<void(bool, CallExpr *)> callback) {
   auto appendingExpr = getAppendingExpr();