From fca4773755a54e1d61101a51abbd2a40856cb80d Mon Sep 17 00:00:00 2001
From: Goncalo Frade <goncalofrade08@gmail.com>
Date: Fri, 22 Nov 2024 16:47:04 +0000
Subject: [PATCH] add option to sort maps by key

so it does not introduce breaking changes this flag is set as default to true

Signed-off-by: Goncalo Frade <goncalofrade08@gmail.com>
---
 Sources/CBOREncoder.swift    | 51 +++++++++++++++++++++++++++---------
 Sources/CBOROptions.swift    |  5 +++-
 Tests/CBOREncoderTests.swift | 15 ++++++++++-
 3 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/Sources/CBOREncoder.swift b/Sources/CBOREncoder.swift
index 8ba85fd..cba2cbe 100644
--- a/Sources/CBOREncoder.swift
+++ b/Sources/CBOREncoder.swift
@@ -136,9 +136,26 @@ extension CBOR {
         res.reserveCapacity(1 + map.count * (MemoryLayout<A>.size + MemoryLayout<B>.size + 2))
         res = map.count.encode(options: options)
         res[0] = res[0] | 0b101_00000
-        for (k, v) in map {
-            res.append(contentsOf: k.encode(options: options))
-            res.append(contentsOf: v.encode(options: options))
+        
+        if options.shouldSortMapKeys {
+            let sortedKeysWithEncodedKeys = map.keys.map {
+                (encoded: $0.encode(options: options), key: $0)
+            }.sorted(by: {
+                $0.encoded.lexicographicallyPrecedes($1.encoded)
+            })
+
+            sortedKeysWithEncodedKeys.forEach { keyTuple in
+                res.append(contentsOf: keyTuple.encoded)
+                guard let value = map[keyTuple.key] else {
+                    return
+                }
+                res.append(contentsOf: value.encode(options: options))
+            }
+        } else {
+            for (k, v) in map {
+                res.append(contentsOf: k.encode(options: options))
+                res.append(contentsOf: v.encode(options: options))
+            }
         }
         return res
     }
@@ -442,16 +459,24 @@ extension CBOR {
         if options.forbidNonStringMapKeys {
             try ensureStringKey(A.self)
         }
-        let sortedKeysWithEncodedKeys = map.keys.map {
-            (encoded: $0.encode(options: options), key: $0)
-        }.sorted(by: {
-            $0.encoded.lexicographicallyPrecedes($1.encoded)
-        })
-
-        try sortedKeysWithEncodedKeys.forEach { keyTuple in
-            res.append(contentsOf: keyTuple.encoded)
-            let encodedVal = try encodeAny(map[keyTuple.key]!, options: options)
-            res.append(contentsOf: encodedVal)
+        if options.shouldSortMapKeys {
+            let sortedKeysWithEncodedKeys = map.keys.map {
+                (encoded: $0.encode(options: options), key: $0)
+            }.sorted(by: {
+                $0.encoded.lexicographicallyPrecedes($1.encoded)
+            })
+
+            try sortedKeysWithEncodedKeys.forEach { keyTuple in
+                res.append(contentsOf: keyTuple.encoded)
+                let encodedVal = try encodeAny(map[keyTuple.key]!, options: options)
+                res.append(contentsOf: encodedVal)
+            }
+        } else {
+            for (k, v) in map {
+                res.append(contentsOf: k.encode(options: options))
+                let encodedVal = try encodeAny(v, options: options)
+                res.append(contentsOf: encodedVal)
+            }
         }
     }
 }
diff --git a/Sources/CBOROptions.swift b/Sources/CBOROptions.swift
index f422857..a0b6244 100644
--- a/Sources/CBOROptions.swift
+++ b/Sources/CBOROptions.swift
@@ -4,17 +4,20 @@ public struct CBOROptions {
     let forbidNonStringMapKeys: Bool
     /// The maximum number of nested items, inclusive, to decode. A maximum set to 0 dissallows anything other than top-level primitives.
     let maximumDepth: Int
+    let shouldSortMapKeys: Bool
 
     public init(
         useStringKeys: Bool = false,
         dateStrategy: DateStrategy = .taggedAsEpochTimestamp,
         forbidNonStringMapKeys: Bool = false,
-        maximumDepth: Int = .max
+        maximumDepth: Int = .max,
+        shouldShortMapKeys: Bool = true
     ) {
         self.useStringKeys = useStringKeys
         self.dateStrategy = dateStrategy
         self.forbidNonStringMapKeys = forbidNonStringMapKeys
         self.maximumDepth = maximumDepth
+        self.shouldSortMapKeys = shouldShortMapKeys
     }
 
     func toCodableEncoderOptions() -> CodableCBOREncoder._Options {
diff --git a/Tests/CBOREncoderTests.swift b/Tests/CBOREncoderTests.swift
index 79f6bc4..dd97788 100644
--- a/Tests/CBOREncoderTests.swift
+++ b/Tests/CBOREncoderTests.swift
@@ -102,7 +102,7 @@ class CBOREncoderTests: XCTestCase {
             "a": 1,
             "b": [2, 3]
         ]
-        let encodedMapToAny = try! CBOR.encodeMap(mapToAny)
+        let encodedMapToAny = try! CBOR.encodeMap(mapToAny, options: .init(shouldShortMapKeys: true))
         XCTAssertEqual(encodedMapToAny, [0xa2, 0x61, 0x61, 0x01, 0x61, 0x62, 0x82, 0x02, 0x03])
 
         let mapToAnyWithIntKeys: [Int: Any] = [
@@ -113,6 +113,19 @@ class CBOREncoderTests: XCTestCase {
             XCTAssertEqual(err as! CBOREncoderError, CBOREncoderError.nonStringKeyInMap)
         }
     }
+    
+    func testEncodeSortedMaps() {
+        XCTAssertEqual(CBOR.encode(Dictionary<Int, Int>()), [0xa0])
+
+        let encoded = CBOR.encode([3: 4, 1: 2])
+        XCTAssert(encoded == [0xa2, 0x01, 0x02, 0x03, 0x04])
+
+        let arr1: CBOR = [1]
+        let arr2: CBOR = [2,3]
+        let nestedEnc: [UInt8] = CBOR.encode(["b": arr2, "a": arr1])
+        let encodedAFirst: [UInt8] = [0xa2, 0x61, 0x61, 0x81, 0x01, 0x61, 0x62, 0x82, 0x02, 0x03]
+        XCTAssert(nestedEnc == encodedAFirst)
+    }
 
     func testEncodeTagged() {
         let bignum: [UInt8] = [0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00] // 2**64