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 an
experimental feature option pending evolution
discussion.

Author: hamishknight

CodeGeneration/Sources/SyntaxSupport/KeywordSpec.swift

@@ -235,6 +235,7 @@ public let KEYWORDS: [KeywordSpec] = [
   KeywordSpec("swift"),
   KeywordSpec("switch", isLexerClassified: true, requiresTrailingSpace: true),
   KeywordSpec("target"),
+  KeywordSpec("then", isExperimental: true),
   KeywordSpec("throw", isLexerClassified: true, requiresTrailingSpace: true),
   KeywordSpec("throws", isLexerClassified: true, requiresTrailingSpace: true),
   KeywordSpec("transpose"),

CodeGeneration/Sources/SyntaxSupport/StmtNodes.swift

@@ -615,4 +615,29 @@ public let STMT_NODES: [Node] = [
     ]
   ),
 
+  // then-stmt -> 'then' expr ';'?
+  Node(
+    kind: .thenStmt,
+    base: .stmt,
+    isExperimental: true,
+    nameForDiagnostics: "'then' statement",
+    documentation: """
+      A statement used to indicate the produced value from an if/switch
+      expression. Written as:
+
+      ```swift
+      then <expr>
+      ```
+      """,
+    children: [
+      Child(
+        name: "ThenKeyword",
+        kind: .token(choices: [.keyword(text: "then")])
+      ),
+      Child(
+        name: "Expression",
+        kind: .node(kind: .expr)
+      ),
+    ]
+  ),
 ]

CodeGeneration/Sources/SyntaxSupport/SyntaxNodeKind.swift

@@ -266,6 +266,7 @@ public enum SyntaxNodeKind: String, CaseIterable {
   case switchDefaultLabel
   case switchExpr
   case ternaryExpr
+  case thenStmt
   case throwStmt
   case tryExpr
   case tupleExpr

Sources/SwiftParser/ExperimentalFeatures.swift

@@ -19,3 +19,8 @@ extension Parser {
     }
   }
 }
+
+extension Parser.ExperimentalFeatures {
+  /// Whether to enable the parsing of 'then' statements.
+  public static let thenStatements = Self(rawValue: 1 << 0)
+}

Sources/SwiftParser/Expressions.swift

