[cxx-interop] Instantiate templated operator== for iterator types

C++ iterator types are often templated, and sometimes declare `operator==` as a non-member templated function. In libc++, an example of this is `__wrap_iter` which is used as an iterator type for `std::vector` and `std::string`.

We don't currently import templated non-member operators into Swift, however, we still want to support common C++ iterator patterns.

This change adds logic to instantiate templated non-member `operator==` for types that define `iterator_category` and are therefore likely to be valid iterator types.

rdar://97915515

Author: egorzhdan

lib/ClangImporter/ClangDerivedConformances.cpp

@@ -15,6 +15,8 @@
 #include "swift/AST/PrettyStackTrace.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/ClangImporter/ClangImporterRequests.h"
+#include "clang/Sema/DelayedDiagnostic.h"
+#include "clang/Sema/Overload.h"
 
 using namespace swift;
 using namespace swift::importer;
@@ -202,6 +204,44 @@ static ValueDecl *getPlusEqualOperator(NominalTypeDecl *decl, Type distanceTy) {
                         isValid);
 }
 
+static void instantiateTemplatedOperator(
+    ClangImporter::Implementation &impl,
+    const clang::ClassTemplateSpecializationDecl *classDecl,
+    clang::BinaryOperatorKind operatorKind) {
+
+  clang::ASTContext &clangCtx = impl.getClangASTContext();
+  clang::Sema &clangSema = impl.getClangSema();
+
+  clang::UnresolvedSet<1> ops;
+  auto qualType = clang::QualType(classDecl->getTypeForDecl(), 0);
+  auto arg = new (clangCtx)
+      clang::CXXThisExpr(clang::SourceLocation(), qualType, false);
+  arg->setType(clang::QualType(classDecl->getTypeForDecl(), 0));
+
+  clang::OverloadedOperatorKind opKind =
+      clang::BinaryOperator::getOverloadedOperator(operatorKind);
+  clang::OverloadCandidateSet candidateSet(
+      classDecl->getLocation(), clang::OverloadCandidateSet::CSK_Operator,
+      clang::OverloadCandidateSet::OperatorRewriteInfo(opKind, false));
+  clangSema.LookupOverloadedBinOp(candidateSet, opKind, ops, {arg, arg}, true);
+
+  clang::OverloadCandidateSet::iterator best;
+  switch (candidateSet.BestViableFunction(clangSema, clang::SourceLocation(),
+                                          best)) {
+  case clang::OR_Success: {
+    if (auto clangCallee = best->Function) {
+      auto lookupTable = impl.findLookupTable(classDecl);
+      addEntryToLookupTable(*lookupTable, clangCallee, impl.getNameImporter());
+    }
+    break;
+  }
+  case clang::OR_No_Viable_Function:
+  case clang::OR_Ambiguous:
+  case clang::OR_Deleted:
+    break;
+  }
+}
+
 bool swift::isIterator(const clang::CXXRecordDecl *clangDecl) {
   return getIteratorCategoryDecl(clangDecl);
 }
