Fix strictDecoding by only propagating not-decoded keys when decoding is successful

Sources/TOMLKit/Decoder/InternalTOMLDecoder.swift

@@ -12,6 +12,7 @@ final class InternalTOMLDecoder: Decoder {
 	var dataDecoder: (TOMLValueConvertible) -> Data?
 	var strictDecoding: Bool = false
 	var notDecodedKeys: NotDecodedKeys
+	let originalNotDecodedKeys: [String: [CodingKey]]
 
 	let tomlValue: TOMLValue
 	init(
@@ -28,6 +29,7 @@ final class InternalTOMLDecoder: Decoder {
 		self.dataDecoder = dataDecoder
 		self.strictDecoding = strictDecoding
 		self.notDecodedKeys = notDecodedKeys
+		self.originalNotDecodedKeys = notDecodedKeys.keys
 	}
 
 	func container<Key>(keyedBy type: Key.Type) throws -> KeyedDecodingContainer<Key> where Key: CodingKey {
@@ -42,6 +44,10 @@ final class InternalTOMLDecoder: Decoder {
 			)
 		}
 
+		// Assume that any previous container creation was related to a failed decoding
+		// attempt, and reset the not-decoded keys back to the value that the parent
+		// decoder passed to us.
+		self.notDecodedKeys.keys = self.originalNotDecodedKeys
 		return KeyedDecodingContainer<Key>(
 			KDC(
 				table: table,
@@ -64,6 +70,11 @@ final class InternalTOMLDecoder: Decoder {
 				)
 			)
 		}
+
+		// Assume that any previous container creation was related to a failed decoding
+		// attempt, and reset the not-decoded keys back to the value that the parent
+		// decoder passed to us.
+		self.notDecodedKeys.keys = self.originalNotDecodedKeys
 		return UDC(
 			array,
 			codingPath: self.codingPath,
@@ -75,7 +86,11 @@ final class InternalTOMLDecoder: Decoder {
 	}
 
 	func singleValueContainer() throws -> SingleValueDecodingContainer {
-		SVDC(
+		// Assume that any previous container creation was related to a failed decoding
+		// attempt, and reset the not-decoded keys back to the value that the parent
+		// decoder passed to us.
+		self.notDecodedKeys.keys = self.originalNotDecodedKeys
+		return SVDC(
 			self.tomlValue,
 			codingPath: self.codingPath,
 			dataDecoder: self.dataDecoder,

Sources/TOMLKit/Decoder/KeyedDecodingContainerProtocol.swift

@@ -228,13 +228,24 @@ extension InternalTOMLDecoder.KDC {
 						codingPath: self.codingPath + key,
 						dataDecoder: self.dataDecoder,
 						strictDecoding: self.strictDecoding,
-						notDecodedKeys: self.notDecodedKeys
+						notDecodedKeys: InternalTOMLDecoder.NotDecodedKeys()
 					)
 
 					self.decodedKeys.append(key.stringValue)
 
-					return try T(from: decoder)
+					let decodedValue = try T(from: decoder)
+
+					// Only propagate not-decoded keys if the decoding was successful.
+					// Otherwise `Decodable` implementations that attempt multiple
+					// decoding strategies in succession (trying the next if the
+					// previous one failed), don't work.
+					for (key, path) in decoder.notDecodedKeys.keys {
+						self.notDecodedKeys.keys[key] = path
+					}
+
+					return decodedValue
 				}
+
 		}
 	}
 

Sources/TOMLKit/Decoder/SingleValueDecodingContainer.swift

@@ -213,20 +213,31 @@ extension InternalTOMLDecoder.SVDC {
 			codingPath: self.codingPath,
 			dataDecoder: self.dataDecoder,
 			strictDecoding: self.strictDecoding,
-			notDecodedKeys: self.notDecodedKeys
+			notDecodedKeys: InternalTOMLDecoder.NotDecodedKeys()
 		)
 
+		let decodedValue: T
 		if type is Data.Type {
 			if let data = self.dataDecoder(self.tomlValue) {
-				return data as! T
+				decodedValue = data as! T
 			} else {
 				throw DecodingError.dataCorrupted(DecodingError.Context(
 					codingPath: self.codingPath,
 					debugDescription: "Unable to decode Data from: \"\(self.tomlValue.debugDescription)\"."
 				))
 			}
 		} else {
-			return try T(from: decoder)
+			decodedValue = try T(from: decoder)
 		}
+
+		// Only propagate not-decoded keys if the decoding was successful.
+		// Otherwise `Decodable` implementations that attempt multiple
+		// decoding strategies in succession (trying the next if the
+		// previous one failed), don't work.
+		for (key, path) in decoder.notDecodedKeys.keys {
+			self.notDecodedKeys.keys[key] = path
+		}
+
+		return decodedValue
 	}
 }

Sources/TOMLKit/Decoder/UnkeyedDecodingContainer.swift

@@ -166,11 +166,21 @@ extension InternalTOMLDecoder.UDC {
 				codingPath: self.codingPath + TOMLCodingKey(index: self.currentIndex),
 				dataDecoder: self.dataDecoder,
 				strictDecoding: self.strictDecoding,
-				notDecodedKeys: self.notDecodedKeys
+				notDecodedKeys: InternalTOMLDecoder.NotDecodedKeys()
 			)
-			let decodable = try T(from: decoder)
+
+			let decodedValue = try T(from: decoder)
 			self.currentIndex += 1
-			return decodable
+
+			// Only propagate not-decoded keys if the decoding was successful.
+			// Otherwise `Decodable` implementations that attempt multiple
+			// decoding strategies in succession (trying the next if the
+			// previous one failed), don't work.
+			for (key, path) in decoder.notDecodedKeys.keys {
+				self.notDecodedKeys.keys[key] = path
+			}
+
+			return decodedValue
 		}
 	}
 

