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

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

CodeGeneration/Sources/SyntaxSupport/StmtNodes.swift

@@ -615,4 +615,19 @@ public let STMT_NODES: [Node] = [
     ]
   ),
 
+  Node(
+    kind: .thenStmt,
+    base: .stmt,
+    nameForDiagnostics: "'then' statement",
+    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

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

Sources/SwiftParser/Statements.swift

@@ -16,16 +16,21 @@ 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 an expression or statement could be parsed, returns
+  ///   `false` if this parameter is set to `true`.
+  ///
   /// - 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 = false) -> Bool {
     var lookahead = self.lookahead()
     if allowRecovery {
       // Attributes are not allowed on statements. But for recovery, skip over
       // misplaced attributes.
       _ = lookahead.consumeAttributeList()
     }
-    return lookahead.isStartOfStatement(allowRecovery: allowRecovery)
+    return lookahead.isStartOfStatement(allowRecovery: allowRecovery, preferExpr: preferExpr)
   }
 }
 
@@ -105,6 +110,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 +637,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 +730,35 @@ 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 +827,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 +842,14 @@ 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 `true`, if either an expression or statement could be
+  ///     parsed, returns `false` such that an expression can be parsed.
+  ///
   /// - Note: This function must be kept in sync with `parseStatement()`.
   /// - Seealso: ``Parser/parseStatement()``
-  mutating func isStartOfStatement(allowRecovery: Bool = false) -> Bool {
+  mutating func isStartOfStatement(allowRecovery: Bool = false, preferExpr: Bool = false) -> 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 +918,66 @@ extension Parser.Lookahead {
         // For example, could be the function call "discard()".
         return false
       }
+
+    case .then:
+      return isStartOfThenStatement(preferExpr: preferExpr)
+
     case nil:
       return false
     }
   }
 
