From dea607c4a7f8963b87797f0346ed14e486690ef8 Mon Sep 17 00:00:00 2001
From: Michel Fortin <michel.fortin@michelf.com>
Date: Thu, 20 Feb 2020 11:54:59 -0500
Subject: [PATCH] Added allowIntKey mode to encode / decode keys as integers in
 dictionaries when the CodingKey implementation provides an intValue.

---
 Sources/MessagePackDecoder.swift              | 42 +++++++++++++++----
 Sources/MessagePackEncoder.swift              | 23 +++++++---
 .../MessagePackType+MapType.swift             | 18 ++++++--
 Tests/MessagePackerTests/MapPackedTests.swift | 22 ++++++++++
 .../MessagePackerTests/MapUnpackedTests.swift | 27 ++++++++++++
 5 files changed, 116 insertions(+), 16 deletions(-)

diff --git a/Sources/MessagePackDecoder.swift b/Sources/MessagePackDecoder.swift
index d38d4d9..15d320c 100644
--- a/Sources/MessagePackDecoder.swift
+++ b/Sources/MessagePackDecoder.swift
@@ -9,15 +9,19 @@
 import Foundation
 
 open class MessagePackDecoder: Decoder {
+    public var allowIntKeys: Bool
     public var codingPath: [CodingKey] = []
     public var userInfo: [CodingUserInfoKey : Any] = [:]
     private var storage = MessagePackStorage()
 
-    public init() {}
+    public init(allowIntKeys: Bool = false) {
+        self.allowIntKeys = allowIntKeys
+    }
 
-    public init(referencing: Data, codingPath: [CodingKey] = []) {
+    public init(referencing: Data, codingPath: [CodingKey] = [], allowIntKeys: Bool = false) {
         storage.push(container: referencing)
         self.codingPath = codingPath
+        self.allowIntKeys = allowIntKeys
     }
 
     public func container<Key: CodingKey>(keyedBy type: Key.Type) throws -> KeyedDecodingContainer<Key> {
@@ -123,25 +127,45 @@ private extension MessagePackDecoder {
 }
 
 extension MessagePackDecoder {
+    private func convert<Key: CodingKey>(packKey: MessagePackType.MapType.Key, to: Key.Type) -> Key? {
+        switch packKey {
+        case .int(let intValue):
+            if allowIntKeys {
+                return Key(intValue: intValue)
+            } else {
+                return Key(stringValue: String(intValue))
+            }
+        case .string(let stringValue):
+            return Key(stringValue: stringValue)
+        }
+    }
+    private func convertToPackKey(from codingKey: CodingKey) -> MessagePackType.MapType.Key {
+        if allowIntKeys, let intValue = codingKey.intValue {
+            return .int(intValue)
+        } else {
+            return .string(codingKey.stringValue)
+        }
+    }
+
     struct KeyedContainer<Key: CodingKey>: KeyedDecodingContainerProtocol {
         private let decoder: MessagePackDecoder
         private(set) var codingPath: [CodingKey]
         private(set) var allKeys: [Key]
-        private var container: [String: Data] = [:]
+        private var container: [MessagePackType.MapType.Key: Data] = [:]
 
-        init(referencing decoder: MessagePackDecoder, container: [String: Data]) {
+        init(referencing decoder: MessagePackDecoder, container: [MessagePackType.MapType.Key: Data]) {
             self.decoder = decoder
             self.codingPath = decoder.codingPath
-            self.allKeys = container.keys.compactMap { Key(stringValue: $0) }
+            self.allKeys = container.keys.compactMap { decoder.convert(packKey: $0, to: Key.self) }
             self.container = container
         }
 
         func contains(_ key: Key) -> Bool {
-            return container[key.stringValue] != nil
+            return container[decoder.convertToPackKey(from: key)] != nil
         }
 
         func findEntry(by key: CodingKey) throws -> Data {
-            guard let entry = self.container[key.stringValue] else {
+            guard let entry = self.container[decoder.convertToPackKey(from: key)] else {
                 let context = DecodingError.Context(codingPath: self.decoder.codingPath, debugDescription: "No value associated with key \(key) (\"\(key.stringValue)\").")
                 throw DecodingError.keyNotFound(key, context)
             }
@@ -542,3 +566,7 @@ extension MessagePackDecoder {
         }
     }
 }
+
+
+
+
diff --git a/Sources/MessagePackEncoder.swift b/Sources/MessagePackEncoder.swift
index 84a5810..c7377de 100644
--- a/Sources/MessagePackEncoder.swift
+++ b/Sources/MessagePackEncoder.swift
@@ -9,11 +9,14 @@
 import Foundation
 
 open class MessagePackEncoder: Encoder {
+    public var allowIntKeys: Bool
     public var codingPath: [CodingKey] = []
     public var userInfo: [CodingUserInfoKey : Any] = [:]
     fileprivate var storage = MessagePackStorage()
 
-    public init() {}
+    public init(allowIntKeys: Bool = false) {
+        self.allowIntKeys = allowIntKeys
+    }
 
     public func container<Key>(keyedBy type: Key.Type) -> KeyedEncodingContainer<Key> where Key : CodingKey {
         return KeyedEncodingContainer(KeyedContainer<Key>(referencing: self, codingPath: codingPath) { [weak self] in
@@ -109,8 +112,13 @@ extension MessagePackEncoder {
         }
 
         fileprivate func add(_ value: Data, forKey key: CodingKey) {
-            let key = encoder.boxMessagePack(key.stringValue)
-            self.packedData.append(key)
+            let keyData: Data
+            if encoder.allowIntKeys, let intValue = key.intValue {
+                keyData = encoder.boxMessagePack(intValue)
+            } else {
+                keyData = encoder.boxMessagePack(key.stringValue)
+            }
+            self.packedData.append(keyData)
             self.packedData.append(value)
             count += 1
         }
@@ -420,11 +428,11 @@ extension MessagePackEncoder {
         }
     }
 
-    class MessagePackReferencingKeyedEncoder<Key: CodingKey>: MessagePackEncoder {
+    class MessagePackReferencingKeyedEncoder<Key: CodingKey, Key2: CodingKey>: MessagePackEncoder {
         private let container: KeyedContainer<Key>
-        private let key: CodingKey
+        private let key: Key2
 
-        init(container: KeyedContainer<Key>, key: CodingKey) {
+        init(container: KeyedContainer<Key>, key: Key2) {
             self.container = container
             self.key = key
             super.init()
@@ -450,3 +458,6 @@ extension MessagePackEncoder {
         }
     }
 }
+
+
+
diff --git a/Sources/MessagePackTypes/MessagePackType+MapType.swift b/Sources/MessagePackTypes/MessagePackType+MapType.swift
index 49e946b..fdc1a69 100644
--- a/Sources/MessagePackTypes/MessagePackType+MapType.swift
+++ b/Sources/MessagePackTypes/MessagePackType+MapType.swift
@@ -102,7 +102,12 @@ extension MessagePackType.MapType {
 }
 
 extension MessagePackType.MapType {
-    static func split(for value: Data) throws -> [String : Data] {
+    enum Key: Hashable {
+        case string(String)
+        case int(Int)
+    }
+
+    static func split(for value: Data) throws -> [Key : Data] {
         guard let firstByte = value.first else { throw MessagePackError.emptyData }
 
         let type = try MessagePackType.MapType(firstByte)
@@ -110,10 +115,16 @@ extension MessagePackType.MapType {
         let length = try type.length(value)
 
         return try (0..<length)
-            .reduce(into: (dictionary: [String : Data](), index: startIndex)) { args, _ in
+            .reduce(into: (dictionary: [Key : Data](), index: startIndex)) { args, _ in
                 let key = try value.subdata(startIndex: args.index).firstMessagePackeValue()
                 let value = try value.subdata(startIndex: args.index + key.count).firstMessagePackeValue()
-                args.dictionary[try String.unpack(for: key)] = value
+                let keyType = try MessagePackType(key.first!)
+                switch keyType {
+                case .signedInteger, .unsignedInteger:
+                    args.dictionary[.int(try Int.unpack(for: key))] = value
+                default:
+                    args.dictionary[.string(try String.unpack(for: key))] = value
+                }
                 args.index += (key.count + value.count)
             }
             .dictionary
@@ -139,3 +150,4 @@ extension MessagePackType.MapType {
         }
     }
 }
+
diff --git a/Tests/MessagePackerTests/MapPackedTests.swift b/Tests/MessagePackerTests/MapPackedTests.swift
index 8621aad..8443e7d 100644
--- a/Tests/MessagePackerTests/MapPackedTests.swift
+++ b/Tests/MessagePackerTests/MapPackedTests.swift
@@ -44,6 +44,26 @@ class MapPackedTests: XCTestCase {
         XCTAssertEqual((try encoder.encode(input)).count, output)
     }
 
+    func testMapAllowIntKey() {
+        let input = [8: 1.1]
+        let output = Data([129, 8, 203, 63, 241, 153, 153, 153, 153, 153, 154])
+        let encoder = MessagePackEncoder(allowIntKeys: true)
+        XCTAssertEqual(try encoder.encode(input), output)
+    }
+
+    func testMapAllowIntKeyWithString() {
+        let input = ["A": 9]
+        let output = Data([129, 161, 65, 9])
+        let encoder = MessagePackEncoder(allowIntKeys: true)
+        XCTAssertEqual(try encoder.encode(input), output)
+    }
+
+    func testMapAllowIntKeyWithIntegerInString() {
+        let input = ["1": 9]
+        let output = Data([129, 1, 9])
+        let encoder = MessagePackEncoder(allowIntKeys: true)
+        XCTAssertEqual(try encoder.encode(input), output)
+    }
 
     func test2DimensionalMap() {
         let input = ["a": ["b": 2]]
@@ -51,3 +71,5 @@ class MapPackedTests: XCTestCase {
         XCTAssertEqual(try encoder.encode(input), output)
     }
 }
+
+
diff --git a/Tests/MessagePackerTests/MapUnpackedTests.swift b/Tests/MessagePackerTests/MapUnpackedTests.swift
index d3adb15..c9acbf9 100644
--- a/Tests/MessagePackerTests/MapUnpackedTests.swift
+++ b/Tests/MessagePackerTests/MapUnpackedTests.swift
@@ -39,6 +39,32 @@ class MapUnpackedTests: XCTestCase {
         }
     }
 
+    func testAllowIntKey() {
+        let input = Data([130, 2, 203, 64, 1, 153, 153, 153, 153, 153, 154, 1, 203, 63, 241, 153, 153, 153, 153, 153, 154])
+        let output: [Int: Double] = [1: 1.1, 2: 2.2]
+
+        do {
+            let decoder = MessagePackDecoder(allowIntKeys: true)
+            let result: [Int: Double] = try decoder.decode(Dictionary.self, from: input)
+            result.forEach { XCTAssertEqual(output[$0.key], $0.value) }
+        } catch {
+            XCTFail(error.localizedDescription)
+        }
+    }
+
+    func testAllowIntKeyWithString() {
+        let input = Data([130, 161, 65, 9, 1, 9])
+        let output: [String: Int] = ["1": 9, "A": 9]
+
+        do {
+            let decoder = MessagePackDecoder(allowIntKeys: true)
+            let result: [String: Int] = try decoder.decode(Dictionary.self, from: input)
+            result.forEach { XCTAssertEqual(output[$0.key], $0.value) }
+        } catch {
+            XCTFail(error.localizedDescription)
+        }
+    }
+
     func test2DimensionalMap() {
         let input = Data([129, 161, 97, 129, 161, 98, 2])
         let output = ["a": ["b": 2]]
@@ -46,3 +72,4 @@ class MapUnpackedTests: XCTestCase {
         XCTAssertEqual(try decoder.decode(Dictionary.self, from: input), output)
     }
 }
+