Support recovery in front of declarations

If the current token is not the start of a declaration, check if we can consume unexpected tokens to reach the declaration’s introducer.

Author: ahoppen

Sources/SwiftParser/Declarations.swift

@@ -13,7 +13,7 @@
 @_spi(RawSyntax) import SwiftSyntax
 
 extension TokenConsumer {
-  func atStartOfDeclaration() -> Bool {
+  func atStartOfDeclaration(isAtTopLevel: Bool = false, allowRecovery: Bool = false) -> Bool {
     if self.at(anyIn: PoundDeclarationStart.self) != nil {
       return true
     }
@@ -44,8 +44,17 @@ extension TokenConsumer {
       }
     }
 
-    switch subparser.at(anyIn: DeclarationStart.self) {
-    case (.caseKeyword, _)?, nil:
+    let declStartKeyword: DeclarationStart?
+    if allowRecovery {
+      declStartKeyword = subparser.canRecoverTo(
+        anyIn: DeclarationStart.self,
+        recoveryPrecedence: isAtTopLevel ? nil : .strongBracketedClose
+      )?.0
+    } else {
+      declStartKeyword = subparser.at(anyIn: DeclarationStart.self)?.0
+    }
+    switch declStartKeyword {
+    case .caseKeyword, nil:
       // When 'case' appears inside a function, it's probably a switch
       // case, not an enum case declaration.
       return false
@@ -113,7 +122,7 @@ extension Parser {
     let attrs = DeclAttributes(
       attributes: self.parseAttributeList(),
       modifiers: self.parseModifierList())
-    switch self.at(anyIn: DeclarationStart.self) {
+    switch self.canRecoverTo(anyIn: DeclarationStart.self) {
     case (.importKeyword, _)?:
       return RawDeclSyntax(self.parseImportDeclaration(attrs))
     case (.classKeyword, _)?:
@@ -1010,8 +1019,7 @@ extension Parser {
   ///     actor-member → declaration | compiler-control-statement
   @_spi(RawSyntax)
   public mutating func parseActorDeclaration(_ attrs: DeclAttributes) -> RawActorDeclSyntax {
-    assert(self.atContextualKeyword("actor"))
-    let actorKeyword = self.expectIdentifierWithoutRecovery()
+    let (unexpectedBeforeActorKeyword, actorKeyword) = self.expectContextualKeyword("actor", precedence: .declKeyword)
     let name = self.expectIdentifierWithoutRecovery()
 
     let generics: RawGenericParameterClauseSyntax?
@@ -1041,6 +1049,7 @@ extension Parser {
     return RawActorDeclSyntax(
       attributes: attrs.attributes,
       modifiers: attrs.modifiers,
+      unexpectedBeforeActorKeyword,
       actorKeyword: actorKeyword,
       identifier: name,
       genericParameterClause: generics,

Sources/SwiftParser/Parser.swift

@@ -203,6 +203,76 @@ extension Parser {
 
 }
 
+// MARK: Check if we can recover to a token
+
+extension Parser {
+  /// Checks if it can reach a token of the given `kind` by skipping unexpected
+  /// tokens that have lower ``TokenPrecedence`` than expected token.
+  @_spi(RawSyntax)
+  public func canRecoverTo(_ kind: RawTokenKind) -> Bool {
+    if self.at(kind) {
+      return true
+    }
+    var lookahead = self.lookahead()
+    return lookahead.canRecoverTo([kind])
+  }
+
+  /// Checks if it can reach a contextual keyword with the given `name` by
+  /// skipping unexpected tokens that have lower ``TokenPrecedence`` than
+  /// `precedence`.
+  @_spi(RawSyntax)
+  public func canRecoverToContextualKeyword(
+    _ name: SyntaxText,
+    precedence: TokenPrecedence = TokenPrecedence(.contextualKeyword)
+  ) -> Bool {
+    if self.atContextualKeyword(name) {
+      return true
+    }
+    var lookahead = self.lookahead()
+    return lookahead.canRecoverTo([], contextualKeywords: [name])
+  }
+
+
+  /// Checks if it can reach a token whose kind is in `kinds` by skipping
+  /// unexpected tokens that have lower ``TokenPrecedence`` than `precedence`.
+  @_spi(RawSyntax)
+  public func canRecoverTo(
+    any kinds: [RawTokenKind]
+  ) -> Bool {
+    if self.at(any: kinds) {
+      return true
+    }
+    var lookahead = self.lookahead()
+    return lookahead.canRecoverTo(kinds)
+  }
+
+
+  /// Checks if we can reach a token in `subset` by skipping tokens that have
+  /// a precedence that have a lower ``TokenPrecedence`` than the minimum
+  /// precedence of a token in that subset.
+  /// If so, return the token that we can recover to and a handle that can be
+  /// used to consume the unexpected tokens and the token we recovered to.
+  func canRecoverTo<Subset: RawTokenKindSubset>(
+    anyIn subset: Subset.Type
+  ) -> (Subset, RecoveryConsumptionHandle)? {
+    if let (kind, handle) = self.at(anyIn: subset) {
+      return (kind, RecoveryConsumptionHandle(unexpectedTokens: 0, tokenConsumptionHandle: handle))
+    }
+    var lookahead = self.lookahead()
+    return lookahead.canRecoverTo(anyIn: subset)
+  }
+
+  /// Eat a token that we know we are currently positioned at, based on `canRecoverTo(anyIn:)`.
+  mutating func eat(_ handle: RecoveryConsumptionHandle) -> (RawUnexpectedNodesSyntax, Token) {
+    var unexpectedTokens = [RawSyntax]()
+    for _ in 0..<handle.unexpectedTokens {
+      unexpectedTokens.append(RawSyntax(self.consumeAnyToken()))
+    }
+    let token = self.eat(handle.tokenConsumptionHandle)
+    return (RawUnexpectedNodesSyntax(elements: unexpectedTokens, arena: self.arena), token)
+  }
+}
+
 // MARK: Expecting Tokens with Recovery
 
 extension Parser {
@@ -229,7 +299,7 @@ extension Parser {
 
   /// Attempts to consume a token of the given kind.
   /// If it cannot be found, the parser tries
-  ///  1. To each unexpected tokens that have lower ``TokenPrecedence`` than the
+  ///  1. To eat unexpected tokens that have lower ``TokenPrecedence`` than the
   ///     expected token and see if the token occurs after that unexpected.
   ///  2. If the token couldn't be found after skipping unexpected, it synthesizes
   ///     a missing token of the requested kind.
@@ -246,24 +316,24 @@ extension Parser {
 
   /// Attempts to consume a contextual keyword whose text is `name`.
   /// If it cannot be found, the parser tries
-  ///  1. To each unexpected tokens that have lower ``TokenPrecedence`` than the
-  ///     expected token and see if the token occurs after that unexpected.
+  ///  1. To eat unexpected tokens that have lower ``TokenPrecedence`` than `precedence`.
   ///  2. If the token couldn't be found after skipping unexpected, it synthesizes
   ///     a missing token of the requested kind.
   @_spi(RawSyntax)
   public mutating func expectContextualKeyword(
-    _ name: SyntaxText
+    _ name: SyntaxText,
+    precedence: TokenPrecedence = TokenPrecedence(.contextualKeyword)
   ) -> (unexpected: RawUnexpectedNodesSyntax?, token: RawTokenSyntax) {
     return expectImpl(
       consume: { $0.consumeIfContextualKeyword(name) },
-      canRecoverTo: { $0.canRecoverTo([], contextualKeywords: [name]) },
+      canRecoverTo: { $0.canRecoverTo([], contextualKeywords: [name], recoveryPrecedence: precedence) },
       makeMissing: { $0.missingToken(.contextualKeyword, text: name) }
     )
   }
 
   /// Attempts to consume a token whose kind is in `kinds`.
   /// If it cannot be found, the parser tries
-  ///  1. To each unexpected tokens that have lower ``TokenPrecedence`` than the
+  ///  1. To eat unexpected tokens that have lower ``TokenPrecedence`` than the
   ///     lowest precedence of the expected token kinds and see if a token of
   ///     the requested kinds occurs after the unexpected.
   ///  2. If the token couldn't be found after skipping unexpected, it synthesizes

Sources/SwiftParser/RawTokenKindSubset.swift

@@ -24,6 +24,11 @@ protocol RawTokenKindSubset: CaseIterable {
   /// If not `nil`, the token's will be remapped to this kind when the handle is eaten.
   var remappedKind: RawTokenKind? { get }
 
+  /// The precedence of this token that determines which tokens can be skipped
+  /// trying to reach it. If this returns `nil`, the precedence of `rawTokenKind`
+  /// is used. This is mostly overwritten for contextual keywords.
+  var precedence: TokenPrecedence? { get }
+
   /// Allows more flexible rejection of further token kinds based on the token's
   /// contents. Useful to e.g. look for contextual keywords.
   func accepts(lexeme: Lexer.Lexeme) -> Bool
@@ -42,6 +47,10 @@ extension RawTokenKindSubset {
     }
   }
 
+  var precedence: TokenPrecedence? {
+    return nil
+  }
+
   func accepts(lexeme: Lexer.Lexeme) -> Bool {
     return true
   }
@@ -392,6 +401,14 @@ enum DeclarationStart: RawTokenKindSubset {
     default: return nil
     }
   }
+
+  var precedence: TokenPrecedence? {
+    switch self {
+    case .actorContextualKeyword: return .declKeyword
+    case .caseKeyword: return .declKeyword
+    default: return nil
+    }
+  }
 }
 
 enum EffectsSpecifier: RawTokenKindSubset {

Sources/SwiftParser/Recovery.swift

@@ -14,6 +14,15 @@
 
 // MARK: Lookahead
 
+
+/// After calling `consume(ifAnyFrom:)` we know which token we are positioned
+/// at based on that function's return type. This handle allows consuming that
+/// token.
+struct RecoveryConsumptionHandle {
+  var unexpectedTokens: Int
+  var tokenConsumptionHandle: TokenConsumptionHandle
+}
+
 extension Parser.Lookahead {
   /// Tries eating tokens until it finds a token whose kind is in `kinds` or a
   /// contextual keyword with a text in `contextualKeywords` without skipping
@@ -25,13 +34,15 @@ extension Parser.Lookahead {
   /// `lookahead.tokensConsumed` as unexpected.
   mutating func canRecoverTo(
     _ kinds: [RawTokenKind],
-    contextualKeywords: [SyntaxText] = []
+    contextualKeywords: [SyntaxText] = [],
+    recoveryPrecedence: TokenPrecedence? = nil
   ) -> Bool {
-    assert(!kinds.isEmpty)
-    var recoveryPrecedence = kinds.map(TokenPrecedence.init).min()!
+    var precedences = kinds.map(TokenPrecedence.init)
     if !contextualKeywords.isEmpty {
-      recoveryPrecedence = min(recoveryPrecedence, TokenPrecedence(.identifier), TokenPrecedence(.contextualKeyword))
+      precedences += [TokenPrecedence(.identifier), TokenPrecedence(.contextualKeyword)]
     }
+    let recoveryPrecedence = recoveryPrecedence ?? precedences.min()!
+
     while !self.at(.eof) {
       if !recoveryPrecedence.shouldSkipOverNewlines,
           self.currentToken.isAtStartOfLine {
@@ -55,5 +66,51 @@ extension Parser.Lookahead {
 
     return false
   }
+
+  /// Checks if we can reach a token in `subset` by skipping tokens that have
+  /// a precedence that have a lower ``TokenPrecedence`` than the minimum
+  /// precedence of a token in that subset.
+  /// If so, return the token that we can recover to and a handle that can be
+  /// used to consume the unexpected tokens and the token we recovered to.
+  mutating func canRecoverTo<Subset: RawTokenKindSubset>(
+    anyIn subset: Subset.Type,
+    recoveryPrecedence: TokenPrecedence? = nil
+  ) -> (Subset, RecoveryConsumptionHandle)? {
+    assert(!subset.allCases.isEmpty, "Subset must have at least one case")
+    let recoveryPrecedence = recoveryPrecedence ?? subset.allCases.map({
+      if let precedence = $0.precedence {
+        return precedence
+      } else {
+        return TokenPrecedence($0.rawTokenKind)
+      }
+    }).min()!
+    let initialTokensConsumed = self.tokensConsumed
+    assert(!subset.allCases.isEmpty)
+    while !self.at(.eof) {
+      if !recoveryPrecedence.shouldSkipOverNewlines,
+          self.currentToken.isAtStartOfLine {
+        break
+      }
+      if let (kind, handle) = self.at(anyIn: subset) {
+        return (kind, RecoveryConsumptionHandle(
+          unexpectedTokens: self.tokensConsumed - initialTokensConsumed,
+          tokenConsumptionHandle: handle
+        ))
+      }
+      let currentTokenPrecedence = TokenPrecedence(self.currentToken.tokenKind)
+      if currentTokenPrecedence >= recoveryPrecedence {
+        break
+      }
+      self.consumeAnyToken()
+      if let closingDelimiter = currentTokenPrecedence.closingTokenKind {
+        guard self.canRecoverTo([closingDelimiter]) else {
+          break
+        }
+        self.eat(closingDelimiter)
+      }
+    }
+
+    return nil
+  }
 }
 

Sources/SwiftParser/TokenPrecedence.swift

@@ -87,7 +87,8 @@ public enum TokenPrecedence: Comparable {
     return self >= .stmtKeyword
   }
 
-  init(_ tokenKind: RawTokenKind) {
+  @_spi(RawSyntax)
+  public init(_ tokenKind: RawTokenKind) {
     switch tokenKind {
       // MARK: Identifier like
     case

Sources/SwiftParser/TopLevel.swift

@@ -46,7 +46,7 @@ extension Parser {
   mutating func parseTopLevelCodeBlockItems() -> RawCodeBlockItemListSyntax {
     var elements = [RawCodeBlockItemSyntax]()
     var loopProgress = LoopProgressCondition()
-    while let newElement = self.parseCodeBlockItem(), loopProgress.evaluate(currentToken) {
+    while let newElement = self.parseCodeBlockItem(isAtTopLevel: true), loopProgress.evaluate(currentToken) {
       elements.append(newElement)
     }
     return .init(elements: elements, arena: self.arena)
@@ -113,10 +113,10 @@ extension Parser {
   ///     statement → compiler-control-statement
   ///     statements → statement statements?
   @_spi(RawSyntax)
-  public mutating func parseCodeBlockItem() -> RawCodeBlockItemSyntax? {
+  public mutating func parseCodeBlockItem(isAtTopLevel: Bool = false) -> RawCodeBlockItemSyntax? {
     // FIXME: It is unfortunate that the Swift book refers to these as
     // "statements" and not "items".
-    let item = self.parseItem()
+    let item = self.parseItem(isAtTopLevel: isAtTopLevel)
     let semi = self.consume(if: .semicolon)
 
     if item.raw.byteLength == 0 && semi == nil {
@@ -125,7 +125,12 @@ extension Parser {
     return .init(item: item, semicolon: semi, errorTokens: nil, arena: self.arena)
   }
 
-  private mutating func parseItem() -> RawSyntax {
+  /// `isAtTopLevel` determins whether this is trying to parse an item that's at
+  /// the top level of the source file. If this is the case, we allow skipping
+  /// closing braces while trying to recover to the next item.
+  /// If we are not at the top level, such a closing brace should close the
+  /// wrapping declaration instead of being consumed by lookeahead.
+  private mutating func parseItem(isAtTopLevel: Bool = false) -> RawSyntax {
     if self.at(.poundIfKeyword) {
       return RawSyntax(self.parsePoundIfDirective {
         $0.parseCodeBlockItem()
@@ -142,6 +147,8 @@ extension Parser {
       return RawSyntax(self.parseStatement())
     } else if self.atStartOfExpression() {
       return RawSyntax(self.parseExpression())
+    } else if self.atStartOfDeclaration(isAtTopLevel: isAtTopLevel, allowRecovery: true) {
+      return RawSyntax(self.parseDeclaration())
     } else {
       return RawSyntax(RawMissingExprSyntax(arena: self.arena))
     }