Address review feedback

Clean up some code, allow `if` and `switch` to
count as the start of an expression, and add some
substructure test assertions.

Author: hamishknight

Sources/SwiftParser/Expressions.swift

@@ -505,19 +505,21 @@ extension Parser {
     // First check to see if we have the start of a regex literal `/.../`.
     //    tryLexRegexLiteral(/*forUnappliedOperator*/ false)
 
+    switch self.currentToken.tokenKind {
     // Try parse an 'if' or 'switch' as an expression. Note we do this here in
     // parseUnaryExpression as we don't allow postfix syntax to hang off such
     // expressions to avoid ambiguities such as postfix '.member', which can
     // currently be parsed as a static dot member for a result builder.
-    if self.at(.keyword(.switch)) {
+    case .switchKeyword:
       return RawExprSyntax(
-        parseSwitchExpression(switchHandle: .constant(.keyword(.switch)))
+        parseSwitchExpression(switchHandle: .constant(.switchKeyword))
       )
-    }
-    if self.at(.keyword(.if)) {
+    case .ifKeyword:
       return RawExprSyntax(
-        parseIfExpression(ifHandle: .constant(.keyword(.if)))
+        parseIfExpression(ifHandle: .constant(.ifKeyword))
       )
+    default:
+      break
     }
 
     switch self.at(anyIn: ExpressionPrefixOperator.self) {
@@ -2495,8 +2497,8 @@ extension Parser.Lookahead {
     }
 
     // If this is the start of a switch body, this isn't a trailing closure.
-    if self.peek().rawTokenKind == .keyword(.case) {
-      return false;
+    if self.peek().tokenKind == .caseKeyword {
+      return false
     }
 
     // If this is a normal expression (not an expr-basic) then trailing closures

Sources/SwiftParser/RawTokenKindSubset.swift

@@ -616,17 +616,17 @@ enum IfOrSwitch: RawTokenKindSubset {
   case switchKeyword
 
   init?(lexeme: Lexer.Lexeme) {
-    switch lexeme {
-    case RawTokenKindMatch(.keyword(.if)): self = .ifKeyword
-    case RawTokenKindMatch(.keyword(.switch)): self = .switchKeyword
+    switch lexeme.tokenKind {
+    case .ifKeyword: self = .ifKeyword
+    case .switchKeyword: self = .switchKeyword
     default: return nil
     }
   }
 
   var rawTokenKind: RawTokenKind {
     switch self {
-    case .ifKeyword: return .keyword(.if)
-    case .switchKeyword: return .keyword(.switch)
+    case .ifKeyword: return .ifKeyword
+    case .switchKeyword: return .switchKeyword
     }
   }
 }
@@ -795,6 +795,7 @@ enum ExpressionStart: RawTokenKindSubset {
   case expressionPrefixOperator(ExpressionPrefixOperator)
   case matchingPatternStart(MatchingPatternStart)
   case primaryExpressionStart(PrimaryExpressionStart)
+  case ifOrSwitch(IfOrSwitch)
 
   init?(lexeme: Lexer.Lexeme) {
     if let subset = AwaitTryMove(lexeme: lexeme) {
@@ -805,6 +806,8 @@ enum ExpressionStart: RawTokenKindSubset {
       self = .matchingPatternStart(subset)
     } else if let subset = PrimaryExpressionStart(lexeme: lexeme) {
       self = .primaryExpressionStart(subset)
+    } else if let subset = IfOrSwitch(lexeme: lexeme) {
+      self = .ifOrSwitch(subset)
     } else {
       return nil
     }
@@ -815,6 +818,7 @@ enum ExpressionStart: RawTokenKindSubset {
       + ExpressionPrefixOperator.allCases.map(Self.expressionPrefixOperator)
       + MatchingPatternStart.allCases.map(Self.matchingPatternStart)
       + PrimaryExpressionStart.allCases.map(Self.primaryExpressionStart)
+      + IfOrSwitch.allCases.map(Self.ifOrSwitch)
   }
 
   var rawTokenKind: RawTokenKind {
@@ -823,6 +827,7 @@ enum ExpressionStart: RawTokenKindSubset {
     case .expressionPrefixOperator(let underlyingKind): return underlyingKind.rawTokenKind
     case .matchingPatternStart(let underlyingKind): return underlyingKind.rawTokenKind
     case .primaryExpressionStart(let underlyingKind): return underlyingKind.rawTokenKind
+    case .ifOrSwitch(let underlyingKind): return underlyingKind.rawTokenKind
     }
   }
 
@@ -832,6 +837,7 @@ enum ExpressionStart: RawTokenKindSubset {
     case .expressionPrefixOperator(let underlyingKind): return underlyingKind.contextualKeyword
     case .matchingPatternStart(let underlyingKind): return underlyingKind.contextualKeyword
     case .primaryExpressionStart(let underlyingKind): return underlyingKind.contextualKeyword
+    case .ifOrSwitch(let underlyingKind): return underlyingKind.contextualKeyword
     }
   }
 }

Sources/SwiftParser/TopLevel.swift

@@ -202,11 +202,12 @@ extension Parser {
     // We could also achieve this by more eagerly attempting to parse an 'if'
     // or 'switch' as an expression when in statement position, but that
     // could result in less useful recovery behavior.
-    if at(.keyword(.as)), let stmtExpr = stmt.as(RawExpressionStmtSyntax.self) {
-      let expr = stmtExpr.expression
+    if at(.asKeyword),
+      let expr = stmt.as(RawExpressionStmtSyntax.self)?.expression
+    {
       if expr.is(RawIfExprSyntax.self) || expr.is(RawSwitchExprSyntax.self) {
         let (op, rhs) = parseUnresolvedAsExpr(
-          handle: .init(tokenKind: .keyword(.as))
+          handle: .init(tokenKind: .asKeyword)
         )
         let sequence = RawExprSyntax(
           RawSequenceExprSyntax(