Flight-School · meeloo · Sep 22, 2021
diff --git a/Sources/MessagePack/Decoder/MessagePackDecoder.swift b/Sources/MessagePack/Decoder/MessagePackDecoder.swift
@@ -78,8 +78,9 @@ final class _MessagePackDecoder {
     var codingPath: [CodingKey] = []
 
     var userInfo: [CodingUserInfoKey : Any] = [:]
-
-    var container: MessagePackDecodingContainer?
+
+    var containers: [MessagePackDecodingContainer] = []
+    var container: MessagePackDecodingContainer? { containers.last }
     fileprivate var data: Data
 
     init(data: Data) {
@@ -89,14 +90,14 @@ final class _MessagePackDecoder {
 
 extension _MessagePackDecoder: Decoder {
     fileprivate func assertCanCreateContainer() {
-        precondition(self.container == nil)
+//        precondition(self.container == nil)
     }
 
     func container<Key>(keyedBy type: Key.Type) -> KeyedDecodingContainer<Key> where Key : CodingKey {
         assertCanCreateContainer()
 
         let container = KeyedContainer<Key>(data: self.data, codingPath: self.codingPath, userInfo: self.userInfo)
-        self.container = container
+        self.containers.append(container)
 
         return KeyedDecodingContainer(container)
     }
@@ -105,7 +106,7 @@ extension _MessagePackDecoder: Decoder {
         assertCanCreateContainer()
 
         let container = UnkeyedContainer(data: self.data, codingPath: self.codingPath, userInfo: self.userInfo)
-        self.container = container
+        self.containers.append(container)
 
         return container
     }
@@ -114,7 +115,7 @@ extension _MessagePackDecoder: Decoder {
         assertCanCreateContainer()
 
         let container = SingleValueContainer(data: self.data, codingPath: self.codingPath, userInfo: self.userInfo)
-        self.container = container
+        self.containers.append(container)
 
         return container
     }

diff --git a/Sources/MessagePack/Encoder/KeyedEncodingContainer.swift b/Sources/MessagePack/Encoder/KeyedEncodingContainer.swift
@@ -1,19 +1,41 @@
 import Foundation
 
+protocol KeyedStorage {
+    var keyStorage: KeyStorage { get }
+}
+
+class KeyStorage {
+    var values: [AnyCodingKey: _MessagePackEncodingContainer] = [:]
+
+    var codingPath: [CodingKey] = []
+    var userInfo: [CodingUserInfoKey: Any] = [:]
+}
+
 extension _MessagePackEncoder {
-    final class KeyedContainer<Key> where Key: CodingKey {
-        private var storage: [AnyCodingKey: _MessagePackEncodingContainer] = [:]
-
-        var codingPath: [CodingKey]
-        var userInfo: [CodingUserInfoKey: Any]
-
+    final class KeyedContainer<Key>: KeyedStorage where Key: CodingKey {
+        internal var keyStorage = KeyStorage()
+        var codingPath: [CodingKey] {
+            get {
+                keyStorage.codingPath
+            }
+
+            set {
+                keyStorage.codingPath = newValue
+            }
+        }
+
         func nestedCodingPath(forKey key: CodingKey) -> [CodingKey] {
-            return self.codingPath + [key]
+            return self.keyStorage.codingPath + [key]
         }
+
 
         init(codingPath: [CodingKey], userInfo: [CodingUserInfoKey : Any]) {
-            self.codingPath = codingPath
-            self.userInfo = userInfo
+            self.keyStorage.codingPath = codingPath
+            self.keyStorage.userInfo = userInfo
+        }
+
+        init(keyStorage: KeyStorage) {
+            self.keyStorage = keyStorage
         }
     }
 }
@@ -30,21 +52,21 @@ extension _MessagePackEncoder.KeyedContainer: KeyedEncodingContainerProtocol {
     }
 
     private func nestedSingleValueContainer(forKey key: Key) -> SingleValueEncodingContainer {
-        let container = _MessagePackEncoder.SingleValueContainer(codingPath: self.nestedCodingPath(forKey: key), userInfo: self.userInfo)
-        self.storage[AnyCodingKey(key)] = container
+        let container = _MessagePackEncoder.SingleValueContainer(codingPath: self.nestedCodingPath(forKey: key), userInfo: self.keyStorage.userInfo)
+        self.keyStorage.values[AnyCodingKey(key)] = container
         return container
     }
 
     func nestedUnkeyedContainer(forKey key: Key) -> UnkeyedEncodingContainer {
-        let container = _MessagePackEncoder.UnkeyedContainer(codingPath: self.nestedCodingPath(forKey: key), userInfo: self.userInfo)
-        self.storage[AnyCodingKey(key)] = container
+        let container = _MessagePackEncoder.UnkeyedContainer(codingPath: self.nestedCodingPath(forKey: key), userInfo: self.keyStorage.userInfo)
+        self.keyStorage.values[AnyCodingKey(key)] = container
 
         return container
     }
 
     func nestedContainer<NestedKey>(keyedBy keyType: NestedKey.Type, forKey key: Key) -> KeyedEncodingContainer<NestedKey> where NestedKey : CodingKey {
-        let container = _MessagePackEncoder.KeyedContainer<NestedKey>(codingPath: self.nestedCodingPath(forKey: key), userInfo: self.userInfo)
-        self.storage[AnyCodingKey(key)] = container
+        let container = _MessagePackEncoder.KeyedContainer<NestedKey>(codingPath: self.nestedCodingPath(forKey: key), userInfo: self.keyStorage.userInfo)
+        self.keyStorage.values[AnyCodingKey(key)] = container
 
         return KeyedEncodingContainer(container)
     }
@@ -62,7 +84,7 @@ extension _MessagePackEncoder.KeyedContainer: _MessagePackEncodingContainer {
     var data: Data {
         var data = Data()
 
-        let length = storage.count
+        let length = keyStorage.values.count
         if let uint16 = UInt16(exactly: length) {
             if length <= 15 {
                 data.append(0x80 + UInt8(length))
@@ -77,8 +99,8 @@ extension _MessagePackEncoder.KeyedContainer: _MessagePackEncodingContainer {
             fatalError()
         }
 
-        for (key, container) in self.storage {
-            let keyContainer = _MessagePackEncoder.SingleValueContainer(codingPath: self.codingPath, userInfo: self.userInfo)
+        for (key, container) in self.keyStorage.values {
+            let keyContainer = _MessagePackEncoder.SingleValueContainer(codingPath: self.codingPath, userInfo: self.keyStorage.userInfo)
             try! keyContainer.encode(key.stringValue)
             data.append(keyContainer.data)
 

diff --git a/Sources/MessagePack/Encoder/MessagePackEncoder.swift b/Sources/MessagePack/Encoder/MessagePackEncoder.swift
@@ -49,7 +49,7 @@ extension MessagePackEncoder: TopLevelEncoder {
 
 // MARK: -
 
-protocol _MessagePackEncodingContainer {
+protocol _MessagePackEncodingContainer: KeyedStorage {
     var data: Data { get }
 }
 
@@ -71,12 +71,20 @@ extension _MessagePackEncoder: Encoder {
     }
 
     func container<Key>(keyedBy type: Key.Type) -> KeyedEncodingContainer<Key> where Key : CodingKey {
-        assertCanCreateContainer()
-
-        let container = KeyedContainer<Key>(codingPath: self.codingPath, userInfo: self.userInfo)
-        self.container = container
-
-        return KeyedEncodingContainer(container)
+//        assertCanCreateContainer()
+
+        guard let container = container else {
+            let container = KeyedContainer<Key>(codingPath: self.codingPath, userInfo: self.userInfo)
+            self.container = container
+
+            return KeyedEncodingContainer(container)
+        }
+
+        let newContainer = KeyedContainer<Key>(keyStorage: container.keyStorage)
+        self.container = newContainer
+
+
+        return KeyedEncodingContainer(newContainer)
     }
 
     func unkeyedContainer() -> UnkeyedEncodingContainer {

diff --git a/Sources/MessagePack/Encoder/SingleValueEncodingContainer.swift b/Sources/MessagePack/Encoder/SingleValueEncodingContainer.swift
@@ -3,21 +3,28 @@ import Foundation
 extension _MessagePackEncoder {
     final class SingleValueContainer {
         private var storage: Data = Data()
-
+        var keyStorage: KeyStorage = KeyStorage()
+        var codingPath: [CodingKey] {
+            get {
+                keyStorage.codingPath
+            }
+
+            set {
+                keyStorage.codingPath = newValue
+            }
+        }
+
         fileprivate var canEncodeNewValue = true
         fileprivate func checkCanEncode(value: Any?) throws {
             guard self.canEncodeNewValue else {
-                let context = EncodingError.Context(codingPath: self.codingPath, debugDescription: "Attempt to encode value through single value container when previously value already encoded.")
+                let context = EncodingError.Context(codingPath: self.keyStorage.codingPath, debugDescription: "Attempt to encode value through single value container when previously value already encoded.")
                 throw EncodingError.invalidValue(value as Any, context)
             }
         }
 
-        var codingPath: [CodingKey]
-        var userInfo: [CodingUserInfoKey: Any]
-
         init(codingPath: [CodingKey], userInfo: [CodingUserInfoKey : Any]) {
-            self.codingPath = codingPath
-            self.userInfo = userInfo
+            self.keyStorage.codingPath = codingPath
+            self.keyStorage.userInfo = userInfo
         }
     }
 }

diff --git a/Sources/MessagePack/Encoder/UnkeyedEncodingContainer.swift b/Sources/MessagePack/Encoder/UnkeyedEncodingContainer.swift
@@ -3,22 +3,28 @@ import Foundation
 extension _MessagePackEncoder {
     final class UnkeyedContainer {
         private var storage: [_MessagePackEncodingContainer] = []
-
+        internal var keyStorage = KeyStorage()
+        var codingPath: [CodingKey] {
+            get {
+                keyStorage.codingPath
+            }
+
+            set {
+                keyStorage.codingPath = newValue
+            }
+        }
+
         var count: Int {
             return storage.count
         }
 
-        var codingPath: [CodingKey]
-
         var nestedCodingPath: [CodingKey] {
             return self.codingPath + [AnyCodingKey(intValue: self.count)!]
         }
 
-        var userInfo: [CodingUserInfoKey: Any]
-
         init(codingPath: [CodingKey], userInfo: [CodingUserInfoKey : Any]) {
-            self.codingPath = codingPath
-            self.userInfo = userInfo
+            self.keyStorage.codingPath = codingPath
+            self.keyStorage.userInfo = userInfo
         }
     }
 }
@@ -35,21 +41,21 @@ extension _MessagePackEncoder.UnkeyedContainer: UnkeyedEncodingContainer {
     }
 
     private func nestedSingleValueContainer() -> SingleValueEncodingContainer {
-        let container = _MessagePackEncoder.SingleValueContainer(codingPath: self.nestedCodingPath, userInfo: self.userInfo)
+        let container = _MessagePackEncoder.SingleValueContainer(codingPath: self.nestedCodingPath, userInfo: self.keyStorage.userInfo)
         self.storage.append(container)
 
         return container
     }
 
     func nestedContainer<NestedKey>(keyedBy keyType: NestedKey.Type) -> KeyedEncodingContainer<NestedKey> where NestedKey : CodingKey {
-        let container = _MessagePackEncoder.KeyedContainer<NestedKey>(codingPath: self.nestedCodingPath, userInfo: self.userInfo)
+        let container = _MessagePackEncoder.KeyedContainer<NestedKey>(codingPath: self.nestedCodingPath, userInfo: self.keyStorage.userInfo)
         self.storage.append(container)
 
         return KeyedEncodingContainer(container)
     }
 
     func nestedUnkeyedContainer() -> UnkeyedEncodingContainer {
-        let container = _MessagePackEncoder.UnkeyedContainer(codingPath: self.nestedCodingPath, userInfo: self.userInfo)
+        let container = _MessagePackEncoder.UnkeyedContainer(codingPath: self.nestedCodingPath, userInfo: self.keyStorage.userInfo)
         self.storage.append(container)
 
         return container