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

@@ -520,26 +520,25 @@ extension Parser {
     // First check to see if we have the start of a regex literal `/.../`.
     //    tryLexRegexLiteral(/*forUnappliedOperator*/ false)
 
-    // 'repeat' is the start of a pack expansion expression.
-    if (self.at(.keyword(.repeat))) {
-      return RawExprSyntax(
-        parsePackExpansionExpr(flavor, pattern: pattern)
-      )
-    }
+    switch self.currentToken.rawTokenKind {
+    case .keyword(.repeat):
+      // 'repeat' is the start of a pack expansion expression.
+      return RawExprSyntax(parsePackExpansionExpr(flavor, pattern: pattern))
 
     // 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 .keyword(.switch):
       return RawExprSyntax(
         parseSwitchExpression(switchHandle: .constant(.keyword(.switch)))
       )
-    }
-    if self.at(.keyword(.if)) {
+    case .keyword(.if):
       return RawExprSyntax(
         parseIfExpression(ifHandle: .constant(.keyword(.if)))
       )
+    default:
+      break
     }
 
     switch self.at(anyIn: ExpressionPrefixOperator.self) {
@@ -2109,7 +2108,7 @@ 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;
+      return false
     }
 
     // If this is a normal expression (not an expr-basic) then trailing closures

Sources/SwiftParser/RawTokenKindSubset.swift

@@ -618,8 +618,8 @@ enum IfOrSwitch: RawTokenKindSubset {
 
   init?(lexeme: Lexer.Lexeme) {
     switch lexeme {
-    case RawTokenKindMatch(.keyword(.if)): self = .ifKeyword
-    case RawTokenKindMatch(.keyword(.switch)): self = .switchKeyword
+    case RawTokenKindMatch(.if): self = .ifKeyword
+    case RawTokenKindMatch(.switch): self = .switchKeyword
     default: return nil
     }
   }
@@ -797,6 +797,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) {
@@ -807,6 +808,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
     }
@@ -817,6 +820,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 {
@@ -825,6 +829,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
     }
   }
 }

Sources/SwiftParser/TopLevel.swift

@@ -197,8 +197,9 @@ 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(.keyword(.as)),
+      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))