Recover if func or var is missing for a member decl item

Author: ahoppen

Sources/SwiftParser/Declarations.swift

@@ -99,14 +99,16 @@ extension TokenConsumer {
       // When 'case' appears inside a function, it's probably a switch
       // case, not an enum case declaration.
       return false
+    case .some(_):
+      // All other decl start keywords unconditonally start a decl.
+      return true
     case nil:
       if subparser.at(anyIn: ContextualDeclKeyword.self)?.0 != nil {
         subparser.consumeAnyToken()
         return subparser.atStartOfDeclaration(
           isAtTopLevel: isAtTopLevel, allowRecovery: allowRecovery)
       }
       return false
-    default: return true
     }
   }
 }
@@ -146,8 +148,11 @@ extension Parser {
   ///     declaration → precedence-group-declaration
   ///
   ///     declarations → declaration declarations?
+  ///
+  /// If `allowMissingFuncOrVarKeywordRecovery` is `true`, this methods tries to
+  /// synthesize `func` or `var` keywords where necessary.
   @_spi(RawSyntax)
-  public mutating func parseDeclaration() -> RawDeclSyntax {
+  public mutating func parseDeclaration(allowMissingFuncOrVarKeywordRecovery: Bool = false) -> RawDeclSyntax {
     switch self.at(anyIn: PoundDeclarationStart.self) {
     case (.poundIfKeyword, _)?:
       let directive = self.parsePoundIfDirective { parser in
@@ -212,10 +217,25 @@ extension Parser {
     case (.actorContextualKeyword, let handle)?:
       return RawDeclSyntax(self.parseActorDeclaration(attrs, handle))
     case nil:
+      if allowMissingFuncOrVarKeywordRecovery {
+        let isProbablyVarDecl = self.at(any: [.identifier, .wildcardKeyword]) && self.peek().tokenKind.is(any: [.colon, .equal, .comma])
+        let isProbablyTupleDecl = self.at(.leftParen) && self.peek().tokenKind.is(any: [.identifier, .wildcardKeyword])
+
+        if isProbablyVarDecl || isProbablyTupleDecl {
+          return RawDeclSyntax(self.parseLetOrVarDeclaration(attrs, .missing(.varKeyword)))
+        }
+
+        let isProbablyFuncDecl = self.at(any: [.identifier, .wildcardKeyword]) || self.at(anyIn: Operator.self) != nil
+
+        if isProbablyFuncDecl {
+          return RawDeclSyntax(self.parseFuncDeclaration(attrs, .missing(.funcKeyword)))
+        }
+      }
       return RawDeclSyntax(RawMissingDeclSyntax(
         attributes: attrs.attributes,
         modifiers: attrs.modifiers,
-        arena: self.arena))
+        arena: self.arena
+      ))
     }
   }
 }
@@ -586,7 +606,7 @@ extension Parser {
     if self.at(.poundSourceLocationKeyword) {
       decl = RawDeclSyntax(self.parsePoundSourceLocationDirective())
     } else {
-      decl = self.parseDeclaration()
+      decl = self.parseDeclaration(allowMissingFuncOrVarKeywordRecovery: true)
     }
 
     let semi = self.consume(if: .semicolon)

Sources/SwiftParser/Diagnostics/DiagnosticExtensions.swift

@@ -76,10 +76,22 @@ extension FixIt.Changes {
     if let trailingTrivia = trailingTrivia {
       presentNode = presentNode.withTrailingTrivia(trailingTrivia)
     }
-    return [.replace(
-      oldNode: Syntax(node),
-      newNode: Syntax(presentNode)
-    )]
+    if node.shouldBeInsertedAfterNextTokenTrivia,
+       let nextToken = node.nextToken(viewMode: .sourceAccurate),
+       leadingTrivia == nil {
+      return [
+        .replace(
+          oldNode: Syntax(node),
+          newNode: Syntax(presentNode).withLeadingTrivia(nextToken.leadingTrivia)
+        ),
+        .replaceLeadingTrivia(token: nextToken, newTrivia: [])
+      ]
+    } else {
+      return [.replace(
+        oldNode: Syntax(node),
+        newNode: Syntax(presentNode)
+      )]
+    }
   }
 
   /// Makes the `token` present, moving it in front of the previous token's trivia.

Sources/SwiftParser/Diagnostics/MissingNodesError.swift

@@ -343,9 +343,16 @@ extension ParseDiagnosticsGenerator {
       notes.append(Note(node: Syntax(startToken), message: MatchingOpeningTokenNote(openingToken: startToken)))
     }
 
+    let position: AbsolutePosition
+    if node.shouldBeInsertedAfterNextTokenTrivia, let nextToken = node.nextToken(viewMode: .sourceAccurate) {
+      position = nextToken.positionAfterSkippingLeadingTrivia
+    } else {
+      position = node.endPosition
+    }
+
     addDiagnostic(
       node,
-      position: node.endPosition,
+      position: position,
       MissingNodesError(missingNodes: missingNodes),
       notes: notes,
       fixIts: [fixIt],

Sources/SwiftParser/Diagnostics/SyntaxExtensions.swift

@@ -92,15 +92,32 @@ extension SyntaxProtocol {
 
   /// Returns this node or the first ancestor that satisfies `condition`.
   func ancestorOrSelf(where condition: (Syntax) -> Bool) -> Syntax? {
+    return ancestorOrSelf(mapping: { condition($0) ? $0 : nil })
+  }
+
+  /// Returns this node or the first ancestor that satisfies `condition`.
+  func ancestorOrSelf<T>(mapping map: (Syntax) -> T?) -> T? {
     var walk: Syntax? = Syntax(self)
     while let unwrappedParent = walk {
-      if condition(unwrappedParent) {
-        return walk
+      if let mapped = map(unwrappedParent) {
+        return mapped
       }
       walk = unwrappedParent.parent
     }
     return nil
   }
+
+  /// Returns `true` if the next token's leading trivia should be made leading trivia
+  /// of this mode, when it is switched from being missing to present.
+  var shouldBeInsertedAfterNextTokenTrivia: Bool {
+    if !self.raw.kind.isMissing,
+       let memberDeclItem = self.ancestorOrSelf(mapping: { $0.as(MemberDeclListItemSyntax.self) }),
+       memberDeclItem.firstToken(viewMode: .all) == self.firstToken(viewMode: .all) {
+      return true
+    } else {
+      return false
+    }
+  }
 }
 
 extension TokenKind {

Sources/SwiftParser/Recovery.swift

@@ -29,6 +29,15 @@ public struct RecoveryConsumptionHandle {
       unexpectedTokens: 0,
       tokenConsumptionHandle: TokenConsumptionHandle(tokenKind: token))
   }
+
+  /// A `RecoveryConsumptionHandle` that will not eat any tokens but instead
+  /// synthesize a missing token of kind `token`.
+  @_spi(RawSyntax)
+  public static func missing(_ token: RawTokenKind) -> RecoveryConsumptionHandle {
+    return RecoveryConsumptionHandle(
+      unexpectedTokens: 0,
+      tokenConsumptionHandle: TokenConsumptionHandle(tokenKind: token, missing: true))
+  }
 }
 
 extension Parser.Lookahead {

Sources/SwiftParser/SyntaxUtils.swift

@@ -47,3 +47,9 @@ extension SyntaxText {
     self.hasSuffix(SyntaxText("#>"))
   }
 }
+
+extension RawTokenKind {
+  func `is`(any kinds: [RawTokenKind]) -> Bool {
+    return kinds.contains(self)
+  }
+}

Sources/SwiftParser/TokenConsumer.swift

@@ -45,6 +45,9 @@ struct TokenConsumptionHandle {
   var tokenKind: RawTokenKind
   /// When not `nil`, the token's kind will be remapped to this kind when consumed.
   var remappedKind: RawTokenKind?
+  /// If `true`, the token we should consume should be synthesized as a missing token
+  /// and no tokens should be consumed.
+  var missing: Bool = false
 }
 
 extension TokenConsumer {
@@ -116,10 +119,13 @@ extension TokenConsumer {
 
   /// Eat a token that we know we are currently positioned at, based on `at(anyIn:)`.
   mutating func eat(_ handle: TokenConsumptionHandle) -> Token {
-    assert(self.at(handle.tokenKind))
-    if let remappedKind = handle.remappedKind {
+    if handle.missing {
+      return missingToken(handle.remappedKind ?? handle.tokenKind, text: nil)
+    } else if let remappedKind = handle.remappedKind {
+      assert(self.at(handle.tokenKind))
       return consumeAnyToken(remapping: remappedKind)
     } else {
+      assert(self.at(handle.tokenKind))
       return consumeAnyToken()
     }
   }

Tests/SwiftParserTest/Declarations.swift

@@ -717,11 +717,13 @@ final class DeclarationTests: XCTestCase {
     AssertParse(
       """
       struct S {
-        1️⃣/ ###line 25 "line-directive.swift"
+        1️⃣/ 2️⃣###line 25 "line-directive.swift"
       }
       """,
       diagnostics: [
-        DiagnosticSpec(locationMarker: "1️⃣", message: #"unexpected code '/ ###line 25 "line-directive.swift"' in struct"#)
+        DiagnosticSpec(locationMarker: "1️⃣", message: "expected 'func' in function"),
+        DiagnosticSpec(locationMarker: "2️⃣", message: "expected parameter clause in function signature"),
+        DiagnosticSpec(locationMarker: "2️⃣", message: #"unexpected code '###line 25 "line-directive.swift"' in struct"#)
       ]
     )
   }