@@ -294,6 +334,13 @@ void swift::conformToCxxIteratorIfNeeded(
   if (!successorTy || successorTy->getAnyNominal() != decl)
     return;
 
+  // If this is a templated class, `operator==` might be templated as well.
+  // Try to instantiate it.
+  if (auto templateSpec =
+          dyn_cast<clang::ClassTemplateSpecializationDecl>(clangDecl)) {
+    instantiateTemplatedOperator(impl, templateSpec,
+                                 clang::BinaryOperatorKind::BO_EQ);
+  }
   // Check if present: `func ==`
   auto equalEqual = getEqualEqualOperator(decl);
   if (!equalEqual)
@@ -309,6 +356,11 @@ void swift::conformToCxxIteratorIfNeeded(
 
   // Try to conform to UnsafeCxxRandomAccessIterator if possible.
 
+  if (auto templateSpec =
+          dyn_cast<clang::ClassTemplateSpecializationDecl>(clangDecl)) {
+    instantiateTemplatedOperator(impl, templateSpec,
+                                 clang::BinaryOperatorKind::BO_Sub);
+  }
   auto minus = dyn_cast_or_null<FuncDecl>(getMinusOperator(decl));
   if (!minus)
     return;

lib/ClangImporter/ClangImporter.cpp

@@ -2661,7 +2661,17 @@ getClangOwningModule(ClangNode Node, const clang::ASTContext &ClangCtx) {
   if (const clang::Decl *D = Node.getAsDecl()) {
     auto ExtSource = ClangCtx.getExternalSource();
     assert(ExtSource);
-    return ExtSource->getModule(D->getOwningModuleID());
+
+    auto originalDecl = D;
+    if (auto functionDecl = dyn_cast<clang::FunctionDecl>(D)) {
+      if (auto pattern = functionDecl->getTemplateInstantiationPattern()) {
+        // Function template instantiations don't have an owning Clang module.
+        // Let's use the owning module of the template pattern.
+        originalDecl = pattern;
+      }
+    }
+
+    return ExtSource->getModule(originalDecl->getOwningModuleID());
   }
 
   if (const clang::ModuleMacro *M = Node.getAsModuleMacro())

test/Interop/Cxx/stdlib/libcxx-module-interface.swift

@@ -10,10 +10,10 @@
 
 // CHECK-IOSFWD: enum std {
 // CHECK-IOSFWD:   enum __1 {
-// CHECK-IOSFWD:     struct __CxxTemplateInstNSt3__112basic_stringIcNS_11char_traitsIcEENS_9allocatorIcEEEE {
+// CHECK-IOSFWD:     struct __CxxTemplateInstNSt3__112basic_stringIcNS_11char_traitsIcEENS_9allocatorIcEEEE : CxxRandomAccessCollection {
 // CHECK-IOSFWD:       typealias value_type = CChar
 // CHECK-IOSFWD:     }
-// CHECK-IOSFWD:     struct __CxxTemplateInstNSt3__112basic_stringIwNS_11char_traitsIwEENS_9allocatorIwEEEE {
+// CHECK-IOSFWD:     struct __CxxTemplateInstNSt3__112basic_stringIwNS_11char_traitsIwEENS_9allocatorIwEEEE : CxxRandomAccessCollection {
 // CHECK-IOSFWD:       typealias value_type = CWideChar
 // CHECK-IOSFWD:     }
 // CHECK-IOSFWD:     typealias string = std.__1.__CxxTemplateInstNSt3__112basic_stringIcNS_11char_traitsIcEENS_9allocatorIcEEEE

test/Interop/Cxx/stdlib/libstdcxx-module-interface.swift

@@ -11,10 +11,10 @@
 // REQUIRES: OS=linux-gnu
 
 // CHECK-STD: enum std {
-// CHECK-STRING:   struct {{__CxxTemplateInstNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE|__CxxTemplateInstSs}} {
+// CHECK-STRING:   struct {{__CxxTemplateInstNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE|__CxxTemplateInstSs}} : CxxRandomAccessCollection {
 // CHECK-STRING:     typealias value_type = std.__CxxTemplateInstSt11char_traitsIcE.char_type
 // CHECK-STRING:   }
-// CHECK-STRING:   struct {{__CxxTemplateInstNSt7__cxx1112basic_stringIwSt11char_traitsIwESaIwEEE|__CxxTemplateInstSbIwSt11char_traitsIwESaIwEE}} {
+// CHECK-STRING:   struct {{__CxxTemplateInstNSt7__cxx1112basic_stringIwSt11char_traitsIwESaIwEEE|__CxxTemplateInstSbIwSt11char_traitsIwESaIwEE}} : CxxRandomAccessCollection {
 // CHECK-STRING:     typealias value_type = std.__CxxTemplateInstSt11char_traitsIwE.char_type
 // CHECK-STRING:   }
 

test/Interop/Cxx/stdlib/overlay/Inputs/custom-iterator.h

@@ -496,4 +496,101 @@ struct TemplatedIterator {
 
 using TemplatedIteratorInt = TemplatedIterator<int>;
 
+template <typename T>
+struct TemplatedIteratorOutOfLineEq {
+  T value;
+
+  using iterator_category = std::input_iterator_tag;
+  using value_type = T;
+  using pointer = T *;
+  using reference = const T &;
+  using difference_type = int;
+
+  TemplatedIteratorOutOfLineEq(int value) : value(value) {}
+  TemplatedIteratorOutOfLineEq(const TemplatedIteratorOutOfLineEq &other) =
+      default;
+
+  const int &operator*() const { return value; }
+
+  TemplatedIteratorOutOfLineEq &operator++() {
+    value++;
+    return *this;
+  }
+  TemplatedIteratorOutOfLineEq operator++(int) {
+    auto tmp = TemplatedIteratorOutOfLineEq(value);
+    value++;
+    return tmp;
+  }
+};
+
+template <typename T>
+bool operator==(const TemplatedIteratorOutOfLineEq<T> &lhs,
+                const TemplatedIteratorOutOfLineEq<T> &rhs) {
+  return lhs.value == rhs.value;
+}
+
+using TemplatedIteratorOutOfLineEqInt = TemplatedIteratorOutOfLineEq<int>;
+
+template <typename T>
+struct TemplatedRACIteratorOutOfLineEq {
+  T value;
+  
+  using iterator_category = std::random_access_iterator_tag;
+  using value_type = T;
+  using pointer = T *;
+  using reference = const T &;
+  using difference_type = int;
+
+  TemplatedRACIteratorOutOfLineEq(int value) : value(value) {}
+  TemplatedRACIteratorOutOfLineEq(const TemplatedRACIteratorOutOfLineEq &other) =
+      default;
+
+  const int &operator*() const { return value; }
+
+  TemplatedRACIteratorOutOfLineEq &operator++() {
+    value++;
+    return *this;
+  }
+  TemplatedRACIteratorOutOfLineEq &operator--() {
+    value--;
+    return *this;
+  }
+  TemplatedRACIteratorOutOfLineEq operator++(int) {
+    auto tmp = TemplatedRACIteratorOutOfLineEq(value);
+    value++;
+    return tmp;
+  }
+
+  TemplatedRACIteratorOutOfLineEq &operator+=(difference_type v) {
+    value += v;
+    return *this;
+  }
+  TemplatedRACIteratorOutOfLineEq &operator-=(difference_type v) {
+    value -= v;
+    return *this;
+  }
+
+  TemplatedRACIteratorOutOfLineEq operator+(difference_type v) const {
+    return TemplatedRACIteratorOutOfLineEq(value + v);
+  }
+  TemplatedRACIteratorOutOfLineEq operator-(difference_type v) const {
+    return TemplatedRACIteratorOutOfLineEq(value - v);
+  }
+};
+
+template <typename T>
+bool operator==(const TemplatedRACIteratorOutOfLineEq<T> &lhs,
+                const TemplatedRACIteratorOutOfLineEq<T> &rhs) {
+  return lhs.value == rhs.value;
+}
+
+template <typename T>
+typename TemplatedRACIteratorOutOfLineEq<T>::difference_type
+operator-(const TemplatedRACIteratorOutOfLineEq<T> &lhs,
+          const TemplatedRACIteratorOutOfLineEq<T> &rhs) {
+  return lhs.value - rhs.value;
+}
+
+using TemplatedRACIteratorOutOfLineEqInt = TemplatedRACIteratorOutOfLineEq<int>;
+
 #endif // TEST_INTEROP_CXX_STDLIB_INPUTS_CUSTOM_ITERATOR_H

test/Interop/Cxx/stdlib/overlay/custom-iterator-module-interface.swift

@@ -92,3 +92,16 @@
 // CHECK:   typealias Pointee = Int32
 // CHECK:   static func == (lhs: __CxxTemplateInst17TemplatedIteratorIiE, other: __CxxTemplateInst17TemplatedIteratorIiE) -> Bool
 // CHECK: }
+
+// CHECK: struct __CxxTemplateInst28TemplatedIteratorOutOfLineEqIiE : UnsafeCxxInputIterator {
+// CHECK:   var pointee: Int32 { get }
+// CHECK:   func successor() -> __CxxTemplateInst28TemplatedIteratorOutOfLineEqIiE
+// CHECK:   typealias Pointee = Int32
+// CHECK: }
+
+// CHECK: struct __CxxTemplateInst31TemplatedRACIteratorOutOfLineEqIiE : UnsafeCxxRandomAccessIterator, UnsafeCxxInputIterator {
+// CHECK:   var pointee: Int32 { get }
+// CHECK:   func successor() -> __CxxTemplateInst31TemplatedRACIteratorOutOfLineEqIiE
+// CHECK:   typealias Pointee = Int32
+// CHECK:   typealias Distance = __CxxTemplateInst31TemplatedRACIteratorOutOfLineEqIiE.difference_type
+// CHECK: }

test/Interop/Cxx/stdlib/overlay/std-string-overlay.swift

@@ -37,22 +37,6 @@ StdStringOverlayTestSuite.test("std::string <=> Swift.String") {
   expectEqual(swift6, "xyz\0abc")
 }
 
-extension std.string.const_iterator: UnsafeCxxInputIterator {
-  // This func should not be required.
-  public static func ==(lhs: std.string.const_iterator,
-                        rhs: std.string.const_iterator) -> Bool {
-#if os(Linux)
-    // In libstdc++, `base()` returns UnsafePointer<Optional<UnsafePointer<CChar>>>.
-    return lhs.__baseUnsafe().pointee == rhs.__baseUnsafe().pointee
-#else
-    // In libc++, `base()` returns UnsafePointer<CChar>.
-    return lhs.__baseUnsafe() == rhs.__baseUnsafe()
-#endif
-  }
-}
-
-extension std.string: CxxSequence {}
-
 StdStringOverlayTestSuite.test("std::string as Swift.Sequence") {
   let cxx1 = std.string()
   var iterated = false

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

@@ -28,7 +28,6 @@ StdMapTestSuite.test("subscript") {
   expectEqual(m[3], 3)
 }
 
-extension Map.const_iterator : UnsafeCxxInputIterator { }
 extension Map : CxxSequence { }
 
 StdMapTestSuite.test("first(where:)") {

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

@@ -12,7 +12,6 @@ import Cxx
 
 var StdSetTestSuite = TestSuite("StdSet")
 
-extension SetOfCInt.const_iterator : UnsafeCxxInputIterator { }
 extension SetOfCInt : CxxSequence { }
 
 StdSetTestSuite.test("iterate") {

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

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift(-I %S/Inputs -Xfrontend -enable-experimental-cxx-interop)
+// RUN: %target-run-simple-swift(-I %S/Inputs -Xfrontend -enable-experimental-cxx-interop -Xfrontend -validate-tbd-against-ir=none)
 //
 // REQUIRES: executable_test
 //
@@ -11,11 +11,6 @@ import CxxStdlib.vector
 
 var StdVectorTestSuite = TestSuite("StdVector")
 
-extension Vector : RandomAccessCollection {
-  public var startIndex: Int { 0 }
-  public var endIndex: Int { size() }
-}
-
 StdVectorTestSuite.test("init") {
     let v = Vector()
     expectEqual(v.size(), 0)