Merge pull request swiftlang#63244 from apple/egorzhdan/cxx-set-protocol

[cxx-interop] Add `CxxSet` protocol for `std::set` ergonomics

Author: egorzhdan

include/swift/AST/KnownProtocols.def

@@ -106,6 +106,7 @@ PROTOCOL(DistributedTargetInvocationResultHandler)
 
 // C++ Standard Library Overlay:
 PROTOCOL(CxxConvertibleToCollection)
+PROTOCOL(CxxSet)
 PROTOCOL(CxxRandomAccessCollection)
 PROTOCOL(CxxSequence)
 PROTOCOL(UnsafeCxxInputIterator)

lib/AST/ASTContext.cpp

@@ -1126,6 +1126,7 @@ ProtocolDecl *ASTContext::getProtocol(KnownProtocolKind kind) const {
     break;
   case KnownProtocolKind::CxxConvertibleToCollection:
   case KnownProtocolKind::CxxRandomAccessCollection:
+  case KnownProtocolKind::CxxSet:
   case KnownProtocolKind::CxxSequence:
   case KnownProtocolKind::UnsafeCxxInputIterator:
   case KnownProtocolKind::UnsafeCxxRandomAccessIterator:

lib/ClangImporter/ClangDerivedConformances.cpp

@@ -520,3 +520,41 @@ void swift::conformToCxxSequenceIfNeeded(
     impl.addSynthesizedProtocolAttrs(
         decl, {KnownProtocolKind::CxxConvertibleToCollection});
 }
+
+void swift::conformToCxxSetIfNeeded(ClangImporter::Implementation &impl,
+                                    NominalTypeDecl *decl,
+                                    const clang::CXXRecordDecl *clangDecl) {
+  PrettyStackTraceDecl trace("conforming to CxxSet", decl);
+
+  assert(decl);
+  assert(clangDecl);
+  ASTContext &ctx = decl->getASTContext();
+
+  // Only auto-conform types from the C++ standard library. Custom user types
+  // might have a similar interface but different semantics.
+  if (!clangDecl->isInStdNamespace())
+    return;
+  if (!clangDecl->getIdentifier())
+    return;
+  StringRef name = clangDecl->getName();
+  if (name != "set" && name != "unordered_set" && name != "multiset")
+    return;
+
+  auto valueTypeId = ctx.getIdentifier("value_type");
+  auto valueTypes = lookupDirectWithoutExtensions(decl, valueTypeId);
+  if (valueTypes.size() != 1)
+    return;
+  auto valueType = dyn_cast<TypeAliasDecl>(valueTypes.front());
+
+  auto sizeTypeId = ctx.getIdentifier("size_type");
+  auto sizeTypes = lookupDirectWithoutExtensions(decl, sizeTypeId);
+  if (sizeTypes.size() != 1)
+    return;
+  auto sizeType = dyn_cast<TypeAliasDecl>(sizeTypes.front());
+
+  impl.addSynthesizedTypealias(decl, ctx.Id_Element,
+                               valueType->getUnderlyingType());
+  impl.addSynthesizedTypealias(decl, ctx.getIdentifier("Size"),
+                               sizeType->getUnderlyingType());
+  impl.addSynthesizedProtocolAttrs(decl, {KnownProtocolKind::CxxSet});
+}

lib/ClangImporter/ClangDerivedConformances.h

@@ -32,6 +32,13 @@ void conformToCxxSequenceIfNeeded(ClangImporter::Implementation &impl,
                                   NominalTypeDecl *decl,
                                   const clang::CXXRecordDecl *clangDecl);
 
+/// If the decl is an instantiation of C++ `std::set`, `std::unordered_set` or
+/// `std::multiset`, synthesize a conformance to CxxSet, which is defined in the
+/// Cxx module.
+void conformToCxxSetIfNeeded(ClangImporter::Implementation &impl,
+                             NominalTypeDecl *decl,
+                             const clang::CXXRecordDecl *clangDecl);
+
 } // namespace swift
 
 #endif // SWIFT_CLANG_DERIVED_CONFORMANCES_H

lib/ClangImporter/ImportDecl.cpp

