Fix a parameter label recovery edge case

When `parseFunctionParameter()` parses two argument labels and then doesn’t find a type, it applies a heuristic to decide whether to reinterpret the second label as a type (and recover by inserting a colon between the labels). The code in this path could drop unexpected nodes between the two labels. Correct this issue and, if the unexpected syntax includes a module selector, reconstruct it and attach it to the type.

This was probably a pre-existing bug, but the module selector tests managed to hit it through mutation testing.

Author: beccadax

Sources/SwiftParser/Names.swift

@@ -65,7 +65,11 @@ private enum StructuralTokens: TokenSpecSet {
   case rightAngle
 
   init?(lexeme: Lexer.Lexeme, experimentalFeatures: Parser.ExperimentalFeatures) {
-    switch lexeme.rawTokenKind {
+    self.init(lexeme.rawTokenKind)
+  }
+
+  init?(_ rawTokenKind: RawTokenKind) {
+    switch rawTokenKind {
     case .comma: self = .comma
     case .colon: self = .colon
     case .leftParen: self = .leftParen
@@ -124,6 +128,10 @@ extension TokenConsumer {
     return operation(&self)
   }
 
+  typealias ModuleSelectorTokens = (
+    moduleNameOrUnexpected: Token, colonColonToken: Token, extra: [Token]
+  )
+
   /// If the subsequent tokens have the form of a module selector, valid or otherwise, consume and return them;
   /// otherwise consume nothing and return `nil`. Additionally consumes invalid chained module selectors.
   ///
@@ -134,9 +142,7 @@ extension TokenConsumer {
   ///  - `colonColonToken`: The `::` indicating this module selector. Always `.colonColon`, always present.
   ///  - `extra`: Tokens for additional trailing module selectors. There is no situation in which two module selectors
   ///    can be validly chained.
-  mutating func consumeModuleSelectorTokens() -> (
-    moduleNameOrUnexpected: Token, colonColonToken: Token, extra: [Token]
-  )? {
+  mutating func consumeModuleSelectorTokens() -> ModuleSelectorTokens? {
     guard self.isAtModuleSelector() else {
       return nil
     }
@@ -166,61 +172,100 @@ extension TokenConsumer {
   }
 }
 
-extension Parser {
-  /// Parses one or more module selectors, if present.
-  mutating func parseModuleSelector() -> RawModuleSelectorSyntax? {
-    guard let (moduleNameOrUnexpected, colonColon, extra) = consumeModuleSelectorTokens() else {
-      return nil
+private func tokens(in node: (some RawSyntaxNodeProtocol)?) -> [RawTokenSyntax] {
+  var tokens: [RawTokenSyntax] = []
+
+  func collect(from child: (some RawSyntaxNodeProtocol)?) {
+    guard let child else { return }
+
+    guard let token = child.as(RawTokenSyntax.self) else {
+      for child in child.raw.layoutView!.children {
+        collect(from: child)
+      }
+      return
     }
 
+    if !token.isMissing {
+      tokens.append(token)
+    }
+  }
+
+  collect(from: node)
+
+  return tokens
+}
+
+extension Parser {
+  /// Helper to form a module selector from a set of tokens, divorced from their source.
+  private mutating func makeModuleSelector(from tokens: ModuleSelectorTokens) -> RawModuleSelectorSyntax {
     let leadingUnexpected: [RawSyntax]
     let moduleName: RawTokenSyntax
     let trailingUnexpected: [RawSyntax]
 
-    if moduleNameOrUnexpected.tokenKind == .identifier {
+    if tokens.moduleNameOrUnexpected.tokenKind == .identifier {
       leadingUnexpected = []
-      moduleName = moduleNameOrUnexpected
+      moduleName = tokens.moduleNameOrUnexpected
     } else {
-      leadingUnexpected = [RawSyntax(moduleNameOrUnexpected)]
-      moduleName = RawTokenSyntax.init(missing: .identifier, arena: arena)
+      leadingUnexpected = [RawSyntax(tokens.moduleNameOrUnexpected)]
+      moduleName = self.missingToken(.identifier)
     }
 
-    trailingUnexpected = extra.map { RawSyntax($0) }
+    trailingUnexpected = tokens.extra.map { RawSyntax($0) }
 
     return RawModuleSelectorSyntax(
       RawUnexpectedNodesSyntax(leadingUnexpected, arena: arena),
       moduleName: moduleName,
-      colonColon: colonColon,
+      colonColon: tokens.colonColonToken,
       RawUnexpectedNodesSyntax(trailingUnexpected, arena: arena),
       arena: arena
     )
   }
 
-  /// Convert a raw syntax node to `RawUnexpectedNodesSyntax` by extracting its tokens.
-  func unexpected(_ node: (some RawSyntaxNodeProtocol)?) -> RawUnexpectedNodesSyntax? {
-    var tokens: [RawTokenSyntax] = []
+  /// Parses one or more module selectors, if present.
+  mutating func parseModuleSelector() -> RawModuleSelectorSyntax? {
+    guard let tokens = consumeModuleSelectorTokens() else {
+      return nil
+    }
+    return makeModuleSelector(from: tokens)
+  }
 
-    func collect(from child: (some RawSyntaxNodeProtocol)?) {
-      guard let child else { return }
+  /// If possible, convert a ``RawUnexpectedNodesSyntax`` back into a module selector. Used in some specialized
+  /// recovery paths.
+  mutating func recoverModuleSelector(
+    from unexpected: RawUnexpectedNodesSyntax?
+  ) -> (unexpectedBeforeModuleSelector: RawUnexpectedNodesSyntax?, moduleSelector: RawModuleSelectorSyntax?) {
+    guard let unexpected, unexpected.containsToken(where: { $0.tokenKind == .colonColon }) else {
+      return (unexpected, nil)
+    }
 
-      guard let token = child.as(RawTokenSyntax.self) else {
-        for child in child.raw.layoutView!.children {
-          collect(from: child)
-        }
-        return
-      }
+    var remainingTokens = tokens(in: unexpected)
+    let selectorTokens: ModuleSelectorTokens
 
-      if !token.isMissing {
-        tokens.append(token)
-      }
-    }
+    // Everything after the first `::` should be either extra (invalid) module selectors or just random extra garbage
+    // we skipped during recovery.
+    let afterColonColonIndex = remainingTokens.firstIndex(where: { $0.tokenKind == .colonColon })! + 1
+    selectorTokens.extra = Array(remainingTokens[afterColonColonIndex...])
+    remainingTokens.removeSubrange(afterColonColonIndex...)
 
-    collect(from: node)
+    selectorTokens.colonColonToken = remainingTokens.removeLast()
 
-    guard !tokens.isEmpty else {
-      return nil
+    // This matches logic from `consumeModuleSelectorTokens()` and `isAtModuleSelector()`.
+    if let tokenBeforeColonColon = remainingTokens.last, StructuralTokens(tokenBeforeColonColon.tokenKind) != nil {
+      selectorTokens.moduleNameOrUnexpected = tokenBeforeColonColon
+      remainingTokens.removeLast()
+    } else {
+      selectorTokens.moduleNameOrUnexpected = self.missingToken(.identifier)
     }
-    return RawUnexpectedNodesSyntax(tokens, arena: self.arena)
+
+    return (
+      RawUnexpectedNodesSyntax(remainingTokens, arena: self.arena),
+      self.makeModuleSelector(from: selectorTokens)
+    )
+  }
+
+  /// Convert a raw syntax node to `RawUnexpectedNodesSyntax` by extracting its tokens.
+  func unexpected(_ node: (some RawSyntaxNodeProtocol)?) -> RawUnexpectedNodesSyntax? {
+    return RawUnexpectedNodesSyntax(tokens(in: node), arena: self.arena)
   }
 
   /// Inject `moduleSelector` into `node`, its first child, its first child's first child, etc. If necessary,

Sources/SwiftParser/Parameters.swift

@@ -129,9 +129,11 @@ extension Parser {
       secondName.tokenText.isStartingWithUppercase
     {
       // Synthesize the secondName parameter as a type node.
+      let (stillUnexpected, moduleSelector) = recoverModuleSelector(from: names.unexpectedBeforeSecondName)
       type = RawTypeSyntax(
         RawIdentifierTypeSyntax(
-          moduleSelector: nil,
+          stillUnexpected,
+          moduleSelector: moduleSelector,
           name: secondName,
           genericArgumentClause: nil,
           arena: self.arena

Tests/SwiftParserTest/translated/ModuleSelectorTests.swift

@@ -516,6 +516,64 @@ final class ModuleSelectorTests: ParserTestCase {
         ) {}
         """
     )
+    assertParse(
+      """
+      func decl1(
+        1️⃣main::p12️⃣3️⃣< ::A
+      ) {}
+      """,
+      diagnostics: [
+        DiagnosticSpec(
+          locationMarker: "1️⃣",
+          message: "parameter cannot be qualified with a module selector",
+          fixIts: ["remove 'main::'"]
+        ),
+        DiagnosticSpec(
+          locationMarker: "2️⃣",
+          message: "expected ':' in parameter",
+          fixIts: ["insert ':'"]
+        ),
+        DiagnosticSpec(
+          locationMarker: "3️⃣",
+          message: "'<' is not a valid identifier"
+        ),
+
+      ],
+      fixedSource: """
+        func decl1(
+          p1: < ::A
+        ) {}
+        """
+    )
+    assertParse(
+      """
+      func decl1(
+        1️⃣main::p12️⃣3️⃣> ::A
+      ) {}
+      """,
+      diagnostics: [
+        DiagnosticSpec(
+          locationMarker: "1️⃣",
+          message: "parameter cannot be qualified with a module selector",
+          fixIts: ["remove 'main::'"]
+        ),
+        DiagnosticSpec(
+          locationMarker: "2️⃣",
+          message: "expected ':' in parameter",
+          fixIts: ["insert ':'"]
+        ),
+        DiagnosticSpec(
+          locationMarker: "3️⃣",
+          message: "'>' is not a valid identifier"
+        ),
+
+      ],
+      fixedSource: """
+        func decl1(
+          p1: > ::A
+        ) {}
+        """
+    )
     assertParse(
       """
       func decl1(