Tests/TOMLKitTests/TOMLKitTests.swift

@@ -19,15 +19,6 @@ enum CodableEnum: String, Codable, Equatable {
 	case abc
 	case def
 	case ghi
-
-	public static func == (lhs: Self, rhs: Self) -> Bool {
-		switch (lhs, rhs) {
-			case (.abc, .abc): return true
-			case (.def, .def): return true
-			case (.ghi, .ghi): return true
-			default: return false
-		}
-	}
 }
 
 struct CodableStruct: Codable, Equatable {
@@ -501,4 +492,83 @@ final class TOMLKitTests: XCTestCase {
 		let _ = try udc.nestedUnkeyedContainer()
 		XCTAssertEqual(udc.currentIndex, 1)
 	}
+
+	// Tests for a bug where not-decoded keys would carry over from unsuccessful
+	// decoding attempts, breaking the common 'try one thing then retry another'
+	// decoding pattern. This bug was specific to decoders with
+	// `strictDecoding: true`.
+	func testRetryKeyedOrDictionary() throws {
+		enum StructOrDictionary: Decodable, Equatable {
+			case `struct`(value: Int)
+			case dictionary([String: Int])
+
+			enum CodingKeys: String, CodingKey {
+				case value
+			}
+
+			init(from decoder: Decoder) throws {
+				if let container = try? decoder.container(keyedBy: CodingKeys.self) {
+					if let value = try? container.decode(Int.self, forKey: .value) {
+						self = .struct(value: value)
+					}
+				}
+
+				if let container = try? decoder.singleValueContainer() {
+					self = try .dictionary(container.decode([String: Int].self))
+				} else {
+					throw DecodingError.dataCorrupted(.init(
+						codingPath: [],
+						debugDescription: "Expected int or keyedInt"
+					))
+				}
+			}
+		}
+
+		let toml = "other_value = 1"
+
+		// Before the bug got fixed, this line would throw an unused key error.
+		let value = try TOMLDecoder(strictDecoding: true).decode(StructOrDictionary.self, from: toml)
+
+		// We don't actually expect this to fail, it's not the point of the test, but
+		// might as well assert it just in case.
+		XCTAssertEqual(value, .dictionary(["other_value": 1]))
+	}
+
+	// This is related to the test above, but tests for a different aspect of the
+	// bug (I originally fixed one but not the other so they are kind of independent).
+	func testRetryKeyedOrInt() throws {
+		struct SimpleCodableStruct: Codable, Equatable {
+			var value: Int
+		}
+
+		enum SimpleOrComplex: Decodable, Equatable {
+			case simple(SimpleCodableStruct)
+			case complex(CodableStruct)
+
+			init(from decoder: Decoder) throws {
+				if let container = try? decoder.singleValueContainer() {
+					if let value = try? container.decode(CodableStruct.self) {
+						self = .complex(value)
+						return
+					} else if let value = try? container.decode(SimpleCodableStruct.self) {
+						self = .simple(value)
+						return
+					}
+				}
+				throw DecodingError.dataCorrupted(.init(
+					codingPath: [],
+					debugDescription: "Expected CodableStruct or SimpleCodableStruct"
+				))
+			}
+		}
+
+		let toml = "value = 2"
+
+		// Before the bug got fixed, this line would throw an unused key error.
+		let value = try TOMLDecoder(strictDecoding: true).decode(SimpleOrComplex.self, from: toml)
+
+		// We don't actually expect this to fail, it's not the point of the test, but
+		// might as well assert it just in case.
+		XCTAssertEqual(value, .simple(SimpleCodableStruct(value: 2)))
+	}
 }