From f48955149434fddce82a90df5148ac1f94b85a97 Mon Sep 17 00:00:00 2001
From: Rintaro Ishizaki <rishizaki@apple.com>
Date: Mon, 10 Jul 2023 15:45:12 -0700
Subject: [PATCH] [Parser] Accept 'self' after 'each'

Also move 'repeat', 'each', and 'any' expression parsing to
'parseSequenceExpressionElement'.

rdar://107450487
---
 Sources/SwiftParser/Expressions.swift  | 68 ++++++++++++++++----------
 Sources/SwiftParser/TokenSpecSet.swift |  9 ++++
 2 files changed, 51 insertions(+), 26 deletions(-)

diff --git a/Sources/SwiftParser/Expressions.swift b/Sources/SwiftParser/Expressions.swift
index c9eb8db505a..93a3dd2769c 100644
--- a/Sources/SwiftParser/Expressions.swift
+++ b/Sources/SwiftParser/Expressions.swift
@@ -525,6 +525,46 @@ extension Parser {
           arena: self.arena
         )
       )
+
+    case (.repeat, let handle)?:
+      // 'repeat' is the start of a pack expansion expression.
+      return RawExprSyntax(parsePackExpansionExpr(repeatHandle: handle, flavor, pattern: pattern))
+
+    case (.each, let handle)?:
+      // `each` is only contextually a keyword, if it's followed by an
+      // identifier or 'self' on the same line. We do this to ensure that we do
+      // not break any 'each' functions defined by users. This is following with
+      // what we have done for the consume keyword.
+      switch self.peek() {
+      case TokenSpec(.identifier, allowAtStartOfLine: false),
+        TokenSpec(.dollarIdentifier, allowAtStartOfLine: false),
+        TokenSpec(.self, allowAtStartOfLine: false):
+        break
+      default:
+        // Break out of `outer switch` on failure.
+        break EXPR_PREFIX
+      }
+
+      let each = self.eat(handle)
+      let packRef = self.parseSequenceExpressionElement(flavor, pattern: pattern)
+      return RawExprSyntax(
+        RawPackElementExprSyntax(
+          eachKeyword: each,
+          packRefExpr: packRef,
+          arena: self.arena
+        )
+      )
+
+    case (.any, _)?:
+      // `any` is only contextually a keyword if it's followed by an identifier
+      // on the same line.
+      guard case TokenSpec(.identifier, allowAtStartOfLine: false) = self.peek() else {
+        break EXPR_PREFIX
+      }
+      // 'any' is parsed as a part of 'type'.
+      let type = self.parseType()
+      return RawExprSyntax(RawTypeExprSyntax(type: type, arena: self.arena))
+
     case nil:
       break
     }
@@ -549,10 +589,6 @@ extension Parser {
     //    tryLexRegexLiteral(/*forUnappliedOperator*/ false)
 
     switch self.currentToken {
-    case TokenSpec(.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
@@ -1234,27 +1270,6 @@ extension Parser {
         return RawExprSyntax(RawUnresolvedPatternExprSyntax(pattern: pattern, arena: self.arena))
       }
 
-      // We might have a contextual keyword followed by an identifier.
-      // 'each <identifier>' is a pack element expr, and 'any <identifier>'
-      // is an existential type expr.
-      if self.peek().rawTokenKind == .identifier, !self.peek().isAtStartOfLine {
-        if self.at(.keyword(.any)) {
-          let ty = self.parseType()
-          return RawExprSyntax(RawTypeExprSyntax(type: ty, arena: self.arena))
-        }
-
-        if let each = self.consume(if: .keyword(.each)) {
-          let packRef = self.parseSequenceExpressionElement(flavor, pattern: pattern)
-          return RawExprSyntax(
-            RawPackElementExprSyntax(
-              eachKeyword: each,
-              packRefExpr: packRef,
-              arena: self.arena
-            )
-          )
-        }
-      }
-
       return RawExprSyntax(self.parseIdentifierExpression())
     case (.Self, _)?:  // Self
       return RawExprSyntax(self.parseIdentifierExpression())
@@ -1493,10 +1508,11 @@ extension Parser {
   /// pack-expansion-expression → 'repeat' pattern-expression
   /// pattern-expression → expression
   mutating func parsePackExpansionExpr(
+    repeatHandle: TokenConsumptionHandle,
     _ flavor: ExprFlavor,
     pattern: PatternContext
   ) -> RawPackExpansionExprSyntax {
-    let repeatKeyword = self.consumeAnyToken()
+    let repeatKeyword = self.eat(repeatHandle)
     let patternExpr = self.parseExpression(flavor, pattern: pattern)
 
     return RawPackExpansionExprSyntax(
diff --git a/Sources/SwiftParser/TokenSpecSet.swift b/Sources/SwiftParser/TokenSpecSet.swift
index 5e70284fa57..2578f12aa8b 100644
--- a/Sources/SwiftParser/TokenSpecSet.swift
+++ b/Sources/SwiftParser/TokenSpecSet.swift
@@ -616,6 +616,9 @@ enum ExpressionModifierKeyword: TokenSpecSet {
   case `try`
   case consume
   case copy
+  case `repeat`
+  case each
+  case any
 
   init?(lexeme: Lexer.Lexeme) {
     switch PrepareForKeywordMatch(lexeme) {
@@ -625,6 +628,9 @@ enum ExpressionModifierKeyword: TokenSpecSet {
     case TokenSpec(.try): self = .try
     case TokenSpec(.consume): self = .consume
     case TokenSpec(.copy): self = .copy
+    case TokenSpec(.repeat): self = .repeat
+    case TokenSpec(.each): self = .each
+    case TokenSpec(.any): self = .any
     default: return nil
     }
   }
@@ -637,6 +643,9 @@ enum ExpressionModifierKeyword: TokenSpecSet {
     case .consume: return .keyword(.consume)
     case .copy: return .keyword(.copy)
     case .try: return .keyword(.try)
+    case .repeat: return .keyword(.repeat)
+    case .each: return .keyword(.each)
+    case .any: return .keyword(.any)
     }
   }
 }