+  /// Whether we're currently at a `then` token that should be parsed as a
+  /// `then` statement.
+  mutating func isStartOfThenStatement(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 want to 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 next.rawTokenKind {
+    case .keyword:
+      // Treat 'is' and 'as' like the binary operator case, and parse as an
+      // expr.
+      switch Keyword(next.tokenText) {
+      case .as, .is:
+        return false
+      default:
+        break
+      }
+
+    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 isAtStartOfSwitchCase(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

Sources/SwiftParser/TokenSpecSet.swift

@@ -66,6 +66,7 @@ enum CanBeStatementStart: TokenSpecSet {
   case `repeat`
   case `return`
   case `switch`
+  case then
   case `throw`
   case `while`
   case yield
@@ -85,6 +86,7 @@ enum CanBeStatementStart: TokenSpecSet {
     case TokenSpec(.repeat): self = .repeat
     case TokenSpec(.return): self = .return
     case TokenSpec(.switch): self = .switch
+    case TokenSpec(.then): self = .then
     case TokenSpec(.throw): self = .throw
     case TokenSpec(.while): self = .while
     case TokenSpec(.yield): self = .yield
@@ -107,6 +109,7 @@ enum CanBeStatementStart: TokenSpecSet {
     case .repeat: return .keyword(.repeat)
     case .return: return .keyword(.return)
     case .switch: return .keyword(.switch)
+    case .then: return .keyword(.then)
     case .throw: return .keyword(.throw)
     case .while: return .keyword(.while)
     case .yield: return .keyword(.yield)
@@ -478,48 +481,88 @@ enum Operator: TokenSpecSet {
   }
 }
 
-/// Tokens that can be used in operator declarations
-enum OperatorLike: TokenSpecSet {
-  case `operator`(Operator)
-  case exclamationMark
+/// Tokens that are either binary operators, or can act like binary operators.
+enum BinaryOperatorLike: TokenSpecSet {
+  case binaryOperator
   case infixQuestionMark
-  case postfixQuestionMark
   case equal
   case arrow
 
   init?(lexeme: Lexer.Lexeme) {
-    if let op = Operator(lexeme: lexeme) {
-      self = .operator(op)
-      return
-    }
     switch lexeme.rawTokenKind {
-    case .exclamationMark: self = .exclamationMark
+    case .binaryOperator: self = .binaryOperator
     case .infixQuestionMark: self = .infixQuestionMark
-    case .postfixQuestionMark: self = .postfixQuestionMark
     case .equal: self = .equal
     case .arrow: self = .arrow
     default: return nil
     }
   }
 
-  static var allCases: [OperatorLike] {
-    return Operator.allCases.map(Self.operator) + [
-      .exclamationMark,
-      .infixQuestionMark,
-      .postfixQuestionMark,
-      .equal,
-      .arrow,
-    ]
+  var spec: TokenSpec {
+    switch self {
+    case .binaryOperator: return .binaryOperator
+    case .infixQuestionMark: return TokenSpec(.infixQuestionMark, remapping: .binaryOperator)
+    case .equal: return TokenSpec(.equal, remapping: .binaryOperator)
+    case .arrow: return TokenSpec(.arrow, remapping: .binaryOperator)
+    }
+  }
+}
+
+/// Tokens that are either postfix operators, or can act like postfix operators.
+enum PostfixOperatorLike: TokenSpecSet {
+  case postfixOperator
+  case exclamationMark
+  case postfixQuestionMark
+
+  init?(lexeme: Lexer.Lexeme) {
+    switch lexeme.rawTokenKind {
+    case .postfixOperator: self = .postfixOperator
+    case .exclamationMark: self = .exclamationMark
+    case .postfixQuestionMark: self = .postfixQuestionMark
+    default: return nil
+    }
   }
 
   var spec: TokenSpec {
     switch self {
-    case .operator(let op): return op.spec
+    case .postfixOperator: return .postfixOperator
     case .exclamationMark: return TokenSpec(.exclamationMark, remapping: .postfixOperator)
-    case .infixQuestionMark: return TokenSpec(.infixQuestionMark, remapping: .binaryOperator)
     case .postfixQuestionMark: return TokenSpec(.postfixQuestionMark, remapping: .postfixOperator)
-    case .equal: return TokenSpec(.equal, remapping: .binaryOperator)
-    case .arrow: return TokenSpec(.arrow, remapping: .binaryOperator)
+    }
+  }
+}
+
+/// Tokens that can be used in operator declarations
+enum OperatorLike: TokenSpecSet {
+  case `operator`(Operator)
+  case binaryOperatorLike(BinaryOperatorLike)
+  case postfixOperatorLike(PostfixOperatorLike)
+
+  init?(lexeme: Lexer.Lexeme) {
+    if let op = Operator(lexeme: lexeme) {
+      self = .operator(op)
+      return
+    }
+    if let binOp = BinaryOperatorLike(lexeme: lexeme) {
+      self = .binaryOperatorLike(binOp)
+      return
+    }
+    if let postfixOp = PostfixOperatorLike(lexeme: lexeme) {
+      self = .postfixOperatorLike(postfixOp)
+      return
+    }
+    return nil
+  }
+
+  static var allCases: [OperatorLike] {
+    Operator.allCases.map(Self.operator) + BinaryOperatorLike.allCases.map(Self.binaryOperatorLike) + PostfixOperatorLike.allCases.map(Self.postfixOperatorLike)
+  }
+
+  var spec: TokenSpec {
+    switch self {
+    case .operator(let op): return op.spec
+    case .binaryOperatorLike(let op): return op.spec
+    case .postfixOperatorLike(let op): return op.spec
     }
   }
 }

Sources/SwiftParserDiagnostics/ParseDiagnosticsGenerator.swift

@@ -1471,6 +1471,20 @@ public class ParseDiagnosticsGenerator: SyntaxAnyVisitor {
     return .visitChildren
   }
 
+  public override func visit(_ node: ThenStmtSyntax) -> SyntaxVisitorContinueKind {
+    if shouldSkip(node) {
+      return .skipChildren
+    }
+    exchangeTokens(
+      unexpected: node.unexpectedBeforeThenKeyword,
+      unexpectedTokenCondition: { $0.tokenKind == .keyword(.try) },
+      correctTokens: [node.expression.as(TryExprSyntax.self)?.tryKeyword],
+      message: { _ in .tryMustBePlacedOnThenExpr },
+      moveFixIt: { MoveTokensAfterFixIt(movedTokens: $0, after: .keyword(.then)) }
+    )
+    return .visitChildren
+  }
+
   public override func visit(_ node: SameTypeRequirementSyntax) -> SyntaxVisitorContinueKind {
     if shouldSkip(node) {
       return .skipChildren

Sources/SwiftParserDiagnostics/ParserDiagnosticMessages.swift

@@ -243,6 +243,9 @@ extension DiagnosticMessage where Self == StaticParserError {
   public static var tryMustBePlacedOnThrownExpr: Self {
     .init("'try' must be placed on the thrown expression")
   }
+  public static var tryMustBePlacedOnThenExpr: Self {
+    .init("'try' must be placed on the produced expression")
+  }
   public static var tryOnInitialValueExpression: Self {
     .init("'try' must be placed on the initial value expression")
   }