[cxx-interop] Add CxxMutableRandomAccessCollection protocol

egorzhdan · egorzhdan · commit 0ab681514bf4 · 2024-08-28T12:43:42.000+01:00
This conforms mutable C++ container types, such as `std::vector`, to `MutableCollection` via a new overlay protocol `CxxMutableRandomAccessCollection`.

rdar://134531554
diff --git a/include/swift/AST/KnownProtocols.def b/include/swift/AST/KnownProtocols.def
@@ -132,6 +132,7 @@ PROTOCOL(CxxOptional)
 PROTOCOL(CxxPair)
 PROTOCOL(CxxSet)
 PROTOCOL(CxxRandomAccessCollection)
+PROTOCOL(CxxMutableRandomAccessCollection)
 PROTOCOL(CxxSequence)
 PROTOCOL(CxxUniqueSet)
 PROTOCOL(CxxVector)
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -1419,6 +1419,7 @@ ProtocolDecl *ASTContext::getProtocol(KnownProtocolKind kind) const {
   case KnownProtocolKind::CxxPair:
   case KnownProtocolKind::CxxOptional:
   case KnownProtocolKind::CxxRandomAccessCollection:
+  case KnownProtocolKind::CxxMutableRandomAccessCollection:
   case KnownProtocolKind::CxxSet:
   case KnownProtocolKind::CxxSequence:
   case KnownProtocolKind::CxxUniqueSet:
diff --git a/lib/ClangImporter/ClangDerivedConformances.cpp b/lib/ClangImporter/ClangDerivedConformances.cpp
@@ -780,8 +780,44 @@ void swift::conformToCxxSequenceIfNeeded(
     impl.addSynthesizedTypealias(decl, ctx.getIdentifier("Indices"), indicesTy);
     impl.addSynthesizedTypealias(decl, ctx.getIdentifier("SubSequence"),
                                  sliceTy);
-    impl.addSynthesizedProtocolAttrs(
-        decl, {KnownProtocolKind::CxxRandomAccessCollection});
+
+    auto tryToConformToMutatingRACollection = [&]() -> bool {
+      auto rawMutableIteratorProto = ctx.getProtocol(
+          KnownProtocolKind::UnsafeCxxMutableRandomAccessIterator);
+      if (!rawMutableIteratorProto)
+        return false;
+
+      // Check if present: `func __beginMutatingUnsafe() -> RawMutableIterator`
+      auto beginMutatingId = ctx.getIdentifier("__beginMutatingUnsafe");
+      auto beginMutating =
+          lookupDirectSingleWithoutExtensions<FuncDecl>(decl, beginMutatingId);
+      if (!beginMutating)
+        return false;
+      auto rawMutableIteratorTy = beginMutating->getResultInterfaceType();
+
+      // Check if present: `func __endMutatingUnsafe() -> RawMutableIterator`
+      auto endMutatingId = ctx.getIdentifier("__endMutatingUnsafe");
+      auto endMutating =
+          lookupDirectSingleWithoutExtensions<FuncDecl>(decl, endMutatingId);
+      if (!endMutating)
+        return false;
+
+      if (!checkConformance(rawMutableIteratorTy, rawMutableIteratorProto))
+        return false;
+
+      impl.addSynthesizedTypealias(
+          decl, ctx.getIdentifier("RawMutableIterator"), rawMutableIteratorTy);
+      impl.addSynthesizedProtocolAttrs(
+          decl, {KnownProtocolKind::CxxMutableRandomAccessCollection});
+      return true;
+    };
+
+    bool conformedToMutableRAC = tryToConformToMutatingRACollection();
+
+    if (!conformedToMutableRAC)
+      impl.addSynthesizedProtocolAttrs(
+          decl, {KnownProtocolKind::CxxRandomAccessCollection});
+
     return true;
   };
 
diff --git a/lib/IRGen/GenMeta.cpp b/lib/IRGen/GenMeta.cpp
@@ -6883,6 +6883,7 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::CxxPair:
   case KnownProtocolKind::CxxOptional:
   case KnownProtocolKind::CxxRandomAccessCollection:
+  case KnownProtocolKind::CxxMutableRandomAccessCollection:
   case KnownProtocolKind::CxxSet:
   case KnownProtocolKind::CxxSequence:
   case KnownProtocolKind::CxxUniqueSet:
diff --git a/stdlib/public/Cxx/CxxRandomAccessCollection.swift b/stdlib/public/Cxx/CxxRandomAccessCollection.swift
@@ -36,16 +36,52 @@ extension CxxRandomAccessCollection {
     return Int(__endUnsafe() - __beginUnsafe())
   }
 
+  @inlinable
+  @inline(__always)
+  internal func _getRawIterator(at index: Int) -> RawIterator {
+    var rawIterator = self.__beginUnsafe()
+    rawIterator += RawIterator.Distance(index)
+    precondition(self.__endUnsafe() - rawIterator > 0,
+        "C++ iterator access out of bounds")
+    return rawIterator
+  }
+
   /// A C++ implementation of the subscript might be more performant. This 
   /// overload should only be used if the C++ type does not define `operator[]`.
   @inlinable
   public subscript(_ index: Int) -> Element {
     _read {
-      // Not using CxxIterator here to avoid making a copy of the collection.
-      var rawIterator = __beginUnsafe()
-      rawIterator += RawIterator.Distance(index)
-      precondition(__endUnsafe() - rawIterator > 0, "C++ iterator access out of bounds")  
-      yield rawIterator.pointee
+      yield self._getRawIterator(at: index).pointee
+    }
+  }
+}
+
+public protocol CxxMutableRandomAccessCollection<Element>:
+    CxxRandomAccessCollection, MutableCollection {
+  associatedtype RawMutableIterator: UnsafeCxxMutableRandomAccessIterator
+      where RawMutableIterator.Pointee == Element
+
+  /// Do not implement this function manually in Swift.
+  mutating func __beginMutatingUnsafe() -> RawMutableIterator
+
+  /// Do not implement this function manually in Swift.
+  mutating func __endMutatingUnsafe() -> RawMutableIterator
+}
+
+extension CxxMutableRandomAccessCollection {
+  /// A C++ implementation of the subscript might be more performant. This 
+  /// overload should only be used if the C++ type does not define `operator[]`.
+  @inlinable
+  public subscript(_ index: Int) -> Element {
+    _read {
+      yield self._getRawIterator(at: index).pointee
+    }
+    _modify {
+      var rawIterator = self.__beginMutatingUnsafe()
+      rawIterator += RawMutableIterator.Distance(index)
+      precondition(self.__endMutatingUnsafe() - rawIterator > 0,
+          "C++ iterator access out of bounds")
+      yield &rawIterator.pointee
     }
   }
 }
