Add an extra percent encoding layer when encoding DocumentURIs to LSP requests

Sources/LanguageServerProtocol/SupportTypes/DocumentURI.swift

@@ -84,7 +84,19 @@ public struct DocumentURI: Codable, Hashable, Sendable {
   }
 
   public init(from decoder: Decoder) throws {
-    try self.init(string: decoder.singleValueContainer().decode(String.self))
+    let string = try decoder.singleValueContainer().decode(String.self)
+    guard let url = URL(string: string) else {
+      throw FailedToConstructDocumentURIFromStringError(string: string)
+    }
+    if url.query() != nil, var urlComponents = URLComponents(string: url.absoluteString) {
+      // See comment in `encode(to:)`
+      urlComponents.percentEncodedQuery = urlComponents.percentEncodedQuery!.removingPercentEncoding
+      if let rewrittenQuery = urlComponents.url {
+        self.init(rewrittenQuery)
+        return
+      }
+    }
+    self.init(url)
   }
 
   /// Equality check to handle escape sequences in file URLs.
@@ -97,7 +109,36 @@ public struct DocumentURI: Codable, Hashable, Sendable {
     hasher.combine(self.pseudoPath)
   }
 
+  private static let additionalQueryEncodingCharacterSet = CharacterSet(charactersIn: "?=&%").inverted
+
   public func encode(to encoder: Encoder) throws {
-    try storage.absoluteString.encode(to: encoder)
+    let urlToEncode: URL
+    if let query = storage.query(percentEncoded: true), var components = URLComponents(string: storage.absoluteString) {
+      // The URI standard RFC 3986 is ambiguous about whether percent encoding and their represented characters are
+      // considered equivalent. VS Code considers them equivalent and treats them the same:
+      //
+      // vscode.Uri.parse("x://a?b=xxxx%3Dyyyy").toString() -> 'x://a?b%3Dxxxx%3Dyyyy'
+      // vscode.Uri.parse("x://a?b=xxxx%3Dyyyy").toString(/*skipEncoding=*/true) -> 'x://a?b=xxxx=yyyy'
+      //
+      // This causes issues because SourceKit-LSP's macro expansion URLs encoded by URLComponents use `=` to denote the
+      // separation of a key and a value in the outer query. The value of the `parent` key may itself contain query
+      // items, which use the escaped form '%3D'. Simplified, such a URL may look like
+      // scheme://host?parent=scheme://host?line%3D2
+      // But after running this through VS Code's URI type `=` and `%3D` get canonicalized and are indistinguishable.
+      // To avoid this ambiguity, always percent escape the characters we use to distinguish URL query parameters,
+      // producing the following URL.
+      // scheme://host?parent%3Dscheme://host%3Fline%253D2
+      components.percentEncodedQuery =
+        query
+        .addingPercentEncoding(withAllowedCharacters: Self.additionalQueryEncodingCharacterSet)
+      if let componentsUrl = components.url {
+        urlToEncode = componentsUrl
+      } else {
+        urlToEncode = self.storage
+      }
+    } else {
+      urlToEncode = self.storage
+    }
+    try urlToEncode.absoluteString.encode(to: encoder)
   }
 }

Sources/SKTestSupport/CheckCoding.swift

@@ -43,9 +43,12 @@ package func checkCoding<T: Codable & Equatable>(
   XCTAssertEqual(json, str, file: file, line: line)
 
   let decoder = JSONDecoder()
-  let decodedValue = try! decoder.decode(WrapFragment<T>.self, from: data).value
-
-  XCTAssertEqual(value, decodedValue, file: file, line: line)
+  XCTAssertNoThrow(
+    try {
+      let decodedValue = try decoder.decode(WrapFragment<T>.self, from: data).value
+      XCTAssertEqual(value, decodedValue, file: file, line: line)
+    }()
+  )
 }
 
 /// JSONEncoder requires the top-level value to be encoded as a JSON container (array or object). Give it one.

Tests/LanguageServerProtocolTests/CodingTests.swift

@@ -1343,6 +1343,29 @@ final class CodingTests: XCTestCase {
         """
     )
   }
+
+  func testDocumentUriQueryParameterCoding() throws {
+    checkCoding(
+      try DocumentURI(string: "scheme://host?parent=scheme://host?line%3D2"),
+      json: #"""
+        "scheme:\/\/host?parent%3Dscheme:\/\/host%3Fline%253D2"
+        """#
+    )
+
+    checkCoding(
+      try DocumentURI(string: "scheme://host?parent=scheme://host?parent%3Dscheme://host%3Fkey%253Dvalue"),
+      json: #"""
+        "scheme:\/\/host?parent%3Dscheme:\/\/host%3Fparent%253Dscheme:\/\/host%253Fkey%25253Dvalue"
+        """#
+    )
+
+    checkCoding(
+      try DocumentURI(string: "scheme://host?parent=with%23hash"),
+      json: #"""
+        "scheme:\/\/host?parent%3Dwith%2523hash"
+        """#
+    )
+  }
 }
 
 func with<T>(_ value: T, mutate: (inout T) -> Void) -> T {