Merge pull request #77445 from CrazyFanFan/feature_cxxset_remove

[cxx-interop] Allow removing elements from `std::set`.

Author: egorzhdan

lib/ClangImporter/ClangDerivedConformances.cpp

@@ -928,22 +928,28 @@ void swift::conformToCxxSetIfNeeded(ClangImporter::Implementation &impl,
   if (!isStdDecl(clangDecl, {"set", "unordered_set"}))
     return;
 
-  ProtocolDecl *cxxIteratorProto =
+  ProtocolDecl *cxxInputIteratorProto =
       ctx.getProtocol(KnownProtocolKind::UnsafeCxxInputIterator);
-  if (!cxxIteratorProto)
+  if (!cxxInputIteratorProto)
     return;
 
+  auto rawIteratorType = lookupDirectSingleWithoutExtensions<TypeAliasDecl>(
+      decl, ctx.getIdentifier("const_iterator"));
   auto rawMutableIteratorType =
       lookupDirectSingleWithoutExtensions<TypeAliasDecl>(
           decl, ctx.getIdentifier("iterator"));
-  if (!rawMutableIteratorType)
+  if (!rawIteratorType || !rawMutableIteratorType)
     return;
 
+  auto rawIteratorTy = rawIteratorType->getUnderlyingType();
   auto rawMutableIteratorTy = rawMutableIteratorType->getUnderlyingType();
-  // Check if RawMutableIterator conforms to UnsafeCxxInputIterator.
-  if (!checkConformance(rawMutableIteratorTy, cxxIteratorProto))
+
+  if (!checkConformance(rawIteratorTy, cxxInputIteratorProto) ||
+      !checkConformance(rawMutableIteratorTy, cxxInputIteratorProto))
     return;
 
+  impl.addSynthesizedTypealias(decl, ctx.getIdentifier("RawIterator"),
+                               rawIteratorTy);
   impl.addSynthesizedTypealias(decl, ctx.getIdentifier("RawMutableIterator"),
                                rawMutableIteratorTy);
   impl.addSynthesizedProtocolAttrs(decl, {KnownProtocolKind::CxxUniqueSet});

stdlib/public/Cxx/CxxSet.swift

@@ -67,10 +67,21 @@ extension CxxSet {
 public protocol CxxUniqueSet<Element>: CxxSet {
   override associatedtype Element
   override associatedtype Size: BinaryInteger
+  associatedtype RawIterator: UnsafeCxxInputIterator
+    where RawIterator.Pointee == Element
   associatedtype RawMutableIterator: UnsafeCxxInputIterator
     where RawMutableIterator.Pointee == Element
   override associatedtype InsertionResult
     where InsertionResult: CxxPair<RawMutableIterator, Bool>
+
+  @discardableResult
+  mutating func __findUnsafe(_ value: Element) -> RawIterator
+
+  @discardableResult
+  mutating func __eraseUnsafe(_ iter: RawIterator) -> RawMutableIterator
+
+  @discardableResult
+  mutating func __endUnsafe() -> RawIterator
 }
 
 extension CxxUniqueSet {
@@ -94,4 +105,19 @@ extension CxxUniqueSet {
     let inserted: Bool = insertionResult.second
     return (inserted, rawIterator.pointee)
   }
+
+  /// Removes the given element from the set.
+  ///
+  /// - Parameter member: An element to remove from the set.
+  @discardableResult
+  @inlinable
+  public mutating func remove(_ member: Element) -> Element? {
+    let iter = self.__findUnsafe(member)
+    guard iter != self.__endUnsafe() else {
+      return nil
+    }
+    let value = iter.pointee
+    self.__eraseUnsafe(iter)
+    return value
+  }
 }

test/Interop/Cxx/stdlib/use-std-set.swift

@@ -164,4 +164,22 @@ StdSetTestSuite.test("UnorderedSetOfCInt.erase") {
     expectFalse(s.contains(2))
 }
 
+StdSetTestSuite.test("SetOfCInt.remove") {
+    var s = initSetOfCInt()
+    expectTrue(s.contains(1))
+    expectEqual(s.remove(1), 1)
+    expectFalse(s.contains(1))
+    expectEqual(s.remove(1), nil)
+    expectFalse(s.contains(1))
+}
+
+StdSetTestSuite.test("UnorderedSetOfCInt.remove") {
+    var s = initUnorderedSetOfCInt()
+    expectTrue(s.contains(2))
+    expectEqual(s.remove(2), 2)
+    expectFalse(s.contains(2))
+    expectEqual(s.remove(2), nil)
+    expectFalse(s.contains(2))
+}
+
 runAllTests()