diff --git a/test/Interop/Cxx/stdlib/libcxx-module-interface.swift b/test/Interop/Cxx/stdlib/libcxx-module-interface.swift
@@ -20,10 +20,10 @@
 
 // CHECK-IOSFWD: enum std {
 // CHECK-IOSFWD:   enum __1 {
-// CHECK-IOSFWD:     struct basic_string<CChar, char_traits<CChar>, allocator<CChar>> : CxxRandomAccessCollection {
+// CHECK-IOSFWD:     struct basic_string<CChar, char_traits<CChar>, allocator<CChar>> : CxxMutableRandomAccessCollection {
 // CHECK-IOSFWD:       typealias value_type = CChar
 // CHECK-IOSFWD:     }
-// CHECK-IOSFWD:     struct basic_string<CWideChar, char_traits<CWideChar>, allocator<CWideChar>> : CxxRandomAccessCollection {
+// CHECK-IOSFWD:     struct basic_string<CWideChar, char_traits<CWideChar>, allocator<CWideChar>> : CxxMutableRandomAccessCollection {
 // CHECK-IOSFWD:       typealias value_type = CWideChar
 // CHECK-IOSFWD:     }
 // CHECK-IOSFWD:     typealias string = std.__1.basic_string<CChar, char_traits<CChar>, allocator<CChar>>
diff --git a/test/Interop/Cxx/stdlib/libstdcxx-module-interface.swift b/test/Interop/Cxx/stdlib/libstdcxx-module-interface.swift
@@ -11,10 +11,10 @@
 // REQUIRES: OS=linux-gnu
 
 // CHECK-STD: enum std {
-// CHECK-STRING:   struct basic_string<CChar, char_traits<CChar>, allocator<CChar>> : CxxRandomAccessCollection {
+// CHECK-STRING:   struct basic_string<CChar, char_traits<CChar>, allocator<CChar>> : CxxMutableRandomAccessCollection {
 // CHECK-STRING:     typealias value_type = std.char_traits<CChar>.char_type
 // CHECK-STRING:   }
