Fix up effect specifiers in function signatures

Author: bnbarham

Sources/SwiftParser/Declarations.swift

@@ -1391,26 +1391,19 @@ extension Parser {
       arena: self.arena)
   }
 
-  /// If a `throws` keyword appears right in front of the `arrow`, it is returned as `misplacedThrowsKeyword` so it can be synthesized in front of the arrow.
+  /// If effect specifiers appear in front of the `arrow` or after the return type, they are returned so that they can be synthesized in front of the arrow.
   @_spi(RawSyntax)
-  public mutating func parseFunctionReturnClause() -> (returnClause: RawReturnClauseSyntax, misplacedThrowsKeyword: RawTokenSyntax?) {
-    let (unexpectedBeforeArrow, arrow) = self.expect(.arrow)
-    var misplacedThrowsKeyword: RawTokenSyntax? = nil
-    let unexpectedBeforeReturnType: RawUnexpectedNodesSyntax?
-    if let throwsKeyword = self.consume(ifAny: [.rethrowsKeyword, .throwsKeyword]) {
-      misplacedThrowsKeyword = throwsKeyword
-      unexpectedBeforeReturnType = RawUnexpectedNodesSyntax(elements: [RawSyntax(throwsKeyword)], arena: self.arena)
-    } else {
-      unexpectedBeforeReturnType = nil
-    }
+  public mutating func parseFunctionReturnClause(_ handle: RecoveryConsumptionHandle) -> (returnClause: RawReturnClauseSyntax, misplacedEffects: RawUnexpectedNodesSyntax?) {
+    let (unexpectedBeforeArrow, arrow) = self.eat(handle)
+    let unexpectedBeforeReturnType = RawUnexpectedNodesSyntax(unexpectedTokens: parseAllEffectsSpecifiers(), arena: self.arena)
     let result = self.parseResultType()
     let returnClause = RawReturnClauseSyntax(
       unexpectedBeforeArrow,
       arrow: arrow,
       unexpectedBeforeReturnType,
       returnType: result,
       arena: self.arena)
-    return (returnClause, misplacedThrowsKeyword)
+    return (returnClause, unexpectedBeforeReturnType)
   }
 }
 
