"Enhance CxxDictionary with removeValue(forKey:), subscript(key:default:), merge、merging and init(grouping:by:)"

Author: CrazyFanFan

stdlib/public/Cxx/CxxDictionary.swift

@@ -40,6 +40,10 @@ public protocol CxxDictionary<Key, Value>: ExpressibleByDictionaryLiteral {
   @discardableResult
   mutating func erase(_ key: Key) -> Size
 
+  /// Do not implement this function manually in Swift.
+  @discardableResult
+  mutating func __eraseUnsafe(_ iter: RawMutableIterator) -> RawMutableIterator
+
   /// Do not implement this function manually in Swift.
   func __beginUnsafe() -> RawIterator
 
@@ -73,6 +77,25 @@ extension CxxDictionary {
     }
   }
 
+  @inlinable
+  public init<S: Sequence>(
+    grouping values: __owned S,
+    by keyForValue: (S.Element) throws -> Key
+  ) rethrows where Value: CxxVector<S.Element> {
+    self.init()
+    for value in values {
+      let key = try keyForValue(value)
+      var iter = __findMutatingUnsafe(key)
+      if iter != __endMutatingUnsafe() {
+        iter.pointee.second.push_back(value)
+      } else {
+        var vector = Value()
+        vector.push_back(value)
+        self[key] = vector
+      }
+    }
+  }
+
   @inlinable
   public subscript(key: Key) -> Value? {
     get {
@@ -98,7 +121,29 @@ extension CxxDictionary {
       }
     }
   }
-  
+
+  @inlinable
+  public subscript(
+    key: Key, default defaultValue: @autoclosure () -> Value
+  ) -> Value {
+    get {
+      let iter = __findUnsafe(key)
+      guard iter != __endUnsafe() else {
+        return defaultValue()
+      }
+      return iter.pointee.second
+    }
+    set(newValue) {
+      var iter = self.__findMutatingUnsafe(key)
+      if iter != self.__endMutatingUnsafe() {
+        iter.pointee.second = newValue
+      } else {
+        let keyValuePair = Element(first: key, second: newValue)
+        self.__insertUnsafe(keyValuePair)
+      }
+    }
+  }
+
   public func filter(_ isIncluded: (_ key: Key, _ value: Value) throws -> Bool) rethrows -> Self {
     var filteredDictionary = Self.init()
     var iterator = __beginUnsafe()
@@ -116,4 +161,95 @@ extension CxxDictionary {
 
     return filteredDictionary
   }
+
+  @inlinable
+  @discardableResult
+  public mutating func removeValue(forKey key: Key) -> Value? {
+    var iter = self.__findMutatingUnsafe(key)
+    guard iter != self.__endMutatingUnsafe() else { return nil }
+
+    let value = iter.pointee.second
+    self.__eraseUnsafe(iter)
+    return value
+  }
+
+  @inlinable
+  public mutating func merge<S: Sequence>(
+    _ other: __owned S,
+    uniquingKeysWith combine: (Value, Value) throws -> Value
+  ) rethrows where S.Element == (Key, Value) {
+    for (key, value) in other {
+      var iter = self.__findMutatingUnsafe(key)
+      if iter != self.__endMutatingUnsafe() {
+        iter.pointee.second = try combine(iter.pointee.second, value)
+      } else {
+        let keyValuePair = Element(first: key, second: value)
+        self.__insertUnsafe(keyValuePair)
+      }
+    }
+  }
+
+  @inlinable
+  public mutating func merge(
+    _ other: __owned Dictionary<Key, Value>,
+    uniquingKeysWith combine: (Value, Value) throws -> Value
+  ) rethrows where Key: Hashable {
+    for (key, value) in other {
+      var iter = self.__findMutatingUnsafe(key)
+      if iter != self.__endMutatingUnsafe() {
+        iter.pointee.second = try combine(iter.pointee.second, value)
+      } else {
+        let keyValuePair = Element(first: key, second: value)
+        self.__insertUnsafe(keyValuePair)
+      }
+    }
+  }
+
+  @inlinable
+  public mutating func merge(
+    _ other: __owned Self,
+    uniquingKeysWith combine: (Value, Value) throws -> Value
+  ) rethrows {
+    var iterator = other.__beginUnsafe()
+    while iterator != other.__endUnsafe() {
+      var iter = self.__findMutatingUnsafe(iterator.pointee.first)
+      if iter != self.__endMutatingUnsafe() {
+        iter.pointee.second = try combine(iter.pointee.second, iterator.pointee.second)
+      } else {
+        let keyValuePair = Element(first: iterator.pointee.first, second: iterator.pointee.second)
+        self.__insertUnsafe(keyValuePair)
+      }
+      iterator = iterator.successor()
+    }
+  }
+
+  @inlinable
+  public __consuming func merging<S: Sequence>(
+    _ other: __owned S,
+    uniquingKeysWith combine: (Value, Value) throws -> Value
+  ) rethrows -> Self where S.Element == (Key, Value) {
+    var result = self
+    try result.merge(other, uniquingKeysWith: combine)
+    return result
+  }
+
+  @inlinable
+  public __consuming func merging(
+      _ other: __owned Dictionary<Key, Value>,
+      uniquingKeysWith combine: (Value, Value) throws -> Value
+  ) rethrows -> Self where Key: Hashable {
+    var result = self
+    try result.merge(other, uniquingKeysWith: combine)
+    return result
+  }
+
+  @inlinable
+  public __consuming func merging(
+    _ other: __owned Self,
+    uniquingKeysWith combine: (Value, Value) throws -> Value
+  ) rethrows -> Self {
+    var result = self
+    try result.merge(other, uniquingKeysWith: combine)
+    return result
+  }
 }

test/Interop/Cxx/stdlib/Inputs/std-map.h

@@ -4,12 +4,14 @@
 #include <map>
 #include <string>
 #include <unordered_map>
+#include <vector>
 
 using Map = std::map<int, int>;
 using MapStrings = std::map<std::string, std::string>;
 using NestedMap = std::map<int, Map>;
+using MapGroup = std::map<int, std::vector<int>>;
 using UnorderedMap = std::unordered_map<int, int>;
-
+using UnorderedMapGroup = std::unordered_map<int, std::vector<int>>;
 inline Map initMap() { return {{1, 3}, {2, 2}, {3, 3}}; }
 inline UnorderedMap initUnorderedMap() { return {{1, 3}, {3, 3}, {2, 2}}; }
 inline Map initEmptyMap() { return {}; }