-// CHECK-STRING:   struct basic_string<CWideChar, char_traits<CWideChar>, allocator<CWideChar>> : CxxRandomAccessCollection {
+// CHECK-STRING:   struct basic_string<CWideChar, char_traits<CWideChar>, allocator<CWideChar>> : CxxMutableRandomAccessCollection {
 // CHECK-STRING:     typealias value_type = std.char_traits<CWideChar>.char_type
 // CHECK-STRING:   }
 
diff --git a/test/Interop/Cxx/stdlib/overlay/Inputs/custom-collection.h b/test/Interop/Cxx/stdlib/overlay/Inputs/custom-collection.h
@@ -11,8 +11,8 @@ struct SimpleCollectionNoSubscript {
 public:
   using iterator = ConstRACIterator;
 
-  iterator begin() const { return iterator(*x); }
-  iterator end() const { return iterator(*x + 5); }
+  iterator begin() const { return iterator(x); }
+  iterator end() const { return iterator(x + 5); }
 };
 
 struct SimpleCollectionReadOnly {
@@ -22,12 +22,29 @@ struct SimpleCollectionReadOnly {
 public:
   using iterator = ConstRACIteratorRefPlusEq;
 
-  iterator begin() const { return iterator(*x); }
-  iterator end() const { return iterator(*x + 5); }
+  iterator begin() const { return iterator(x); }
+  iterator end() const { return iterator(x + 5); }
 
   const int& operator[](int index) const { return x[index]; }
 };
 
+struct SimpleCollectionReadWrite {
+private:
+  int x[5] = {1, 2, 3, 4, 5};
+
+public:
+  using const_iterator = ConstRACIterator;
+  using iterator = MutableRACIterator;
+
+  const_iterator begin() const { return const_iterator(x); }
+  const_iterator end() const { return const_iterator(x + 5); }
+  iterator begin() { return iterator(x); }
+  iterator end() { return iterator(x + 5); }
+
+  const int &operator[](int index) const { return x[index]; }
+  int &operator[](int index) { return x[index]; }
+};
+
 template <typename T>
 struct HasInheritedTemplatedConstRACIterator {
 public:
diff --git a/test/Interop/Cxx/stdlib/overlay/Inputs/custom-iterator.h b/test/Interop/Cxx/stdlib/overlay/Inputs/custom-iterator.h
@@ -42,7 +42,7 @@ struct ConstIterator {
 
 struct ConstRACIterator {
 private:
-  int value;
+  const int *value;
 
 public:
   using iterator_category = std::random_access_iterator_tag;
@@ -51,10 +51,10 @@ struct ConstRACIterator {
   using reference = const int &;
   using difference_type = int;
 
-  ConstRACIterator(int value) : value(value) {}
+  ConstRACIterator(const int *value) : value(value) {}
   ConstRACIterator(const ConstRACIterator &other) = default;
 
-  const int &operator*() const { return value; }
+  const int &operator*() const { return *value; }
 
   ConstRACIterator &operator++() {
     value++;
@@ -97,7 +97,7 @@ struct ConstRACIterator {
 // Same as ConstRACIterator, but operator+= returns a reference to this.
 struct ConstRACIteratorRefPlusEq {
 private:
-  int value;
+  const int *value;
 
 public:
   using iterator_category = std::random_access_iterator_tag;
@@ -106,10 +106,10 @@ struct ConstRACIteratorRefPlusEq {
   using reference = const int &;
   using difference_type = int;
 
-  ConstRACIteratorRefPlusEq(int value) : value(value) {}
+  ConstRACIteratorRefPlusEq(const int *value) : value(value) {}
   ConstRACIteratorRefPlusEq(const ConstRACIteratorRefPlusEq &other) = default;
 
-  const int &operator*() const { return value; }
+  const int &operator*() const { return *value; }
 
   ConstRACIteratorRefPlusEq &operator++() {
     value++;
@@ -918,7 +918,7 @@ struct InputOutputConstIterator {
 
 struct MutableRACIterator {
 private:
-  int value;
+  int *value;
 
 public:
   struct iterator_category : std::random_access_iterator_tag,
@@ -928,11 +928,11 @@ struct MutableRACIterator {
   using reference = const int &;
   using difference_type = int;
 
-  MutableRACIterator(int value) : value(value) {}
+  MutableRACIterator(int *value) : value(value) {}
   MutableRACIterator(const MutableRACIterator &other) = default;
 
-  const int &operator*() const { return value; }
-  int &operator*() { return value; }
+  const int &operator*() const { return *value; }
+  int &operator*() { return *value; }
 
   MutableRACIterator &operator++() {
     value++;
diff --git a/test/Interop/Cxx/stdlib/overlay/custom-collection-module-interface.swift b/test/Interop/Cxx/stdlib/overlay/custom-collection-module-interface.swift
@@ -20,6 +20,13 @@
 // CHECK:   typealias RawIterator = SimpleCollectionReadOnly.iterator
 // CHECK: }
 
+// CHECK: struct SimpleCollectionReadWrite : CxxMutableRandomAccessCollection {
+// CHECK:   typealias Element = ConstRACIterator.Pointee
+// CHECK:   typealias Iterator = CxxIterator<SimpleCollectionReadWrite>
+// CHECK:   typealias RawIterator = SimpleCollectionReadWrite.const_iterator
+// CHECK:   typealias RawMutableIterator = SimpleCollectionReadWrite.iterator
+// CHECK: }
+
 // CHECK: struct HasInheritedTemplatedConstRACIterator<CInt> : CxxRandomAccessCollection {
 // CHECK:   typealias Element = InheritedTemplatedConstRACIterator<CInt>.Pointee
 // CHECK:   typealias Iterator = CxxIterator<HasInheritedTemplatedConstRACIterator<CInt>>
diff --git a/test/Interop/Cxx/stdlib/overlay/custom-collection.swift b/test/Interop/Cxx/stdlib/overlay/custom-collection.swift
@@ -37,6 +37,23 @@ CxxCollectionTestSuite.test("SimpleCollectionReadOnly as Swift.Collection") {
   expectEqual(slice.last, 3)
 }
 
+CxxCollectionTestSuite.test("SimpleCollectionReadWrite as Swift.MutableCollection") {
+  var c = SimpleCollectionReadWrite()
+  expectEqual(c.first, 1)
+  expectEqual(c.last, 5)
+
+  c.swapAt(0, 4)
+  expectEqual(c.first, 5)
+  expectEqual(c.last, 1)
+
+  c.reverse()
+  expectEqual(c[0], 1)
+  expectEqual(c[1], 4)
+  expectEqual(c[2], 3)
+  expectEqual(c[3], 2)
+  expectEqual(c[4], 5)
+}
+
 CxxCollectionTestSuite.test("SimpleArrayWrapper as Swift.Collection") {
   let c = SimpleArrayWrapper()
   expectEqual(c.first, 10)
diff --git a/test/Interop/Cxx/stdlib/use-std-vector.swift b/test/Interop/Cxx/stdlib/use-std-vector.swift
@@ -70,6 +70,36 @@ StdVectorTestSuite.test("VectorOfInt as ExpressibleByArrayLiteral") {
     expectEqual(v2[2], 3)
 }
 
+#if !os(Windows) // FIXME: rdar://113704853
+StdVectorTestSuite.test("VectorOfInt as MutableCollection") {
+    var v = Vector([2, 3, 1])
+    v.sort() // Swift function
+    expectEqual(v[0], 1)
+    expectEqual(v[1], 2)
+    expectEqual(v[2], 3)
+
+    v.reverse() // Swift function
+    expectEqual(v[0], 3)
+    expectEqual(v[1], 2)
+    expectEqual(v[2], 1)
+}
+
+StdVectorTestSuite.test("VectorOfString as MutableCollection") {
+    var v = VectorOfString([std.string("xyz"),
+                            std.string("abc"),
+                            std.string("ijk")])
+    v.swapAt(0, 2) // Swift function
+    expectEqual(v[0], std.string("ijk"))
+    expectEqual(v[1], std.string("abc"))
+    expectEqual(v[2], std.string("xyz"))
+
+    v.reverse() // Swift function
+    expectEqual(v[0], std.string("xyz"))
+    expectEqual(v[1], std.string("abc"))
+    expectEqual(v[2], std.string("ijk"))
+}
+#endif
+
 StdVectorTestSuite.test("VectorOfInt.push_back") {
     var v = Vector()
     let _42: CInt = 42
