swiftlang · hamishknight · Sep 1, 2023 · Sep 1, 2023 · Sep 1, 2023 · Sep 1, 2023
@@ -255,6 +255,7 @@ public enum Keyword: CaseIterable {
   case swift
   case `switch`
   case target
+  case then
   case `throw`
   case `throws`
   case transpose
@@ -644,6 +645,8 @@ public enum Keyword: CaseIterable {
       return KeywordSpec("switch", isLexerClassified: true)
     case .target:
       return KeywordSpec("target")
+    case .then:
+      return KeywordSpec("then")
     case .throw:
       return KeywordSpec("throw", isLexerClassified: true)
     case .throws:

@@ -613,4 +613,31 @@ public let STMT_NODES: [Node] = [
     ]
   ),
 
+  // then-stmt -> 'then' expr ';'?
+  Node(
+    kind: .thenStmt,
+    base: .stmt,
+    // FIXME: This should be marked experimental.
+    isExperimental: false,
+    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(.then)])
+      ),
+      Child(
+        name: "expression",
+        kind: .node(kind: .expr)
+      ),
+    ]
+  ),
 ]
@@ -267,6 +267,7 @@ public enum SyntaxNodeKind: String, CaseIterable {
   case switchDefaultLabel
   case switchExpr
   case ternaryExpr
+  case thenStmt
   case throwStmt
   case tryExpr
   case tupleExpr

@@ -22,3 +22,8 @@ extension Parser {
     public static let referenceBindings = Self(rawValue: 1 << 0)
   }
 }
+
+extension Parser.ExperimentalFeatures {
+  /// Whether to enable the parsing of 'then' statements.
+  public static let thenStatements = Self(rawValue: 1 << 0)
+}
@@ -24,7 +24,9 @@ extension TokenConsumer {
       if backtrack.at(anyIn: IfOrSwitch.self) != nil {
         return true
       }
-      if backtrack.atStartOfDeclaration() || backtrack.atStartOfStatement() {
+      // Note we currently 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
@@ -1545,7 +1547,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
         }
@@ -2166,7 +2170,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 {

@@ -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,7 +111,9 @@ extension Parser {
       return label(self.parseDoStatement(doHandle: handle), with: optLabel)
     case (.yield, let handle)?:
       return label(self.parseYieldStatement(yieldHandle: handle), with: optLabel)
-    case nil:
+    case (.then, let handle)? where experimentalFeatures.contains(.thenStatements):
+      return label(self.parseThenStatement(handle: handle), with: optLabel)
+    case nil, (.then, _)?:
       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,36 @@ extension Parser {
   }
 }
 
+extension Parser {
+  /// Parse a `then` statement.
+  mutating func parseThenStatement(handle: RecoveryConsumptionHandle) -> RawStmtSyntax {
+    assert(experimentalFeatures.contains(.thenStatements))
+
+    let (unexpectedBeforeThen, then) = self.eat(handle)
+    let hasMisplacedTry = unexpectedBeforeThen?.containsToken(where: { TokenSpec(.try) ~= $0 }) ?? false
+
+    var expr = self.parseExpression(flavor: .basic, pattern: .none)
+    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 +829,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 +844,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 +921,57 @@ extension Parser.Lookahead {
         // For example, could be the function call "discard()".
         return false
       }
-    case nil:
+
+    case .then where experimentalFeatures.contains(.thenStatements):
+      return atStartOfThenStatement(preferExpr: preferExpr)
+
+    case nil, .then:
       return false
     }
   }
 
+  /// Whether we're currently at a `then` token that should be parsed as a
+  /// `then` statement.
+  mutating func atStartOfThenStatement(preferExpr: Bool) -> Bool {
+    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 {

@@ -210,7 +210,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