Merge pull request #66764 from apple/egorzhdan/cxx-set-insert

[cxx-interop] Allow inserting elements into `std::set` from Swift

Author: egorzhdan

include/swift/AST/KnownProtocols.def

@@ -114,6 +114,7 @@ PROTOCOL(CxxPair)
 PROTOCOL(CxxSet)
 PROTOCOL(CxxRandomAccessCollection)
 PROTOCOL(CxxSequence)
+PROTOCOL(CxxUniqueSet)
 PROTOCOL(UnsafeCxxInputIterator)
 PROTOCOL(UnsafeCxxRandomAccessIterator)
 

lib/AST/ASTContext.cpp

@@ -1138,6 +1138,7 @@ ProtocolDecl *ASTContext::getProtocol(KnownProtocolKind kind) const {
   case KnownProtocolKind::CxxRandomAccessCollection:
   case KnownProtocolKind::CxxSet:
   case KnownProtocolKind::CxxSequence:
+  case KnownProtocolKind::CxxUniqueSet:
   case KnownProtocolKind::UnsafeCxxInputIterator:
   case KnownProtocolKind::UnsafeCxxRandomAccessIterator:
     M = getLoadedModule(Id_Cxx);

lib/ClangImporter/ClangDerivedConformances.cpp

@@ -614,6 +614,34 @@ void swift::conformToCxxSetIfNeeded(ClangImporter::Implementation &impl,
   impl.addSynthesizedTypealias(decl, ctx.getIdentifier("InsertionResult"),
                                insert->getResultInterfaceType());
   impl.addSynthesizedProtocolAttrs(decl, {KnownProtocolKind::CxxSet});
+
+  // If this isn't a std::multiset, try to also synthesize the conformance to
+  // CxxUniqueSet.
+  if (!isStdDecl(clangDecl, {"set", "unordered_set"}))
+    return;
+
+  ProtocolDecl *cxxIteratorProto =
+      ctx.getProtocol(KnownProtocolKind::UnsafeCxxInputIterator);
+  if (!cxxIteratorProto)
+    return;
+
+  auto rawMutableIteratorType =
+      lookupDirectSingleWithoutExtensions<TypeAliasDecl>(
+          decl, ctx.getIdentifier("iterator"));
+  if (!rawMutableIteratorType)
+    return;
+
+  auto rawMutableIteratorTy = rawMutableIteratorType->getUnderlyingType();
+  // Check if RawMutableIterator conforms to UnsafeCxxInputIterator.
+  ModuleDecl *module = decl->getModuleContext();
+  auto rawIteratorConformanceRef =
+      module->lookupConformance(rawMutableIteratorTy, cxxIteratorProto);
+  if (!isConcreteAndValid(rawIteratorConformanceRef, module))
+    return;
+
+  impl.addSynthesizedTypealias(decl, ctx.getIdentifier("RawMutableIterator"),
+                               rawMutableIteratorTy);
+  impl.addSynthesizedProtocolAttrs(decl, {KnownProtocolKind::CxxUniqueSet});
 }
 
 void swift::conformToCxxPairIfNeeded(ClangImporter::Implementation &impl,

lib/IRGen/GenMeta.cpp

@@ -6288,6 +6288,7 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::CxxRandomAccessCollection:
   case KnownProtocolKind::CxxSet:
   case KnownProtocolKind::CxxSequence:
+  case KnownProtocolKind::CxxUniqueSet:
   case KnownProtocolKind::UnsafeCxxInputIterator:
   case KnownProtocolKind::UnsafeCxxRandomAccessIterator:
   case KnownProtocolKind::Executor:

stdlib/public/Cxx/CxxSet.swift

@@ -13,7 +13,10 @@
 public protocol CxxSet<Element> {
   associatedtype Element
   associatedtype Size: BinaryInteger
-  associatedtype InsertionResult // std::pair<iterator, bool>
+
+  // std::pair<iterator, bool> for std::set and std::unordered_set
+  // iterator for std::multiset
+  associatedtype InsertionResult 
 
   init()
 
@@ -43,3 +46,25 @@ extension CxxSet {
     return count(element) > 0
   }
 }
+
+public protocol CxxUniqueSet<Element>: CxxSet {
+  override associatedtype Element
+  override associatedtype Size: BinaryInteger
+  associatedtype RawMutableIterator: UnsafeCxxInputIterator
+    where RawMutableIterator.Pointee == Element
+  override associatedtype InsertionResult
+    where InsertionResult: CxxPair<RawMutableIterator, Bool>
+}
+
+extension CxxUniqueSet {
+  @inlinable
+  @discardableResult
+  public mutating func insert(
+    _ element: Element
+  ) -> (inserted: Bool, memberAfterInsert: Element) {
+    let insertionResult = self.__insertUnsafe(element)
+    let rawIterator: RawMutableIterator = insertionResult.first
+    let inserted: Bool = insertionResult.second
+    return (inserted, rawIterator.pointee)
+  }
+}

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

@@ -61,4 +61,32 @@ StdSetTestSuite.test("UnorderedSetOfCInt.init()") {
     expectTrue(s.contains(3))
 }
 
+StdSetTestSuite.test("SetOfCInt.insert") {
+    var s = SetOfCInt()
+    expectFalse(s.contains(123))
+
+    let res1 = s.insert(123)
+    expectTrue(res1.inserted)
+    expectTrue(s.contains(123))
+
+    let res2 = s.insert(123)
+    expectFalse(res2.inserted)
+    expectTrue(s.contains(123))
+}
+
+#if !os(Linux) // FIXME: https://github.com/apple/swift/issues/66767 / rdar://105220600
+StdSetTestSuite.test("UnorderedSetOfCInt.insert") {
+    var s = UnorderedSetOfCInt()
+    expectFalse(s.contains(123))
+
+    let res1 = s.insert(123)
+    expectTrue(res1.inserted)
+    expectTrue(s.contains(123))
+
+    let res2 = s.insert(123)
+    expectFalse(res2.inserted)
+    expectTrue(s.contains(123))
+}
+#endif
+
 runAllTests()