From e67583b007490e4a152bfa244bafc465183278c6 Mon Sep 17 00:00:00 2001
From: Shiva Huang <shiva.huang@cocoazen.com>
Date: Fri, 26 Sep 2025 21:15:55 +0800
Subject: [PATCH] Fix JSONEncoder key conversion for CodingKeyRepresentable
 dictionary keys

When encoding dictionaries with CodingKeyRepresentable keys (e.g., custom
enum or struct keys), JSONEncoder was incorrectly applying key encoding
strategies like convertToSnakeCase to the dictionary keys themselves,
rather than treating them as semantic keys that should remain unchanged.

This change:
- Expands _JSONStringDictionaryEncodableMarker protocol coverage from
  String-keyed dictionaries to any CodingKeyRepresentable-keyed dictionaries
- Renames the protocol to _JSONCodingKeyRepresentableDictionaryEncodableMarker
  to better reflect its expanded scope
- Updates the encoding logic to handle CodingKeyRepresentable keys by
  converting them directly to their string representation

Fixes encoding behavior where Dictionary<CustomKey, Value> keys were
being transformed (e.g., "leaveMeAlone" -> "leave_me_alone") when they
should preserve their original form.
---
 .../JSON/JSONEncoder.swift                     | 18 ++++++++++--------
 .../JSONEncoderTests.swift                     | 16 ++++++++++++++++
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/Sources/FoundationEssentials/JSON/JSONEncoder.swift b/Sources/FoundationEssentials/JSON/JSONEncoder.swift
index 84159bbd1..046345fd8 100644
--- a/Sources/FoundationEssentials/JSON/JSONEncoder.swift
+++ b/Sources/FoundationEssentials/JSON/JSONEncoder.swift
@@ -1184,14 +1184,16 @@ private extension __JSONEncoder {
         }
     }
 
-    func wrap(_ dict: [String : Encodable], for additionalKey: (some CodingKey)? = _CodingKey?.none) throws -> JSONEncoderValue? {
+    func wrap(_ dict: _JSONCodingKeyRepresentableDictionaryEncodableMarker, for additionalKey: (some CodingKey)? = _CodingKey?.none) throws -> JSONEncoderValue? {
+        let dict = dict as! [AnyHashable: Encodable]
         var result = [String: JSONEncoderValue]()
         result.reserveCapacity(dict.count)
 
         let encoder = __JSONEncoder(options: self.options, ownerEncoder: self)
         for (key, value) in dict {
-            encoder.codingKey = _CodingKey(stringValue: key)
-            result[key] = try encoder.wrap(value)
+            let stringKey = (key.base as! CodingKeyRepresentable).codingKey.stringValue
+            encoder.codingKey = _CodingKey(stringValue: stringKey)
+            result[stringKey] = try encoder.wrap(value)
         }
 
         return .object(result)
@@ -1214,8 +1216,8 @@ private extension __JSONEncoder {
             return self.wrap(url.absoluteString)
         } else if let decimal = value as? Decimal {
             return .number(decimal.description)
-        } else if let encodable = value as? _JSONStringDictionaryEncodableMarker {
-            return try self.wrap(encodable as! [String:Encodable], for: additionalKey)
+        } else if let encodable = value as? _JSONCodingKeyRepresentableDictionaryEncodableMarker {
+            return try self.wrap(encodable, for: additionalKey)
         } else if let array = value as? _JSONDirectArrayEncodable {
             if options.outputFormatting.contains(.prettyPrinted) {
                 let (bytes, lengths) = try array.individualElementRepresentation(encoder: self, additionalKey)
@@ -1362,11 +1364,11 @@ extension JSONEncoder : @unchecked Sendable {}
 // Special-casing Support
 //===----------------------------------------------------------------------===//
 
-/// A marker protocol used to determine whether a value is a `String`-keyed `Dictionary`
+/// A marker protocol used to determine whether a value is a `CodingKeyRepresentable`-keyed `Dictionary`
 /// containing `Encodable` values (in which case it should be exempt from key conversion strategies).
-private protocol _JSONStringDictionaryEncodableMarker { }
+private protocol _JSONCodingKeyRepresentableDictionaryEncodableMarker { }
 
-extension Dictionary : _JSONStringDictionaryEncodableMarker where Key == String, Value: Encodable { }
+extension Dictionary : _JSONCodingKeyRepresentableDictionaryEncodableMarker where Key: CodingKeyRepresentable, Value: Encodable { }
 
 /// A protocol used to determine whether a value is an `Array` containing values that allow
 /// us to bypass UnkeyedEncodingContainer overhead by directly encoding the contents as
diff --git a/Tests/FoundationEssentialsTests/JSONEncoderTests.swift b/Tests/FoundationEssentialsTests/JSONEncoderTests.swift
index 566919ece..a2a76ca2e 100644
--- a/Tests/FoundationEssentialsTests/JSONEncoderTests.swift
+++ b/Tests/FoundationEssentialsTests/JSONEncoderTests.swift
@@ -2333,6 +2333,22 @@ extension JSONEncoderTests {
 
         #expect(expected == resultString)
     }
+    
+    @Test func encodingDictionaryCodingKeyRepresentableKeyConversionUntouched() throws {
+        struct Key: RawRepresentable, CodingKeyRepresentable, Hashable, Codable {
+            let rawValue: String
+        }
+        
+        let expected = "{\"leaveMeAlone\":\"test\"}"
+        let toEncode: [Key: String] = [Key(rawValue: "leaveMeAlone"): "test"]
+
+        let encoder = JSONEncoder()
+        encoder.keyEncodingStrategy = .convertToSnakeCase
+        let resultData = try encoder.encode(toEncode)
+        let resultString = String(bytes: resultData, encoding: .utf8)
+
+        #expect(expected == resultString)
+    }
 
     @Test func keyStrategySnakeGeneratedAndCustom() throws {
         // Test that this works with a struct that has automatically generated keys