Fix key encoding/decoding strategy applied to CodingKeyRepresentable dictionary keys (#1526)

- Fix JSONEncoder incorrectly applying key encoding strategy to CodingKeyRepresentable-keyed
  dictionary keys; only nested Encodable properties should be affected
- Fix JSONDecoder incorrectly applying key decoding strategy to CodingKeyRepresentable-keyed
  dictionary keys; only nested Decodable properties should be affected
- Extend _JSONCodingKeyRepresentableDictionaryEncodableMarker to replace the String-only
  marker, avoiding unnecessary AnyHashable boxing
- Extend decoder's dictionary marker protocol to cover all CodingKeyRepresentable keys
- Add tests verifying keys bypass the strategy while nested struct properties still apply it

Author: ShivaHuang

Sources/FoundationEssentials/JSON/JSONDecoder.swift

@@ -19,17 +19,44 @@ import Darwin
 internal import _FoundationCShims
 internal import Synchronization
 
-/// 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 `Decodable` values (in which case it should be exempt from key conversion strategies).
 ///
-/// The marker protocol also provides access to the type of the `Decodable` values,
-/// which is needed for the implementation of the key conversion strategy exemption.
-private protocol _JSONStringDictionaryDecodableMarker {
+/// The protocol provides `_fromStringKeyedDictionary` to convert a `[String: Any]` dictionary
+/// to the proper `[Key: Value]` type using `Key.init(codingKey:)`.
+private protocol _JSONCodingKeyRepresentableDictionaryDecodableMarker {
     static var elementType: Decodable.Type { get }
+    static func _fromStringKeyedDictionary(_ dict: [String: Any]) -> Self?
 }
 
-extension Dictionary : _JSONStringDictionaryDecodableMarker where Key == String, Value: Decodable {
+extension Dictionary : _JSONCodingKeyRepresentableDictionaryDecodableMarker where Key: CodingKeyRepresentable, Value: Decodable {
     static var elementType: Decodable.Type { return Value.self }
+
+    static func _fromStringKeyedDictionary(_ dict: [String: Any]) -> Self? {
+        // Fast path for String keys - no conversion needed
+        if Key.self == String.self {
+            return dict as? Self
+        }
+        // Convert keys for other CodingKeyRepresentable types
+        var result = Self()
+        result.reserveCapacity(dict.count)
+        for (stringKey, value) in dict {
+            guard let key = Key(codingKey: _DictionaryCodingKey(stringValue: stringKey)),
+                  let typedValue = value as? Value else {
+                return nil
+            }
+            result[key] = typedValue
+        }
+        return result
+    }
+}
+
+/// A simple CodingKey implementation for string-to-key conversion.
+private struct _DictionaryCodingKey: CodingKey {
+    var stringValue: String
+    var intValue: Int? { nil }
+    init(stringValue: String) { self.stringValue = stringValue }
+    init?(intValue: Int) { nil }
 }
 
 //===----------------------------------------------------------------------===//
@@ -611,7 +638,7 @@ extension JSONDecoderImpl: Decoder {
         if type == Decimal.self {
             return try self.unwrapDecimal(from: mapValue, for: codingPathNode, additionalKey) as! T
         }
-        if !options.keyDecodingStrategy.isDefault, T.self is _JSONStringDictionaryDecodableMarker.Type {
+        if !options.keyDecodingStrategy.isDefault, T.self is _JSONCodingKeyRepresentableDictionaryDecodableMarker.Type {
             return try self.unwrapDictionary(from: mapValue, as: type, for: codingPathNode, additionalKey)
         }
 
@@ -770,8 +797,8 @@ extension JSONDecoderImpl: Decoder {
     private func unwrapDictionary<T: Decodable>(from mapValue: JSONMap.Value, as type: T.Type, for codingPathNode: _CodingPathNode, _ additionalKey: (some CodingKey)? = nil) throws -> T {
         try checkNotNull(mapValue, expectedType: [String:Any].self, for: codingPathNode, additionalKey)
 
-        guard let dictType = type as? (_JSONStringDictionaryDecodableMarker & Decodable).Type else {
-            preconditionFailure("Must only be called if T implements __JSONStringDictionaryDecodableMarker")
+        guard let dictType = type as? _JSONCodingKeyRepresentableDictionaryDecodableMarker.Type else {
+            preconditionFailure("Must only be called if T implements _JSONCodingKeyRepresentableDictionaryDecodableMarker")
         }
 
         guard case let .object(region) = mapValue else {
@@ -781,8 +808,8 @@ extension JSONDecoderImpl: Decoder {
             ))
         }
 
-        var result = [String: Any]()
-        result.reserveCapacity(region.count / 2)
+        var stringKeyedResult = [String: Any]()
+        stringKeyedResult.reserveCapacity(region.count / 2)
 
         let dictCodingPathNode = codingPathNode.appending(additionalKey)
 
@@ -791,10 +818,17 @@ extension JSONDecoderImpl: Decoder {
             // We know these values are keys, but UTF-8 decoding could still fail.
             let key = try self.unwrapString(from: keyValue, for: dictCodingPathNode, _CodingKey?.none)
             let value = try self.unwrap(value, as: dictType.elementType, for: dictCodingPathNode, _CodingKey(stringValue: key)!)
-            result[key]._setIfNil(to: value)
+            stringKeyedResult[key]._setIfNil(to: value)
         }
 
-        return result as! T
+        // Convert [String: Any] to [Key: Value] using the marker protocol
+        guard let result = dictType._fromStringKeyedDictionary(stringKeyedResult) as? T else {
+            throw DecodingError.dataCorrupted(DecodingError.Context(
+                codingPath: codingPathNode.path(byAppending: additionalKey),
+                debugDescription: "Failed to create dictionary with CodingKeyRepresentable keys"
+            ))
+        }
+        return result
     }
 
     private func unwrapString(from value: JSONMap.Value, for codingPathNode: _CodingPathNode, _ additionalKey: (some CodingKey)? = nil) throws -> String {

Sources/FoundationEssentials/JSON/JSONEncoder.swift

@@ -1216,8 +1216,8 @@ private extension __JSONEncoder {
             return self.wrap(url.absoluteString)
         } else if let decimal = value as? Decimal {
             return .number(decimal.description)
-        } else if !options.keyEncodingStrategy.isDefault, let encodable = value as? _JSONStringDictionaryEncodableMarker {
-            return try self.wrap(encodable as! [String:Encodable], for: additionalKey)
+        } else if !options.keyEncodingStrategy.isDefault, let encodable = value as? _JSONCodingKeyRepresentableDictionaryEncodableMarker {
+            return try self.wrap(encodable._stringKeyedDictionary, for: additionalKey)
         } else if let array = _asDirectArrayEncodable(value) {
             if options.outputFormatting.contains(.prettyPrinted) {
                 let (bytes, lengths) = try array.individualElementRepresentation(encoder: self, additionalKey)
@@ -1400,11 +1400,30 @@ 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 { }
+///
+/// The protocol provides `_stringKeyedDictionary` to convert keys using their `codingKey.stringValue`,
+/// allowing the encoder to use the optimized String-keyed encoding path.
+private protocol _JSONCodingKeyRepresentableDictionaryEncodableMarker {
+    var _stringKeyedDictionary: [String: Encodable] { get }
+}
 
-extension Dictionary : _JSONStringDictionaryEncodableMarker where Key == String, Value: Encodable { }
+extension Dictionary: _JSONCodingKeyRepresentableDictionaryEncodableMarker where Key: CodingKeyRepresentable, Value: Encodable {
+    var _stringKeyedDictionary: [String: Encodable] {
+        // Fast path for String keys - return self without creating a new dictionary
+        if Key.self == String.self {
+            return self as! [String: Encodable]
+        }
+        // Convert other CodingKeyRepresentable keys to String
+        var result = [String: Encodable]()
+        result.reserveCapacity(count)
+        for (key, value) in self {
+            result[key.codingKey.stringValue] = value
+        }
+        return result
+    }
+}
 
 /// 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

Tests/FoundationEssentialsTests/JSONEncoderTests.swift

@@ -470,6 +470,10 @@ private struct JSONEncoderTests {
         let outerValue: EncodeNested
     }
 
+    private struct CodingKeyRepresentableKey: RawRepresentable, CodingKeyRepresentable, Hashable, Codable {
+        let rawValue: String
+    }
+
     @Test func encodingKeyStrategyPath() throws {
         // Make sure a more complex path shows up the way we want
         // Make sure the path reflects keys in the Swift, not the resulting ones in the JSON
@@ -548,6 +552,18 @@ private struct JSONEncoderTests {
 
         #expect(["leave_me_alone": "test"] == result)
     }
+    
+    @Test func decodingDictionaryCodingKeyRepresentableKeyConversionUntouched() throws {
+        // CodingKeyRepresentable dictionary keys should NOT be converted by key decoding strategy,
+        // but nested struct properties should still be converted.
+        let input = "{\"leave_me_alone\":{\"this_is_camel_case\":\"test\"}}".data(using: .utf8)!
+        let decoder = JSONDecoder()
+        decoder.keyDecodingStrategy = .convertFromSnakeCase
+        let result = try decoder.decode([CodingKeyRepresentableKey: DecodeMe3].self, from: input)
+
+        let value = try #require(result[CodingKeyRepresentableKey(rawValue: "leave_me_alone")])
+        #expect(value.thisIsCamelCase == "test")
+    }
 
     @Test func decodingDictionaryFailureKeyPath() {
         let input = "{\"leave_me_alone\":\"test\"}".data(using: .utf8)!
@@ -2333,6 +2349,22 @@ extension JSONEncoderTests {
 
         #expect(expected == resultString)
     }
+    
+    @Test func encodingDictionaryCodingKeyRepresentableKeyConversionUntouched() throws {
+        // CodingKeyRepresentable dictionary keys should NOT be converted by key encoding strategy,
+        // but nested struct properties should still be converted.
+        let expected = "{\"leaveMeAlone\":{\"this_is_camel_case\":\"test\"}}"
+        let toEncode: [CodingKeyRepresentableKey: DecodeMe3] = [
+            CodingKeyRepresentableKey(rawValue: "leaveMeAlone"): DecodeMe3(thisIsCamelCase: "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