[cxx-interop] Allow removing elements from std::vector

This adds `func remove(at index: Int)` to all instantiations of `std::vector` via an extension for `protocol CxxVector`.

The original C++ method `std::vector::erase` is not visible in Swift because all of its overloads return unsafe iterators.

rdar://113704853

Author: egorzhdan

lib/ClangImporter/ClangDerivedConformances.cpp

@@ -990,15 +990,20 @@ void swift::conformToCxxVectorIfNeeded(ClangImporter::Implementation &impl,
       decl, ctx.getIdentifier("value_type"));
   auto iterType = lookupDirectSingleWithoutExtensions<TypeAliasDecl>(
       decl, ctx.getIdentifier("const_iterator"));
-  if (!valueType || !iterType)
+  auto mutableIterType = lookupDirectSingleWithoutExtensions<TypeAliasDecl>(
+      decl, ctx.getIdentifier("iterator"));
+  if (!valueType || !iterType || !mutableIterType)
     return;
 
   ProtocolDecl *cxxRandomAccessIteratorProto =
       ctx.getProtocol(KnownProtocolKind::UnsafeCxxRandomAccessIterator);
-  if (!cxxRandomAccessIteratorProto)
+  ProtocolDecl *cxxMutableRandomAccessIteratorProto =
+      ctx.getProtocol(KnownProtocolKind::UnsafeCxxMutableRandomAccessIterator);
+  if (!cxxRandomAccessIteratorProto || !cxxMutableRandomAccessIteratorProto)
     return;
 
   auto rawIteratorTy = iterType->getUnderlyingType();
+  auto rawMutableIteratorTy = mutableIterType->getUnderlyingType();
 
   // Check if RawIterator conforms to UnsafeCxxRandomAccessIterator.
   ModuleDecl *module = decl->getModuleContext();
@@ -1007,11 +1012,19 @@ void swift::conformToCxxVectorIfNeeded(ClangImporter::Implementation &impl,
   if (!isConcreteAndValid(rawIteratorConformanceRef, module))
     return;
 
+  // Check if RawMutableIterator conforms to UnsafeCxxMutableInputIterator.
+  auto rawMutableIteratorConformanceRef = module->lookupConformance(
+      rawMutableIteratorTy, cxxMutableRandomAccessIteratorProto);
+  if (!isConcreteAndValid(rawMutableIteratorConformanceRef, module))
+    return;
+
   impl.addSynthesizedTypealias(decl, ctx.Id_Element,
                                valueType->getUnderlyingType());
   impl.addSynthesizedTypealias(decl, ctx.Id_ArrayLiteralElement,
                                valueType->getUnderlyingType());
   impl.addSynthesizedTypealias(decl, ctx.getIdentifier("RawIterator"),
                                rawIteratorTy);
+  impl.addSynthesizedTypealias(decl, ctx.getIdentifier("RawMutableIterator"),
+                               rawMutableIteratorTy);
   impl.addSynthesizedProtocolAttrs(decl, {KnownProtocolKind::CxxVector});
 }

stdlib/public/Cxx/CxxVector.swift

@@ -17,10 +17,19 @@ public protocol CxxVector<Element>: ExpressibleByArrayLiteral {
   associatedtype Element
   associatedtype RawIterator: UnsafeCxxRandomAccessIterator
     where RawIterator.Pointee == Element
+  associatedtype RawMutableIterator: UnsafeCxxMutableRandomAccessIterator
+    where RawMutableIterator.Pointee == Element
 
   init()
 
+  /// Do not implement this function manually in Swift.
+  mutating func __beginUnsafe() -> RawIterator
+
   mutating func push_back(_ element: Element)
+
+  /// Do not implement this function manually in Swift.
+  @discardableResult
+  mutating func __eraseUnsafe(_ iterator: RawIterator) -> RawMutableIterator
 }
 
 extension CxxVector {
@@ -37,11 +46,20 @@ extension CxxVector {
       self.push_back(item)
     }
   }
-}
 
-extension CxxVector {
   @inlinable
   public init(arrayLiteral elements: Element...) {
     self.init(elements)
   }
+
+  @discardableResult
+  @inlinable
+  public mutating func remove(at index: Int) -> Element {
+    // Not using CxxIterator here to avoid making a copy of the collection.
+    var rawIterator = self.__beginUnsafe()
+    rawIterator += RawIterator.Distance(index)
+    let element = rawIterator.pointee
+    self.__eraseUnsafe(rawIterator)
+    return element
+  }
 }

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

@@ -73,6 +73,48 @@ func fill(vector v: inout Vector) {
     v.push_back(CInt(3))
 }
 
+StdVectorTestSuite.test("VectorOfInt.remove(at:)") {
+    var v = Vector()
+    fill(vector: &v)
+
+    let rm1 = v.remove(at: 1)
+    expectEqual(rm1, 2)
+    expectEqual(v.size(), 2)
+    expectEqual(v[0], 1)
+    expectEqual(v[1], 3)
+
+    let rm2 = v.remove(at: 0)
+    expectEqual(rm2, 1)
+    expectEqual(v.size(), 1)
+    expectEqual(v[0], 3)
+}
+
+StdVectorTestSuite.test("VectorOfString.remove(at:)") {
+    var v = VectorOfString()
+    v.push_back(std.string())
+    v.push_back(std.string("123"))
+    v.push_back(std.string("abc"))
+    v.push_back(std.string("qwe"))
+
+    let rm1 = v.remove(at: 3)
+    expectEqual(rm1, std.string("qwe"))
+    expectEqual(v.size(), 3)
+    expectEqual(v[0], std.string())
+    expectEqual(v[1], std.string("123"))
+    expectEqual(v[2], std.string("abc"))
+
+    let rm2 = v.remove(at: 1)
+    expectEqual(rm2, std.string("123"))
+    expectEqual(v.size(), 2)
+
+    v.remove(at: 0)
+    expectEqual(v.size(), 1)
+
+    v.remove(at: 0)
+    expectEqual(v.size(), 0)
+    expectTrue(v.empty())
+}
+
 StdVectorTestSuite.test("VectorOfInt for loop") {
     var v = Vector()
     fill(vector: &v)