[SE-0320] Added protocol CodingKeyRepresentable (#34458)

mortenbekditlevsen · web-flow · commit 0251957ae382 · 2021-10-30T11:33:42.000+01:00
diff --git a/stdlib/public/core/Codable.swift b/stdlib/public/core/Codable.swift
@@ -5510,23 +5510,102 @@ internal struct _DictionaryCodingKey: CodingKey {
   internal let stringValue: String
   internal let intValue: Int?
 
-  internal init?(stringValue: String) {
+  internal init(stringValue: String) {
     self.stringValue = stringValue
     self.intValue = Int(stringValue)
   }
 
-  internal init?(intValue: Int) {
+  internal init(intValue: Int) {
     self.stringValue = "\(intValue)"
     self.intValue = intValue
   }
+
+  fileprivate init(codingKey: CodingKey) {
+    self.stringValue = codingKey.stringValue
+    self.intValue = codingKey.intValue
+  }
+}
+
+/// A type that can be converted to and from a coding key.
+///
+/// With a `CodingKeyRepresentable` type, you can losslessly convert between a
+/// custom type and a `CodingKey` type.
+///
+/// Conforming a type to `CodingKeyRepresentable` lets you opt in to encoding
+/// and decoding `Dictionary` values keyed by the conforming type to and from
+/// a keyed container, rather than encoding and decoding the dictionary as an
+/// unkeyed container of alternating key-value pairs.
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+public protocol CodingKeyRepresentable {
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  var codingKey: CodingKey { get }
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  init?<T: CodingKey>(codingKey: T)
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+extension RawRepresentable where Self: CodingKeyRepresentable, RawValue == String {
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  public var codingKey: CodingKey {
+    _DictionaryCodingKey(stringValue: rawValue)
+  }
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  public init?<T: CodingKey>(codingKey: T) {
+    self.init(rawValue: codingKey.stringValue)
+  }
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+extension RawRepresentable where Self: CodingKeyRepresentable, RawValue == Int {
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  public var codingKey: CodingKey {
+    _DictionaryCodingKey(intValue: rawValue)
+  }
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  public init?<T: CodingKey>(codingKey: T) {
+    if let intValue = codingKey.intValue {
+      self.init(rawValue: intValue)
+    } else {
+      return nil
+    }
+  }
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+extension Int: CodingKeyRepresentable {
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  public var codingKey: CodingKey {
+    _DictionaryCodingKey(intValue: self)
+  }
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  public init?<T: CodingKey>(codingKey: T) {
+    if let intValue = codingKey.intValue {
+      self = intValue
+    } else {
+      return nil
+    }
+  }
+}
+
+@available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+extension String: CodingKeyRepresentable {
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  public var codingKey: CodingKey {
+    _DictionaryCodingKey(stringValue: self)
+  }
+  @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+  public init?<T: CodingKey>(codingKey: T) {
+    self = codingKey.stringValue
+  }
 }
 
 extension Dictionary: Encodable where Key: Encodable, Value: Encodable {
   /// Encodes the contents of this dictionary into the given encoder.
   ///
-  /// If the dictionary uses `String` or `Int` keys, the contents are encoded
-  /// in a keyed container. Otherwise, the contents are encoded as alternating
-  /// key-value pairs in an unkeyed container.
+  /// If the dictionary uses keys that are `String`, `Int`, or a type conforming
+  /// to `CodingKeyRepresentable`, the contents are encoded in a keyed container.
+  /// Otherwise, the contents are encoded as alternating key-value pairs in an
+  /// unkeyed container.
   ///
   /// This function throws an error if any values are invalid for the given
   /// encoder's format.
@@ -5537,16 +5616,26 @@ extension Dictionary: Encodable where Key: Encodable, Value: Encodable {
       // Since the keys are already Strings, we can use them as keys directly.
       var container = encoder.container(keyedBy: _DictionaryCodingKey.self)
       for (key, value) in self {
-        let codingKey = _DictionaryCodingKey(stringValue: key as! String)!
+        let codingKey = _DictionaryCodingKey(stringValue: key as! String)
         try container.encode(value, forKey: codingKey)
       }
     } else if Key.self == Int.self {
       // Since the keys are already Ints, we can use them as keys directly.
       var container = encoder.container(keyedBy: _DictionaryCodingKey.self)
       for (key, value) in self {
-        let codingKey = _DictionaryCodingKey(intValue: key as! Int)!
+        let codingKey = _DictionaryCodingKey(intValue: key as! Int)
         try container.encode(value, forKey: codingKey)
       }
+    } else if #available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *),
+              Key.self is CodingKeyRepresentable.Type {
+      // Since the keys are CodingKeyRepresentable, we can use the `codingKey`
+      // to create `_DictionaryCodingKey` instances.
+      var container = encoder.container(keyedBy: _DictionaryCodingKey.self)
+      for (key, value) in self {
+        let codingKey = (key as! CodingKeyRepresentable).codingKey
+        let dictionaryCodingKey = _DictionaryCodingKey(codingKey: codingKey)
+        try container.encode(value, forKey: dictionaryCodingKey)
+      }
     } else {
       // Keys are Encodable but not Strings or Ints, so we cannot arbitrarily
       // convert to keys. We can encode as an array of alternating key-value
@@ -5601,6 +5690,22 @@ extension Dictionary: Decodable where Key: Decodable, Value: Decodable {
         let value = try container.decode(Value.self, forKey: key)
         self[key.intValue! as! Key] = value
       }
+    } else if #available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *),
+              let keyType = Key.self as? CodingKeyRepresentable.Type {
+      // The keys are CodingKeyRepresentable, so we should be able to expect
+      // a keyed container.
+      let container = try decoder.container(keyedBy: _DictionaryCodingKey.self)
+      for codingKey in container.allKeys {
+        guard let key: Key = keyType.init(codingKey: codingKey) as? Key else {
+          throw DecodingError.dataCorruptedError(
+            forKey: codingKey,
+            in: container,
+            debugDescription: "Could not convert key to type \(Key.self)"
+          )
+        }
+        let value: Value = try container.decode(Value.self, forKey: codingKey)
+        self[key] = value
+      }
     } else {
       // We should have encoded as an array of alternating key-value pairs.
       var container = try decoder.unkeyedContainer()
diff --git a/test/api-digester/Outputs/cake-abi.json b/test/api-digester/Outputs/cake-abi.json
@@ -1873,6 +1873,14 @@
           "usr": "s:Se",
           "mangledName": "$sSe"
         },
+        {
+          "kind": "Conformance",
+          "name": "CodingKeyRepresentable",
+          "printedName": "CodingKeyRepresentable",
+          "usr": "s:s22CodingKeyRepresentableP",
+          "mangledName": "$ss22CodingKeyRepresentableP",
+          "isABIPlaceholder": true
+        },
         {
           "kind": "Conformance",
           "name": "CustomReflectable",
diff --git a/test/api-digester/Outputs/cake.json b/test/api-digester/Outputs/cake.json
@@ -1740,6 +1740,14 @@
           "usr": "s:Se",
           "mangledName": "$sSe"
         },
+        {
+          "kind": "Conformance",
+          "name": "CodingKeyRepresentable",
+          "printedName": "CodingKeyRepresentable",
+          "usr": "s:s22CodingKeyRepresentableP",
+          "mangledName": "$ss22CodingKeyRepresentableP",
+          "isABIPlaceholder": true
+        },
         {
           "kind": "Conformance",
           "name": "CustomReflectable",
diff --git a/test/stdlib/CodableTests.swift b/test/stdlib/CodableTests.swift
@@ -76,14 +76,23 @@ func expectRoundTripEquality<T : Codable>(of value: T, encode: (T) throws -> Dat
     expectEqual(value, decoded, "\(#file):\(lineNumber): Decoded \(T.self) <\(debugDescription(decoded))> not equal to original <\(debugDescription(value))>")
 }
 
-func expectRoundTripEqualityThroughJSON<T : Codable>(for value: T, lineNumber: Int) where T : Equatable {
+func expectRoundTripEqualityThroughJSON<T : Codable>(for value: T, expectedJSON: String? = nil, lineNumber: Int) where T : Equatable {
     let inf = "INF", negInf = "-INF", nan = "NaN"
     let encode = { (_ value: T) throws -> Data in
         let encoder = JSONEncoder()
         encoder.nonConformingFloatEncodingStrategy = .convertToString(positiveInfinity: inf,
                                                                       negativeInfinity: negInf,
                                                                       nan: nan)
-        return try encoder.encode(value)
+        if #available(macOS 10.13, iOS 11.0, watchOS 4.0, tvOS 11.0, *) {
+            encoder.outputFormatting = .sortedKeys
+        }
+        let encoded = try encoder.encode(value)
+
+        if let expectedJSON = expectedJSON {
+            let actualJSON = String(decoding: encoded, as: UTF8.self)
+            expectEqual(expectedJSON, actualJSON, line: UInt(lineNumber))
+        }
+        return encoded
     }
 
     let decode = { (_ type: T.Type, _ data: Data) throws -> T in
@@ -111,13 +120,22 @@ func expectRoundTripEqualityThroughPlist<T : Codable>(for value: T, lineNumber:
 
 // MARK: - Helper Types
 // A wrapper around a UUID that will allow it to be encoded at the top level of an encoder.
-struct UUIDCodingWrapper : Codable, Equatable {
+struct UUIDCodingWrapper : Codable, Equatable, Hashable, CodingKeyRepresentable {
     let value: UUID
 
     init(_ value: UUID) {
         self.value = value
     }
 
+    init?<T: CodingKey>(codingKey: T) {
+        guard let uuid = UUID(uuidString: codingKey.stringValue) else { return nil }
+        self.value = uuid
+    }
+
+    var codingKey: CodingKey {
+        GenericCodingKey(stringValue: value.uuidString)
+    }
+
     static func ==(_ lhs: UUIDCodingWrapper, _ rhs: UUIDCodingWrapper) -> Bool {
         return lhs.value == rhs.value
     }
@@ -462,6 +480,91 @@ class TestCodable : TestCodableSuper {
         }
     }
 
+    @available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *)
+    func test_Dictionary_JSON() {
+        enum X: String, Codable { case a, b }
+        enum Y: String, Codable, CodingKeyRepresentable { case a, b }
+        enum Z: Codable, CodingKeyRepresentable {
+            case a
+            case b
+            init?<T: CodingKey>(codingKey: T) {
+                switch codingKey.stringValue {
+                case "α":
+                    self = .a
+                case "β":
+                    self = .b
+                default:
+                    return nil
+                }
+            }
+
+            var codingKey: CodingKey {
+                GenericCodingKey(stringValue: encoded)
+            }
+
+            var encoded: String {
+                switch self {
+                case .a: return "α"
+                case .b: return "β"
+                }
+            }
+        }
+        enum S: Character, Codable, CodingKeyRepresentable {
+            case a = "a"
+            case b = "b"
+
+            init?<T: CodingKey>(codingKey: T) {
+                guard codingKey.stringValue.count == 1 else { return nil }
+                self.init(rawValue: codingKey.stringValue.first!)
+            }
+
+            var codingKey: CodingKey {
+                GenericCodingKey(stringValue: "\(self.rawValue)")
+            }
+        }
+
+        enum U: Int, Codable { case a = 0, b}
+        enum V: Int, Codable, CodingKeyRepresentable { case a = 0, b }
+        enum W: Codable, CodingKeyRepresentable {
+            case a
+            case b
+            init?<T: CodingKey>(codingKey: T) {
+                guard let intValue = codingKey.intValue else { return nil }
+                switch intValue {
+                case 42:
+                    self = .a
+                case 64:
+                    self = .b
+                default:
+                    return nil
+                }
+            }
+            var codingKey: CodingKey {
+                GenericCodingKey(intValue: self.encoded)
+            }
+            var encoded: Int {
+                switch self {
+                case .a: return 42
+                case .b: return 64
+                }
+            }
+        }
+
+        let uuid = UUID(uuidString: "E621E1F8-C36C-495A-93FC-0C247A3E6E5F")!
+        let uuidWrapper = UUIDCodingWrapper(uuid)
+
+        expectRoundTripEqualityThroughJSON(for: [X.a: true],             expectedJSON: #"["a",true]"#,                                    lineNumber: #line)
+        expectRoundTripEqualityThroughJSON(for: [Y.a: true, Y.b: false], expectedJSON: #"{"a":true,"b":false}"#,                          lineNumber: #line)
+        expectRoundTripEqualityThroughJSON(for: [Z.a: true, Z.b: false], expectedJSON: #"{"α":true,"β":false}"#,                          lineNumber: #line)
+        expectRoundTripEqualityThroughJSON(for: [S.a: true, S.b: false], expectedJSON: #"{"a":true,"b":false}"#,                          lineNumber: #line)
+        expectRoundTripEqualityThroughJSON(for: [uuidWrapper: true],     expectedJSON: #"{"E621E1F8-C36C-495A-93FC-0C247A3E6E5F":true}"#, lineNumber: #line)
+        expectRoundTripEqualityThroughJSON(for: [uuid: true],            expectedJSON: #"["E621E1F8-C36C-495A-93FC-0C247A3E6E5F",true]"#, lineNumber: #line)
+        expectRoundTripEqualityThroughJSON(for: [U.a: true],             expectedJSON: #"[0,true]"#,                                      lineNumber: #line)
+        expectRoundTripEqualityThroughJSON(for: [V.a: true, V.b: false], expectedJSON: #"{"0":true,"1":false}"#,                          lineNumber: #line)
+        expectRoundTripEqualityThroughJSON(for: [W.a: true, W.b: false], expectedJSON: #"{"42":true,"64":false}"#,                        lineNumber: #line)
+    }
+
+
     // MARK: - IndexPath
     lazy var indexPathValues: [Int : IndexPath] = [
         #line : IndexPath(), // empty
@@ -853,6 +956,20 @@ class TestCodable : TestCodableSuper {
 
 // MARK: - Helper Types
 
+struct GenericCodingKey: CodingKey {
+    var stringValue: String
+    var intValue: Int?
+
+    init(stringValue: String) {
+        self.stringValue = stringValue
+    }
+
+    init(intValue: Int) {
+        self.stringValue = "\(intValue)"
+        self.intValue = intValue
+    }
+}
+
 struct TopLevelWrapper<T> : Codable, Equatable where T : Codable, T : Equatable {
     let value: T
 
@@ -937,6 +1054,10 @@ if #available(macOS 10.12, iOS 10.0, watchOS 3.0, tvOS 10.0, *) {
     tests["test_URLComponents_Plist"] = TestCodable.test_URLComponents_Plist
 }
 
+if #available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *) {
+    tests["test_Dictionary_JSON"] = TestCodable.test_Dictionary_JSON
+}
+
 var CodableTests = TestSuite("TestCodable")
 for (name, test) in tests {
     CodableTests.test(name) { test(TestCodable())() }