@@ -1495,32 +1488,27 @@ extension Parser {
   public mutating func parseFunctionSignature() -> RawFunctionSignatureSyntax {
     let input = self.parseParameterClause(for: .functionParameters)
 
-    let async: RawTokenSyntax?
-    if let asyncTok = self.consumeIfContextualKeyword("async") {
-      async = asyncTok
-    } else if let reasync = self.consumeIfContextualKeyword("reasync") {
-      async = reasync
-    } else {
-      async = nil
-    }
-
-    var throwsKeyword = self.consume(ifAny: [.throwsKeyword, .rethrowsKeyword])
+    var asyncOrReasync: RawTokenSyntax?
+    var throwsOrRethrows: RawTokenSyntax?
+    let unexpectedAfterThrows: RawUnexpectedNodesSyntax?
+    (asyncOrReasync, throwsOrRethrows, unexpectedAfterThrows) = parseExpectedEffectsSpecifiers()
 
     let output: RawReturnClauseSyntax?
-    if self.at(.arrow) {
-      let result = self.parseFunctionReturnClause()
+    if let handle = self.canRecoverTo(.arrow) {
+      let result = self.parseFunctionReturnClause(handle)
       output = result.returnClause
-      if let misplacedThrowsKeyword = result.misplacedThrowsKeyword, throwsKeyword == nil {
-        throwsKeyword = RawTokenSyntax(missing: misplacedThrowsKeyword.tokenKind, arena: self.arena)
-      }
+      addMissingEffectSpecifiers(from: result.misplacedEffects, asyncOrReasync: &asyncOrReasync, throwsOrRethrows: &throwsOrRethrows)
     } else {
       output = nil
     }
 
+    // Check for more specifiers?
+
     return RawFunctionSignatureSyntax(
       input: input,
-      asyncOrReasyncKeyword: async,
-      throwsOrRethrowsKeyword: throwsKeyword,
+      asyncOrReasyncKeyword: asyncOrReasync,
+      throwsOrRethrowsKeyword: throwsOrRethrows,
+      unexpectedAfterThrows,
       output: output,
       arena: self.arena)
   }
@@ -1553,8 +1541,8 @@ extension Parser {
     let indices = self.parseParameterClause(for: .indices)
 
     let result: RawReturnClauseSyntax
-    if self.at(.arrow) {
-      result = self.parseFunctionReturnClause().returnClause
+    if let handle = self.canRecoverTo(.arrow) {
+      result = self.parseFunctionReturnClause(handle).returnClause
     } else {
       result = RawReturnClauseSyntax(
         arrow: RawTokenSyntax(missing: .arrow, arena: self.arena),
@@ -1728,33 +1716,14 @@ extension Parser {
       attributes: attrs, modifier: modifier, kind: kind, token: introducer)
   }
 
-  @_spi(RawSyntax)
-  public mutating func parseEffectsSpecifier() -> RawTokenSyntax? {
-    // 'async'
-    if let async = self.consumeIfContextualKeyword("async") {
-      return async
-    }
-
-    // 'reasync'
-    if let reasync = self.consumeIfContextualKeyword("reasync") {
-      return reasync
+  mutating func parseEffectsSpecifier() -> RawTokenSyntax? {
+    if let (_, handle) = self.at(anyIn: EffectsSpecifier.self) {
+      return self.eat(handle)
     }
-
-    // 'throws'/'rethrows'
-    if let throwsRethrows = self.consume(ifAny: [.throwsKeyword, .rethrowsKeyword]) {
-      return throwsRethrows
-    }
-
-    // diagnose 'throw'/'try'.
-    if let throwTry = self.consume(ifAny: [.throwKeyword, .tryKeyword], where: { !$0.isAtStartOfLine }) {
-      return throwTry
-    }
-
     return nil
   }
 
-  @_spi(RawSyntax)
-  public mutating func parseEffectsSpecifiers() -> [RawTokenSyntax] {
+  mutating func parseAllEffectsSpecifiers() -> [RawTokenSyntax] {
     var specifiers = [RawTokenSyntax]()
     var loopProgress = LoopProgressCondition()
     while let specifier = self.parseEffectsSpecifier(), loopProgress.evaluate(currentToken) {
@@ -1763,6 +1732,69 @@ extension Parser {
     return specifiers
   }
 
+  private func addMissingEffectSpecifiers(from unexpectedTokens: RawUnexpectedNodesSyntax?, asyncOrReasync: inout RawTokenSyntax?, throwsOrRethrows: inout RawTokenSyntax?) {
+    guard let unexpectedTokens = unexpectedTokens else {
+      return
+    }
+
+    for unexpectedToken in unexpectedTokens.elements {
+      guard asyncOrReasync == nil || throwsOrRethrows == nil, let specifier = unexpectedToken.as(RawTokenSyntax.self) else {
+        break
+      }
+
+      switch specifier.tokenKind {
+      case .contextualKeyword:
+        if asyncOrReasync == nil, specifier.tokenText == "async" || specifier.tokenText == "reasync" {
+          asyncOrReasync = RawTokenSyntax(missing: specifier.tokenKind, text: specifier.tokenText, arena: self.arena)
+        }
+        break
+      case .throwsKeyword:
+        fallthrough
+      case .rethrowsKeyword:
+        if throwsOrRethrows == nil {
+          throwsOrRethrows = RawTokenSyntax(missing: specifier.tokenKind, text: specifier.tokenText, arena: self.arena)
+        }
+        break
+      default:
+        break
+      }
+    }
+  }
+
+  mutating func parseExpectedEffectsSpecifiers() -> (asyncOrReasync: RawTokenSyntax?, throwsOrRethrows: RawTokenSyntax?, unexpected: RawUnexpectedNodesSyntax?) {
+    guard let firstSpecifier = parseEffectsSpecifier() else {
+      return (nil, nil, nil)
+    }
+
+    var asyncOrReasync: RawTokenSyntax? = nil
+    var throwsOrRethrows: RawTokenSyntax? = nil
+    var unexpectedTokens = [RawTokenSyntax]()
+
+    // Not a fan of the number of places contextual + async/reasync is checked (or throws/rethrows for that matter). Can we use EffectSpecifier somehow?
+    if firstSpecifier.tokenKind == .contextualKeyword,
+       firstSpecifier.tokenText == "async" || firstSpecifier.tokenText == "reasync" {
+      asyncOrReasync = firstSpecifier
+
+      if let secondSpecifier = parseEffectsSpecifier() {
+        if secondSpecifier.tokenKind == .throwsKeyword || secondSpecifier.tokenKind == .rethrowsKeyword {
+          throwsOrRethrows = firstSpecifier
+        } else {
+          unexpectedTokens.append(secondSpecifier)
+        }
+      }
+    } else if firstSpecifier.tokenKind == .throwsKeyword || firstSpecifier.tokenKind == .rethrowsKeyword {
+      throwsOrRethrows = firstSpecifier
+    } else {
+      unexpectedTokens.append(firstSpecifier)
+    }
+    unexpectedTokens.append(contentsOf: parseAllEffectsSpecifiers())
+
+    let unexpected = RawUnexpectedNodesSyntax(unexpectedTokens: unexpectedTokens, arena: self.arena)
+    addMissingEffectSpecifiers(from: unexpected, asyncOrReasync: &asyncOrReasync, throwsOrRethrows: &throwsOrRethrows)
+
+    return (asyncOrReasync, throwsOrRethrows, unexpected)
+  }
+
   /// Parse the body of a variable declaration. This can include explicit
   /// getters, setters, and observers, or the body of a computed property.
   ///

Sources/SwiftParser/RawTokenKindSubset.swift

@@ -418,6 +418,15 @@ enum EffectsSpecifier: RawTokenKindSubset {
     default: return nil
     }
   }
+
+  var remappedKind: RawTokenKind? {
+    switch self {
+    case .asyncContextualKeyword: return .contextualKeyword
+    case .awaitContextualKeyword: return .contextualKeyword
+    case .reasyncContextualKeyword: return .contextualKeyword
+    default: return nil
+    }
+  }
 }
 
 enum IdentifierTokens: RawTokenKindSubset {

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -774,3 +774,11 @@ extension RawSyntax {
   }
 }
 
+extension RawUnexpectedNodesSyntax {
+  public init?(unexpectedTokens: [RawTokenSyntax], arena: SyntaxArena) {
+    if unexpectedTokens.isEmpty {
+      return nil
+    }
+    self.init(elements: unexpectedTokens.map {RawSyntax($0)}, arena: arena)
+  }
+}