[Parser] Add 'atContextualKeywordPrefixedSyntax' method

Unified disambiguation method for contextual-keyword prefixed syntax.

Author: rintaro

Sources/SwiftParser/Expressions.swift

@@ -381,20 +381,6 @@ extension Parser {
     }
   }
 
-  /// Whether the current token is a valid contextual exprssion modifier like
-  /// `copy`, `consume`.
-  ///
-  /// `copy` etc. are only contextually a keyword if they are followed by an
-  /// identifier or keyword on the same line. We do this to ensure that we do
-  /// not break any copy functions defined by users.
-  private mutating func atContextualExpressionModifier() -> Bool {
-    return self.peek(
-      isAt: TokenSpec(.identifier, allowAtStartOfLine: false),
-      TokenSpec(.dollarIdentifier, allowAtStartOfLine: false),
-      TokenSpec(.self, allowAtStartOfLine: false)
-    )
-  }
-
   /// Parse an expression sequence element.
   mutating func parseSequenceExpressionElement(
     flavor: ExprFlavor,
@@ -445,27 +431,7 @@ extension Parser {
         )
       )
     case (.unsafe, let handle)?:
-      if self.peek().isAtStartOfLine
-        // Closing paired syntax
-        || self.peek(isAt: .rightParen, .rightSquare, .rightBrace)
-        // Assignment
-        || self.peek(isAt: .equal)
-        // As an argument label or in a list context.
-        || self.peek(isAt: .colon, .comma)
-        // Start of a closure in a context where it should be interpreted as
-        // being part of a statement.
-        || (flavor == .stmtCondition && self.peek(isAt: .leftBrace))
-        // Avoid treating as an "unsafe" expression when there is no trivia
-        // following the "unsafe" and the following token could either be a
-        // postfix expression or a subexpression:
-        //   - Member access vs. leading .
-        //   - Call vs. tuple expression.
-        //   - Subscript vs. array or dictionary expression
-        || (self.peek(isAt: .period, .leftParen, .leftSquare) && self.peek().leadingTriviaByteLength == 0
-          && self.currentToken.trailingTriviaByteLength == 0)
-        // End of file
-        || self.peek(isAt: .endOfFile)
-      {
+      if !atContextualKeywordPrefixedSyntax(exprFlavor: flavor, acceptClosure: true, preferPostfixExpr: false) {
         break EXPR_PREFIX
       }
 
@@ -486,7 +452,7 @@ extension Parser {
       assert(self.experimentalFeatures.contains(.oldOwnershipOperatorSpellings))
       fallthrough
     case (.borrow, let handle)?:
-      if !atContextualExpressionModifier() {
+      if !atContextualKeywordPrefixedSyntax(exprFlavor: flavor) {
         break EXPR_PREFIX
       }
       let borrowTok = self.eat(handle)
@@ -503,7 +469,7 @@ extension Parser {
       )
 
     case (.copy, let handle)?:
-      if !atContextualExpressionModifier() {
+      if !atContextualKeywordPrefixedSyntax(exprFlavor: flavor) {
         break EXPR_PREFIX
       }
 
@@ -524,7 +490,7 @@ extension Parser {
       assert(self.experimentalFeatures.contains(.oldOwnershipOperatorSpellings))
       fallthrough
     case (.consume, let handle)?:
-      if !atContextualExpressionModifier() {
+      if !atContextualKeywordPrefixedSyntax(exprFlavor: flavor) {
         break EXPR_PREFIX
       }
 
@@ -546,7 +512,7 @@ extension Parser {
       return RawExprSyntax(parsePackExpansionExpr(repeatHandle: handle, flavor: flavor, pattern: pattern))
 
     case (.each, let handle)?:
-      if !atContextualExpressionModifier() {
+      if !atContextualKeywordPrefixedSyntax(exprFlavor: flavor) {
         break EXPR_PREFIX
       }
 
@@ -561,7 +527,7 @@ extension Parser {
       )
 
     case (.any, _)?:
-      if !atContextualExpressionModifier() && !self.peek().isContextualPunctuator("~") {
+      if !atContextualKeywordPrefixedSyntax(exprFlavor: flavor) && !self.peek().isContextualPunctuator("~") {
         break EXPR_PREFIX
       }
 

Sources/SwiftParser/Statements.swift

@@ -900,6 +900,119 @@ extension Parser {
   }
 }
 
+extension TokenConsumer {
+  /// Disambiguate the word at the cursor looks like a keyword-prefixed syntax.
+  ///
+  /// - Parameters:
+  ///   - exprFlavor: The expression context. When using this function for a statement, e.g. 'yield',
+  ///     use `.basic`.
+  ///   - acceptClosure: When the next token is '{' and it looks like a closure, use this value as the result.
+  ///   - preferPostfixExpr: When the next token is '.', '(', or '[' and there is a space between the word,
+  ///     use `!preferPostfixExpr` as the result.
+  mutating func atContextualKeywordPrefixedSyntax(
+    exprFlavor: Parser.ExprFlavor,
+    acceptClosure: Bool = false,
+    preferPostfixExpr: Bool = true
+  ) -> Bool {
+    let next = peek()
+
+    // The next token must be at the same line.
+    if next.isAtStartOfLine {
+      return false
+    }
+
+    switch next.rawTokenKind {
+
+    case .identifier, .dollarIdentifier, .wildcard:
+      // E.g. <word> foo
+      return true
+
+    case .integerLiteral, .floatLiteral,
+      .stringQuote, .multilineStringQuote, .singleQuote, .rawStringPoundDelimiter,
+      .regexSlash, .regexPoundDelimiter:
+      // E.g. <word> 1
+      return true
+
+    case .prefixAmpersand, .prefixOperator, .atSign, .backslash, .pound:
+      // E.g. <word> !<expr>
+      return true
+
+    case .keyword:
+      switch Keyword(next.tokenText) {
+      case .as, .is, .in:
+        // E.g. <word> is <expr>
+        return false
+      default:
+        // Other lexer-classified keywords are identifier-like.
+        // E.g. <word> self
+        return true
+      }
+
+    case .binaryOperator, .equal, .arrow, .infixQuestionMark:
+      // E.g. <word> != <expr>
+      return false
+    case .postfixOperator, .postfixQuestionMark, .exclamationMark, .ellipsis:
+      // E.g. <word>++
+      return false
+    case .rightBrace, .rightParen, .rightSquare:
+      // E.g. <word>]
+      return false
+    case .colon, .comma:
+      // E.g. <word>,
+      return false
+    case .semicolon, .endOfFile, .poundElse, .poundElseif, .poundEndif:
+      return false
+
+    case .leftAngle, .rightAngle:
+      // Lexer never produce these token kinds.
+      return false
+
+    case .stringSegment, .regexLiteralPattern:
+      // Calling this function inside a string/regex literal?
+      return false
+
+    case .backtick, .poundAvailable, .poundUnavailable,
+      .poundSourceLocation, .poundIf, .shebang, .unknown:
+      // These are invalid for both cases
+      // E.g. <word> #available
+      return false
+
+    case .period, .leftParen, .leftSquare:
+      // These are truly ambiguous. They can be both start of postfix expression
+      // suffix or start of primary expression:
+      //
+      //   - Member access vs. implicit member expression
+      //   - Call vs. tuple expression
+      //   - Subscript vs. collection literal
+      //
+      let hasSpace = (next.leadingTriviaByteLength + currentToken.trailingTriviaByteLength) != 0
+      if !hasSpace {
+        // No space, the word is an decl-ref expression
+        return false
+      }
+      return !preferPostfixExpr
+
+    case .leftBrace:
+      // E.g. <word> { ... }
+      // Trailing closure is also ambiguous:
+      //
+      //   - Trailing closure vs. immediately-invoked closure
+      //
+      // Checking whitespace between the word cannot help this because people
+      // usually put a space before trailing closures. Even though that is source
+      // breaking, we prefer parsing it as a keyword if the syntax accepts
+      // immediately-invoked closure patterns. E.g. 'unsafe { ... }()'
+      if !acceptClosure {
+        return false
+      }
+      return self.withLookahead {
+        $0.consumeAnyToken()
+        return $0.atValidTrailingClosure(flavor: exprFlavor)
+      }
+    }
+  }
+}
+
 // MARK: Lookahead
 
 extension Parser.Lookahead {
@@ -949,91 +1062,16 @@ extension Parser.Lookahead {
       // FIXME: 'repeat' followed by '{' could be a pack expansion
       // with a closure pattern.
       return self.peek().rawTokenKind == .leftBrace
-    case .yield?:
-      switch self.peek().rawTokenKind {
-      case .prefixAmpersand:
-        // "yield &" always denotes a yield statement.
-        return true
-      case .leftParen:
-        // "yield (", by contrast, must be disambiguated with additional
-        // context. We always consider it an apply expression of a function
-        // called `yield` for the purposes of the parse.
-        return false
-      case .binaryOperator:
-        // 'yield &= x' treats yield as an identifier.
-        return false
-      default:
-        // "yield" followed immediately by any other token is likely a
-        // yield statement of some singular expression.
-        return !self.peek().isAtStartOfLine
-      }
-    case .discard?:
-      let next = peek()
-      // The thing to be discarded must be on the same line as `discard`.
-      if next.isAtStartOfLine {
-        return false
-      }
-      switch next.rawTokenKind {
-      case .identifier, .keyword:
-        // Since some identifiers like "self" are classified as keywords,
-        // we want to recognize those too, to handle "discard self". We also
-        // accept any identifier since we want to emit a nice error message
-        // later on during type checking.
-        return true
-      default:
-        // any other token following "discard" means it's not the statement.
-        // For example, could be the function call "discard()".
-        return false
-      }
-
-    case .then:
-      return atStartOfThenStatement(preferExpr: preferExpr)
+    case .yield?, .discard?:
+      return atContextualKeywordPrefixedSyntax(exprFlavor: .basic, preferPostfixExpr: true)
+    case .then?:
+      return atContextualKeywordPrefixedSyntax(exprFlavor: .basic, preferPostfixExpr: false)
 
     case nil:
       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
-    }
-
-    // If 'then' is followed by a binary or postfix operator, prefer to parse as
-    // an expr.
-    if peek(isAtAnyIn: BinaryOperatorLike.self) != nil || peek(isAtAnyIn: PostfixOperatorLike.self) != nil {
-      return false
-    }
-
-    switch PrepareForKeywordMatch(peek()) {
-    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 || !peek().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 {

Tests/SwiftParserTest/ThenStatementTests.swift

@@ -297,13 +297,7 @@ final class ThenStatementTests: ParserTestCase {
       """
       then1️⃣
       """,
-      diagnostics: [
-        DiagnosticSpec(
-          message: "expected expression in 'then' statement",
-          fixIts: ["insert expression"]
-        )
-      ],
-      fixedSource: "then <#expression#>"
+      substructure: DeclReferenceExprSyntax(baseName: .identifier("then"))
     )
   }
 
@@ -312,13 +306,7 @@ final class ThenStatementTests: ParserTestCase {
       """
       then1️⃣;
       """,
-      diagnostics: [
-        DiagnosticSpec(
-          message: "expected expression in 'then' statement",
-          fixIts: ["insert expression"]
-        )
-      ],
-      fixedSource: "then <#expression#>;"
+      substructure: DeclReferenceExprSyntax(baseName: .identifier("then"))
     )
   }
 
@@ -342,7 +330,7 @@ final class ThenStatementTests: ParserTestCase {
       then
       0
       """,
-      substructure: ThenStmtSyntax(expression: IntegerLiteralExprSyntax(0))
+      substructure: DeclReferenceExprSyntax(baseName: .identifier("then"))
     )
   }
 
@@ -685,7 +673,10 @@ final class ThenStatementTests: ParserTestCase {
       then
         .foo
       """,
-      substructure: ThenStmtSyntax(expression: MemberAccessExprSyntax(name: .identifier("foo")))
+      substructure: MemberAccessExprSyntax(
+        base: DeclReferenceExprSyntax(baseName: .identifier("then")),
+        name: .identifier("foo")
+      )
     )
   }