@@ -24,7 +24,9 @@ extension TokenConsumer {
       if backtrack.at(anyIn: IfOrSwitch.self) != nil {
         return true
       }
-      if backtrack.atStartOfDeclaration() || backtrack.atStartOfStatement() {
+      // Note we current pass `preferExpr: false` to prefer diagnosing `try then`
+      // as needing to be `then try`, rather than parsing `then` as an expression.
+      if backtrack.atStartOfDeclaration() || backtrack.atStartOfStatement(preferExpr: false) {
         // If after the 'try' we are at a declaration or statement, it can't be a valid expression.
         // Decide how we want to consume the 'try':
         // If the declaration or statement starts at a new line, the user probably just forgot to write the expression after 'try' -> parse it as a TryExpr
@@ -1552,7 +1554,9 @@ extension Parser {
         // If the next token is at the beginning of a new line and can never start
         // an element, break.
         if self.atStartOfLine
-          && (self.at(.rightBrace, .poundEndif) || self.atStartOfDeclaration() || self.atStartOfStatement())
+          && (self.at(.rightBrace, .poundEndif)
+            || self.atStartOfDeclaration()
+            || self.atStartOfStatement(preferExpr: false))
         {
           break
         }
@@ -2168,7 +2172,10 @@ extension Parser {
             )
           )
         )
-      } else if allowStandaloneStmtRecovery && (self.atStartOfExpression() || self.atStartOfStatement() || self.atStartOfDeclaration()) {
+      } else if allowStandaloneStmtRecovery
+        && (self.atStartOfExpression() || self.atStartOfStatement(preferExpr: false)
+          || self.atStartOfDeclaration())
+      {
         // Synthesize a label for the statement or declaration that isn't covered by a case right now.
         let statements = parseSwitchCaseBody()
         if statements.isEmpty {

Sources/SwiftParser/Statements.swift

@@ -16,16 +16,22 @@ extension TokenConsumer {
   /// Returns `true` if the current token represents the start of a statement
   /// item.
   ///
+  /// - Parameters:
+  ///   - allowRecovery: Whether to attempt to perform recovery.
+  ///   - preferExpr: If either an expression or statement could be
+  ///     parsed and this parameter is `true`, the function returns `false`
+  ///     such that an expression can be parsed.
+  ///
   /// - Note: This function must be kept in sync with `parseStatement()`.
   /// - Seealso: ``Parser/parseStatement()``
-  func atStartOfStatement(allowRecovery: Bool = false) -> Bool {
+  func atStartOfStatement(allowRecovery: Bool = false, preferExpr: Bool) -> Bool {
     var lookahead = self.lookahead()
     if allowRecovery {
       // Attributes are not allowed on statements. But for recovery, skip over
       // misplaced attributes.
       _ = lookahead.consumeAttributeList()
     }
-    return lookahead.atStartOfStatement(allowRecovery: allowRecovery)
+    return lookahead.atStartOfStatement(allowRecovery: allowRecovery, preferExpr: preferExpr)
   }
 }
 
@@ -105,6 +111,8 @@ extension Parser {
       return label(self.parseDoStatement(doHandle: handle), with: optLabel)
     case (.yield, let handle)?:
       return label(self.parseYieldStatement(yieldHandle: handle), with: optLabel)
+    case (.then, let handle)?:
+      return label(self.parseThenStatement(handle: handle), with: optLabel)
     case nil:
       let missingStmt = RawStmtSyntax(RawMissingStmtSyntax(arena: self.arena))
       return label(missingStmt, with: optLabel)
@@ -630,7 +638,7 @@ extension Parser {
     if self.at(anyIn: IfOrSwitch.self) != nil {
       return true
     }
-    if self.atStartOfStatement() || self.atStartOfDeclaration() {
+    if self.atStartOfStatement(preferExpr: true) || self.atStartOfDeclaration() {
       return false
     }
     return true
@@ -723,6 +731,37 @@ extension Parser {
   }
 }
 
+extension Parser {
+  /// Parse a `then` statement.
+  mutating func parseThenStatement(handle: RecoveryConsumptionHandle) -> RawStmtSyntax {
+    guard experimentalFeatures.contains(.thenStatements) else {
+      return RawStmtSyntax(RawMissingStmtSyntax(arena: self.arena))
+    }
+    let (unexpectedBeforeThen, then) = self.eat(handle)
+    let hasMisplacedTry = unexpectedBeforeThen?.containsToken(where: { TokenSpec(.try) ~= $0 }) ?? false
+
+    var expr = self.parseExpression()
+    if hasMisplacedTry && !expr.is(RawTryExprSyntax.self) {
+      expr = RawExprSyntax(
+        RawTryExprSyntax(
+          tryKeyword: missingToken(.try),
+          questionOrExclamationMark: nil,
+          expression: expr,
+          arena: self.arena
+        )
+      )
+    }
+    return RawStmtSyntax(
+      RawThenStmtSyntax(
+        unexpectedBeforeThen,
+        thenKeyword: then,
+        expression: expr,
+        arena: self.arena
+      )
+    )
+  }
+}
+
 extension Parser {
   struct StatementLabel {
     var label: RawTokenSyntax
@@ -791,7 +830,7 @@ extension Parser {
     }
 
     guard
-      self.at(.identifier) && !self.atStartOfStatement() && !self.atStartOfDeclaration()
+      self.at(.identifier) && !self.atStartOfStatement(preferExpr: true) && !self.atStartOfDeclaration()
     else {
       return nil
     }
@@ -806,9 +845,15 @@ extension Parser.Lookahead {
   /// Returns `true` if the current token represents the start of a statement
   /// item.
   ///
+  /// - Parameters:
+  ///   - allowRecovery: Whether to attempt to perform recovery.
+  ///   - preferExpr: If either an expression or statement could be
+  ///     parsed and this parameter is `true`, the function returns `false`
+  ///     such that an expression can be parsed.
+  ///
   /// - Note: This function must be kept in sync with `parseStatement()`.
   /// - Seealso: ``Parser/parseStatement()``
-  mutating func atStartOfStatement(allowRecovery: Bool = false) -> Bool {
+  mutating func atStartOfStatement(allowRecovery: Bool = false, preferExpr: Bool) -> Bool {
     if (self.at(anyIn: SwitchCaseStart.self) != nil || self.at(.atSign)) && withLookahead({ $0.atStartOfSwitchCaseItem() }) {
       // We consider SwitchCaseItems statements so we don't parse the start of a new case item as trailing parts of an expression.
       return true
@@ -877,11 +922,61 @@ extension Parser.Lookahead {
         // For example, could be the function call "discard()".
         return false
       }
+
+    case .then:
+      return atStartOfThenStatement(preferExpr: preferExpr)
+
     case nil:
       return false
     }
   }
 
+  /// Whether we're currently at a `then` token that should be parsed as a
+  /// `then` statement.
+  mutating func atStartOfThenStatement(preferExpr: Bool) -> Bool {
+    // If the feature is disabled, don't parse it.
+    guard self.experimentalFeatures.contains(.thenStatements) else {
+      return false
+    }
+    guard self.at(.keyword(.then)) else {
+      return false
+    }
+
+    // If we prefer an expr and aren't at the start of a newline, then don't
+    // parse a ThenStmt.
+    if preferExpr && !self.atStartOfLine {
+      return false
+    }
+
+    let next = peek()
+
+    // If 'then' is followed by a binary or postfix operator, prefer to parse as
+    // an expr.
+    if BinaryOperatorLike(lexeme: next) != nil || PostfixOperatorLike(lexeme: next) != nil {
+      return false
+    }
+
+    switch PrepareForKeywordMatch(next) {
+    case TokenSpec(.is), TokenSpec(.as):
+      // Treat 'is' and 'as' like the binary operator case, and parse as an
+      // expr.
+      return false
+
+    case .leftBrace:
+      // This is a trailing closure.
+      return false
+
+    case .leftParen, .leftSquare, .period:
+      // These are handled based on whether there is trivia between the 'then'
+      // and the token. If so, it's a 'then' statement. Otherwise it should
+      // be treated as an expression, e.g `then(...)`, `then[...]`, `then.foo`.
+      return !self.currentToken.trailingTriviaText.isEmpty || !next.leadingTriviaText.isEmpty
+    default:
+      break
+    }
+    return true
+  }
+
   /// Returns whether the parser's current position is the start of a switch case,
   /// given that we're in the middle of a switch already.
   mutating func atStartOfSwitchCase(allowRecovery: Bool = false) -> Bool {

Sources/SwiftParser/TokenPrecedence.swift

@@ -212,7 +212,7 @@ enum TokenPrecedence: Comparable {
       // Secondary parts of control-flow constructs
       .case, .catch, .default, .else,
       // Return-like statements
-      .break, .continue, .fallthrough, .return, .throw, .yield:
+      .break, .continue, .fallthrough, .return, .throw, .then, .yield:
       self = .stmtKeyword
     // MARK: Decl keywords
     case  // Types