Add Key to combine closure

Author: amomchilov

Guides/Keyed.md

@@ -22,7 +22,7 @@ Duplicate keys will trigger a runtime error by default. To handle this, you can
 let fruits = ["Apricot", "Banana", "Apple", "Cherry", "Blackberry", "Avocado", "Coconut"]
 let fruitsByLetter = fruits.keyed(
     by: { $0.first! },
-    uniquingKeysWith: { old, new in new } // Always pick the latest fruit
+    uniquingKeysWith: { key, old, new in new } // Always pick the latest fruit
 )
 // Results in:
 [
@@ -40,7 +40,7 @@ The `keyed(by:)` method is declared as a `Sequence` extension returning `[Key: E
 extension Sequence {
     public func keyed<Key>(
         by keyForValue: (Element) throws -> Key,
-        uniquingKeysWith combine: ((Element, Element) throws -> Element)? = nil
+        uniquingKeysWith combine: ((Key, Element, Element) throws -> Element)? = nil
     ) rethrows -> [Key: Element]
 }
 ```

Sources/Algorithms/Keyed.swift

@@ -30,18 +30,41 @@ extension Sequence {
   @inlinable
   public func keyed<Key>(
     by keyForValue: (Element) throws -> Key,
-    // TODO: pass `Key` into `combine`: (Key, Element, Element) throws -> Element
-    uniquingKeysWith combine: ((Element, Element) throws -> Element)? = nil
+    uniquingKeysWith combine: ((Key, Element, Element) throws -> Element)? = nil
   ) rethrows -> [Key: Element] {
-    // Note: This implementation is a bit convoluted, but it's just aiming to reuse the existing stdlib logic,
-    // to ensure consistent behaviour, error messages, etc.
-    // If this API ends up in the stdlib itself, it could just call the underlying `_NativeDictionary` methods.
-    try withoutActuallyEscaping(keyForValue) { keyForValue in
-      if let combine {
-        return try Dictionary(self.lazy.map { (try keyForValue($0), $0) }, uniquingKeysWith: combine)
-      } else {
-        return try Dictionary(uniqueKeysWithValues: self.lazy.map { (try keyForValue($0), $0) } )
+    var result = [Key: Element]()
+
+    if combine != nil {
+      // We have a `combine` closure. Use it to resolve duplicate keys.
+
+      for element in self {
+        let key = try keyForValue(element)
+        
+        if let oldValue = result.updateValue(element, forKey: key) {
+          // Can't use a conditional binding to unwrap this, because the newly bound variable
+          // doesn't play nice with the `rethrows` system.
+          let valueToKeep = try combine!(key, oldValue, element)
+          
+          // This causes a second look-up for the same key. The standard library can avoid that
+          // by calling `mutatingFind` to get access to the bucket where the value will end up,
+          // and updating in place.
+          // Swift Algorithms doesn't have access to that API, so we make due.
+          // When this gets merged into the standard library, we should optimize this.
+          result[key] = valueToKeep
+        }
+      }
+    } else {
+      // There's no `combine` closure. Duplicate keys are disallowed.
+
+      for element in self {
+        let key = try keyForValue(element)
+
+        guard result.updateValue(element, forKey: key) == nil else {
+          fatalError("Duplicate values for key: '\(key)'")
+        }
       }
     }
+
+    return result
   }
 }

Tests/SwiftAlgorithmsTests/KeyedTests.swift

@@ -38,16 +38,30 @@ final class KeyedTests: XCTestCase {
   }
 
   func testNonUniqueKeysWithMergeFunction() {
+    var combineCallHistory = [(key: Character, current: String, new: String)]()
+    let expectedCallHistory = [
+      (key: "A", current: "Apple", new: "Avocado"),
+      (key: "C", current: "Cherry", new: "Coconut"),
+    ]
+
     let d = ["Apple", "Avocado", "Banana", "Cherry", "Coconut"].keyed(
       by: { $0.first! },
-      uniquingKeysWith: { older, newer in "\(older)-\(newer)"}
+      uniquingKeysWith: { key, older, newer in
+        combineCallHistory.append((key, older, newer))
+        return "\(older)-\(newer)"
+      }
     )
 
     XCTAssertEqual(d.count, 3)
     XCTAssertEqual(d["A"]!, "Apple-Avocado")
     XCTAssertEqual(d["B"]!, "Banana")
     XCTAssertEqual(d["C"]!, "Cherry-Coconut")
     XCTAssertNil(d["D"])
+
+    XCTAssertEqual(
+      combineCallHistory.map(String.init(describing:)), // quick/dirty workaround: tuples aren't Equatable
+      expectedCallHistory.map(String.init(describing:))
+    )
   }
 
   func testThrowingFromKeyFunction() {
@@ -66,7 +80,7 @@ final class KeyedTests: XCTestCase {
     let error = SampleError()
 
     XCTAssertThrowsError(
-      try input.keyed(by: { $0.first! }, uniquingKeysWith: { _, _ in throw error })
+      try input.keyed(by: { $0.first! }, uniquingKeysWith: { _, _, _ in throw error })
     ) { thrownError in
       XCTAssertIdentical(error, thrownError as? SampleError)
     }