[cxx-interop] Add UnsafeCxxContiguousIterator & UnsafeCxxMutableContiguousIterator protocols

egorzhdan · egorzhdan · commit cbd533db68e7 · 2024-10-14T16:05:26.000+01:00
This adds a pair of Swift protocols that represents C++ iterator types conforming to `std::contiguous_iterator_tag` requirements. These are random access iterators that guarantee that the values are stored in consequent memory addresses.

This will be used to optimize usage of C++ containers such as `std::vector` from Swift, for instance, by providing an overload of `withContiguousStorageIfAvailable` for contiguous containers.
diff --git a/include/swift/AST/KnownProtocols.def b/include/swift/AST/KnownProtocols.def
@@ -142,6 +142,8 @@ PROTOCOL(UnsafeCxxInputIterator)
 PROTOCOL(UnsafeCxxMutableInputIterator)
 PROTOCOL(UnsafeCxxRandomAccessIterator)
 PROTOCOL(UnsafeCxxMutableRandomAccessIterator)
+PROTOCOL(UnsafeCxxContiguousIterator)
+PROTOCOL(UnsafeCxxMutableContiguousIterator)
 
 PROTOCOL(AsyncSequence)
 PROTOCOL(AsyncIteratorProtocol)
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -1444,6 +1444,8 @@ ProtocolDecl *ASTContext::getProtocol(KnownProtocolKind kind) const {
   case KnownProtocolKind::UnsafeCxxMutableInputIterator:
   case KnownProtocolKind::UnsafeCxxRandomAccessIterator:
   case KnownProtocolKind::UnsafeCxxMutableRandomAccessIterator:
+  case KnownProtocolKind::UnsafeCxxContiguousIterator:
+  case KnownProtocolKind::UnsafeCxxMutableContiguousIterator:
     M = getLoadedModule(Id_Cxx);
     break;
   case KnownProtocolKind::Copyable:
diff --git a/lib/ClangImporter/ClangDerivedConformances.cpp b/lib/ClangImporter/ClangDerivedConformances.cpp
@@ -462,19 +462,28 @@ void swift::conformToCxxIteratorIfNeeded(
   auto isRandomAccessIteratorDecl = [&](const clang::CXXRecordDecl *base) {
     return isIteratorCategoryDecl(base, "random_access_iterator_tag");
   };
+  auto isContiguousIteratorDecl = [&](const clang::CXXRecordDecl *base) {
+    return isIteratorCategoryDecl(base, "contiguous_iterator_tag"); // C++20
+  };
 
   // Traverse all transitive bases of `underlyingDecl` to check if
   // it inherits from `std::input_iterator_tag`.
   bool isInputIterator = isInputIteratorDecl(underlyingCategoryDecl);
   bool isRandomAccessIterator =
       isRandomAccessIteratorDecl(underlyingCategoryDecl);
+  bool isContiguousIterator = isContiguousIteratorDecl(underlyingCategoryDecl);
   underlyingCategoryDecl->forallBases([&](const clang::CXXRecordDecl *base) {
     if (isInputIteratorDecl(base)) {
       isInputIterator = true;
     }
     if (isRandomAccessIteratorDecl(base)) {
       isRandomAccessIterator = true;
       isInputIterator = true;
+    }
+    if (isContiguousIteratorDecl(base)) {
+      isContiguousIterator = true;
+      isRandomAccessIterator = true;
+      isInputIterator = true;
       return false;
     }
     return true;
@@ -594,6 +603,15 @@ void swift::conformToCxxIteratorIfNeeded(
   else
     impl.addSynthesizedProtocolAttrs(
         decl, {KnownProtocolKind::UnsafeCxxRandomAccessIterator});
+
+  if (isContiguousIterator) {
+    if (pointeeSettable)
+      impl.addSynthesizedProtocolAttrs(
+          decl, {KnownProtocolKind::UnsafeCxxMutableContiguousIterator});
+    else
+      impl.addSynthesizedProtocolAttrs(
+          decl, {KnownProtocolKind::UnsafeCxxContiguousIterator});
+  }
 }
 
 void swift::conformToCxxConvertibleToBoolIfNeeded(
diff --git a/lib/IRGen/GenMeta.cpp b/lib/IRGen/GenMeta.cpp
@@ -6967,6 +6967,8 @@ SpecialProtocol irgen::getSpecialProtocolID(ProtocolDecl *P) {
   case KnownProtocolKind::UnsafeCxxMutableInputIterator:
   case KnownProtocolKind::UnsafeCxxRandomAccessIterator:
   case KnownProtocolKind::UnsafeCxxMutableRandomAccessIterator:
+  case KnownProtocolKind::UnsafeCxxContiguousIterator:
+  case KnownProtocolKind::UnsafeCxxMutableContiguousIterator:
   case KnownProtocolKind::Executor:
   case KnownProtocolKind::SerialExecutor:
   case KnownProtocolKind::TaskExecutor:
diff --git a/stdlib/public/Cxx/UnsafeCxxIterators.swift b/stdlib/public/Cxx/UnsafeCxxIterators.swift
@@ -87,3 +87,8 @@ public protocol UnsafeCxxMutableRandomAccessIterator:
 UnsafeCxxRandomAccessIterator, UnsafeCxxMutableInputIterator {}
 
 extension UnsafeMutablePointer: UnsafeCxxMutableRandomAccessIterator {}
+
+public protocol UnsafeCxxContiguousIterator: UnsafeCxxRandomAccessIterator {}
+
+public protocol UnsafeCxxMutableContiguousIterator:
+UnsafeCxxContiguousIterator, UnsafeCxxMutableRandomAccessIterator {}
diff --git a/test/Interop/Cxx/stdlib/overlay/Inputs/custom-iterator.h b/test/Interop/Cxx/stdlib/overlay/Inputs/custom-iterator.h
@@ -286,6 +286,181 @@ struct HasTypedefIteratorTag {
   }
 };
 
+#if __cplusplus >= 202002L
+struct ConstContiguousIterator {
+private:
+  const int *value;
+
+public:
+  using iterator_category = std::contiguous_iterator_tag;
+  using value_type = int;
+  using pointer = int *;
+  using reference = const int &;
+  using difference_type = int;
+
+  ConstContiguousIterator(const int *value) : value(value) {}
+  ConstContiguousIterator(const ConstContiguousIterator &other) = default;
+
+  const int &operator*() const { return *value; }
+
+  ConstContiguousIterator &operator++() {
+    value++;
+    return *this;
+  }
+  ConstContiguousIterator operator++(int) {
+    auto tmp = ConstContiguousIterator(value);
+    value++;
+    return tmp;
+  }
+
+  void operator+=(difference_type v) { value += v; }
+  void operator-=(difference_type v) { value -= v; }
+  ConstContiguousIterator operator+(difference_type v) const {
+    return ConstContiguousIterator(value + v);
+  }
+  ConstContiguousIterator operator-(difference_type v) const {
+    return ConstContiguousIterator(value - v);
+  }
+  friend ConstContiguousIterator operator+(difference_type v,
+                                           const ConstContiguousIterator &it) {
+    return it + v;
+  }
+  int operator-(const ConstContiguousIterator &other) const {
+    return value - other.value;
+  }
+
+  bool operator<(const ConstContiguousIterator &other) const {
+    return value < other.value;
+  }
+
+  bool operator==(const ConstContiguousIterator &other) const {
+    return value == other.value;
+  }
+  bool operator!=(const ConstContiguousIterator &other) const {
+    return value != other.value;
+  }
+};
+
+struct HasCustomContiguousIteratorTag {
+private:
+  const int *value;
+
+public:
+  struct CustomTag : std::contiguous_iterator_tag {};
+  using iterator_category = CustomTag;
+  using value_type = int;
+  using pointer = int *;
+  using reference = const int &;
+  using difference_type = int;
+
+  HasCustomContiguousIteratorTag(const int *value) : value(value) {}
+  HasCustomContiguousIteratorTag(const HasCustomContiguousIteratorTag &other) =
+      default;
+
+  const int &operator*() const { return *value; }
+
+  HasCustomContiguousIteratorTag &operator++() {
+    value++;
+    return *this;
+  }
+  HasCustomContiguousIteratorTag operator++(int) {
+    auto tmp = HasCustomContiguousIteratorTag(value);
+    value++;
+    return tmp;
+  }
+
+  void operator+=(difference_type v) { value += v; }
+  void operator-=(difference_type v) { value -= v; }
+  HasCustomContiguousIteratorTag operator+(difference_type v) const {
+    return HasCustomContiguousIteratorTag(value + v);
+  }
+  HasCustomContiguousIteratorTag operator-(difference_type v) const {
+    return HasCustomContiguousIteratorTag(value - v);
+  }
+  friend HasCustomContiguousIteratorTag
+  operator+(difference_type v, const HasCustomContiguousIteratorTag &it) {
+    return it + v;
+  }
+  int operator-(const HasCustomContiguousIteratorTag &other) const {
+    return value - other.value;
+  }
+
+  bool operator<(const HasCustomContiguousIteratorTag &other) const {
+    return value < other.value;
+  }
+
+  bool operator==(const HasCustomContiguousIteratorTag &other) const {
+    return value == other.value;
+  }
+  bool operator!=(const HasCustomContiguousIteratorTag &other) const {
+    return value != other.value;
+  }
+};
+
+struct ConstContiguousIteratorInheritedFromConstRACIterator : ConstRACIterator {
+  using iterator_category = std::contiguous_iterator_tag;
+};
+
+struct MutableContiguousIterator {
+private:
+  int *value;
+
+public:
+  using iterator_category = std::contiguous_iterator_tag;
+  using value_type = int;
+  using pointer = int *;
+  using reference = const int &;
+  using difference_type = int;
+
+  MutableContiguousIterator(int *value) : value(value) {}
+  MutableContiguousIterator(const MutableRACIterator &other) = default;
+
+  const int &operator*() const { return *value; }
+  int &operator*() { return *value; }
+
+  MutableContiguousIterator &operator++() {
+    value++;
+    return *this;
+  }
+  MutableContiguousIterator operator++(int) {
+    auto tmp = MutableRACIterator(value);
+    value++;
+    return tmp;
+  }
+
+  void operator+=(difference_type v) { value += v; }
+  void operator-=(difference_type v) { value -= v; }
+  MutableContiguousIterator operator+(difference_type v) const {
+    return MutableContiguousIterator(value + v);
+  }
+  MutableContiguousIterator operator-(difference_type v) const {
+    return MutableContiguousIterator(value - v);
+  }
+  friend MutableContiguousIterator
+  operator+(difference_type v, const MutableContiguousIterator &it) {
+    return it + v;
+  }
+  int operator-(const MutableContiguousIterator &other) const {
+    return value - other.value;
+  }
+
+  bool operator<(const MutableContiguousIterator &other) const {
+    return value < other.value;
+  }
+
+  bool operator==(const MutableContiguousIterator &other) const {
+    return value == other.value;
+  }
+  bool operator!=(const MutableContiguousIterator &other) const {
+    return value != other.value;
+  }
+};
+
+struct MutableContiguousInheritedFromMutableRACIterator : MutableRACIterator {
+  using iterator_category = std::contiguous_iterator_tag;
+};
+#endif
+
 // MARK: Types that are not actually iterators
 
 struct HasNoIteratorCategory {
diff --git a/test/Interop/Cxx/stdlib/overlay/custom-iterator-module-interface.swift b/test/Interop/Cxx/stdlib/overlay/custom-iterator-module-interface.swift
@@ -1,6 +1,8 @@
 // RUN: %target-swift-ide-test -print-module -module-to-print=CustomIterator -source-filename=x -I %S/Inputs -enable-experimental-cxx-interop | %FileCheck %s
 // RUN: %target-swift-ide-test -print-module -module-to-print=CustomIterator -source-filename=x -I %S/Inputs -cxx-interoperability-mode=swift-6 | %FileCheck %s
+// RUN: %target-swift-ide-test -print-module -module-to-print=CustomIterator -source-filename=x -I %S/Inputs -cxx-interoperability-mode=swift-6 -Xcc -std=c++20 | %FileCheck %s --check-prefix=CHECK-CXX20
 // RUN: %target-swift-ide-test -print-module -module-to-print=CustomIterator -source-filename=x -I %S/Inputs -cxx-interoperability-mode=upcoming-swift | %FileCheck %s
+// RUN: %target-swift-ide-test -print-module -module-to-print=CustomIterator -source-filename=x -I %S/Inputs -cxx-interoperability-mode=upcoming-swift -Xcc -std=c++20 | %FileCheck %s --check-prefix=CHECK-CXX20 
 
 // CHECK: struct ConstIterator : UnsafeCxxInputIterator {
 // CHECK:   func successor() -> ConstIterator
@@ -80,6 +82,41 @@
 // CHECK:   static func == (lhs: HasTypedefIteratorTag, other: HasTypedefIteratorTag) -> Bool
 // CHECK: }
 
+// CHECK-CXX20: struct ConstContiguousIterator : UnsafeCxxContiguousIterator, UnsafeCxxRandomAccessIterator, UnsafeCxxInputIterator {
+// CHECK-CXX20:   func successor() -> ConstContiguousIterator
+// CHECK-CXX20:   var pointee: Int32
+// CHECK-CXX20:   typealias Pointee = Int32
+// CHECK-CXX20:   typealias Distance = Int32
+// CHECK-CXX20: }
+
+// CHECK-CXX20: struct HasCustomContiguousIteratorTag : UnsafeCxxContiguousIterator, UnsafeCxxRandomAccessIterator, UnsafeCxxInputIterator {
+// CHECK-CXX20:   func successor() -> HasCustomContiguousIteratorTag
+// CHECK-CXX20:   var pointee: Int32
+// CHECK-CXX20:   typealias Pointee = Int32
+// CHECK-CXX20:   typealias Distance = Int32
+// CHECK-CXX20: }
+
+// CHECK-CXX20: struct ConstContiguousIteratorInheritedFromConstRACIterator : UnsafeCxxContiguousIterator, UnsafeCxxRandomAccessIterator, UnsafeCxxInputIterator {
+// CHECK-CXX20:   func successor() -> ConstContiguousIteratorInheritedFromConstRACIterator
+// CHECK-CXX20:   var pointee: Int32
+// CHECK-CXX20:   typealias Pointee = Int32
+// CHECK-CXX20:   typealias Distance = Int32
+// CHECK-CXX20: }
+
+// CHECK-CXX20: struct MutableContiguousIterator : UnsafeCxxMutableContiguousIterator, UnsafeCxxMutableRandomAccessIterator, UnsafeCxxMutableInputIterator {
+// CHECK-CXX20:   func successor() -> MutableContiguousIterator
+// CHECK-CXX20:   var pointee: Int32
+// CHECK-CXX20:   typealias Pointee = Int32
+// CHECK-CXX20:   typealias Distance = Int32
+// CHECK-CXX20: }
+
+// CHECK-CXX20: struct MutableContiguousInheritedFromMutableRACIterator : UnsafeCxxMutableContiguousIterator, UnsafeCxxMutableRandomAccessIterator, UnsafeCxxMutableInputIterator {
+// CHECK-CXX20:   func successor() -> MutableContiguousInheritedFromMutableRACIterator
+// CHECK-CXX20:   var pointee: Int32
+// CHECK-CXX20:   typealias Pointee = Int32
+// CHECK-CXX20:   typealias Distance = Int32
+// CHECK-CXX20: }
+
 // CHECK-NOT: struct HasNoIteratorCategory : UnsafeCxxInputIterator
 // CHECK-NOT: struct HasInvalidIteratorCategory : UnsafeCxxInputIterator
 // CHECK-NOT: struct HasNoEqualEqual : UnsafeCxxInputIterator