@@ -2609,6 +2609,7 @@ namespace {
         auto nominalDecl = cast<NominalTypeDecl>(result);
         conformToCxxIteratorIfNeeded(Impl, nominalDecl, decl);
         conformToCxxSequenceIfNeeded(Impl, nominalDecl, decl);
+        conformToCxxSetIfNeeded(Impl, nominalDecl, decl);
       }
 
       if (auto *ntd = dyn_cast<NominalTypeDecl>(result))

lib/IRGen/GenMeta.cpp

@@ -5876,6 +5876,7 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::DistributedTargetInvocationResultHandler:
   case KnownProtocolKind::CxxConvertibleToCollection:
   case KnownProtocolKind::CxxRandomAccessCollection:
+  case KnownProtocolKind::CxxSet:
   case KnownProtocolKind::CxxSequence:
   case KnownProtocolKind::UnsafeCxxInputIterator:
   case KnownProtocolKind::UnsafeCxxRandomAccessIterator:

stdlib/public/Cxx/CMakeLists.txt

@@ -5,6 +5,7 @@ endif()
 
 add_swift_target_library(swiftCxx ${SWIFT_CXX_LIBRARY_KIND} NO_LINK_NAME IS_STDLIB
     CxxConvertibleToCollection.swift
+    CxxSet.swift
     CxxRandomAccessCollection.swift
     CxxSequence.swift
 

stdlib/public/Cxx/CxxSet.swift

@@ -0,0 +1,24 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+public protocol CxxSet<Element> {
+  associatedtype Element
+  associatedtype Size: BinaryInteger
+
+  func count(_ element: Element) -> Size
+}
+
+extension CxxSet {
+  public func contains(_ element: Element) -> Bool {
+    return count(element) > 0
+  }
+}

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

@@ -2,9 +2,14 @@
 #define TEST_INTEROP_CXX_STDLIB_INPUTS_STD_SET_H
 
 #include <set>
+#include <unordered_set>
 
 using SetOfCInt = std::set<int>;
+using UnorderedSetOfCInt = std::unordered_set<int>;
+using MultisetOfCInt = std::multiset<int>;
 
 inline SetOfCInt initSetOfCInt() { return {1, 5, 3}; }
+inline UnorderedSetOfCInt initUnorderedSetOfCInt() { return {2, 4, 6}; }
+inline MultisetOfCInt initMultisetOfCInt() { return {2, 2, 4, 6}; }
 
 #endif // TEST_INTEROP_CXX_STDLIB_INPUTS_STD_SET_H

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

@@ -12,10 +12,8 @@ import Cxx
 
 var StdSetTestSuite = TestSuite("StdSet")
 
-extension SetOfCInt : CxxSequence { }
-
-StdSetTestSuite.test("iterate") {
-    let s = initSetOfCInt()
+StdSetTestSuite.test("iterate over Swift.Array") {
+    let s = Array(initSetOfCInt())
     var result = [CInt]()
     for x in s {
         result.append(x)
@@ -25,4 +23,28 @@ StdSetTestSuite.test("iterate") {
     expectEqual(result[2], 5)
 }
 
+StdSetTestSuite.test("SetOfCInt.contains") {
+    // This relies on the `std::set` conformance to `CxxSet` protocol.
+    let s = initSetOfCInt()
+    expectTrue(s.contains(1))
+    expectFalse(s.contains(2))
+    expectTrue(s.contains(3))
+}
+
+StdSetTestSuite.test("UnorderedSetOfCInt.contains") {
+    // This relies on the `std::unordered_set` conformance to `CxxSet` protocol.
+    let s = initUnorderedSetOfCInt()
+    expectFalse(s.contains(1))
+    expectTrue(s.contains(2))
+    expectFalse(s.contains(3))
+}
+
+StdSetTestSuite.test("MultisetOfCInt.contains") {
+    // This relies on the `std::multiset` conformance to `CxxSet` protocol.
+    let s = initMultisetOfCInt()
+    expectFalse(s.contains(1))
+    expectTrue(s.contains(2))
+    expectFalse(s.contains(3))
+}
+
 runAllTests()