Merge pull request #60155 from apple/egorzhdan/cxx-iterator-ool-equals

[cxx-interop] Allow iterators with out-of-class `operator==`

Author: egorzhdan

lib/ClangImporter/ClangDerivedConformances.cpp

@@ -28,6 +28,51 @@ getIteratorCategoryDecl(const clang::CXXRecordDecl *clangDecl) {
   return dyn_cast_or_null<clang::TypeDecl>(iteratorCategory);
 }
 
+static ValueDecl *getEqualEqualOperator(NominalTypeDecl *decl) {
+  auto id = decl->getASTContext().Id_EqualsOperator;
+
+  auto isValid = [&](ValueDecl *equalEqualOp) -> bool {
+    auto equalEqual = dyn_cast<FuncDecl>(equalEqualOp);
+    if (!equalEqual || !equalEqual->hasParameterList())
+      return false;
+    auto params = equalEqual->getParameters();
+    if (params->size() != 2)
+      return false;
+    auto lhs = params->get(0);
+    auto rhs = params->get(1);
+    if (lhs->isInOut() || rhs->isInOut())
+      return false;
+    auto lhsTy = lhs->getType();
+    auto rhsTy = rhs->getType();
+    if (!lhsTy || !rhsTy)
+      return false;
+    auto lhsNominal = lhsTy->getAnyNominal();
+    auto rhsNominal = rhsTy->getAnyNominal();
+    if (lhsNominal != rhsNominal || lhsNominal != decl)
+      return false;
+    return true;
+  };
+
+  // First look for `func ==` declared as a member.
+  auto memberResults = decl->lookupDirect(id);
+  for (const auto &member : memberResults) {
+    if (isValid(member))
+      return member;
+  }
+
+  // If no member `func ==` was found, look for out-of-class definitions in the
+  // same module.
+  auto module = decl->getModuleContext();
+  llvm::SmallVector<ValueDecl *> nonMemberResults;
+  module->lookupValue(id, NLKind::UnqualifiedLookup, nonMemberResults);
+  for (const auto &nonMember : nonMemberResults) {
+    if (isValid(nonMember))
+      return nonMember;
+  }
+
+  return nullptr;
+}
+
 bool swift::isIterator(const clang::CXXRecordDecl *clangDecl) {
   return getIteratorCategoryDecl(clangDecl);
 }
@@ -103,24 +148,8 @@ void swift::conformToCxxIteratorIfNeeded(
     return;
 
   // Check if present: `func ==`
-  // FIXME: this only detects `operator==` declared as a member.
-  auto equalEquals = decl->lookupDirect(ctx.Id_EqualsOperator);
-  if (equalEquals.empty())
-    return;
-  auto equalEqual = dyn_cast<FuncDecl>(equalEquals.front());
-  if (!equalEqual || !equalEqual->hasParameterList())
-    return;
-  auto equalEqualParams = equalEqual->getParameters();
-  if (equalEqualParams->size() != 2)
-    return;
-  auto equalEqualLHS = equalEqualParams->get(0);
-  auto equalEqualRHS = equalEqualParams->get(1);
-  if (equalEqualLHS->isInOut() || equalEqualRHS->isInOut())
-    return;
-  auto equalEqualLHSTy = equalEqualLHS->getType();
-  auto equalEqualRHSTy = equalEqualRHS->getType();
-  if (!equalEqualLHSTy || !equalEqualRHSTy ||
-      equalEqualLHSTy->getAnyNominal() != equalEqualRHSTy->getAnyNominal())
+  auto equalEqual = getEqualEqualOperator(decl);
+  if (!equalEqual)
     return;
 
   impl.addSynthesizedTypealias(decl, ctx.getIdentifier("Pointee"),

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

@@ -191,6 +191,19 @@ struct HasNoEqualEqual {
   }
 };
 
+struct HasInvalidEqualEqual {
+  int value;
+  using iterator_category = std::input_iterator_tag;
+  const int &operator*() const { return value; }
+  HasInvalidEqualEqual &operator++() {
+    value++;
+    return *this;
+  }
+  bool operator==(const int &other) const { // wrong type
+    return value == other;
+  }
+};
+
 struct HasNoIncrementOperator {
   int value;
   using iterator_category = std::input_iterator_tag;

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

@@ -9,8 +9,7 @@
 // CHECK:   static func == (lhs: ConstIterator, other: ConstIterator) -> Bool
 // CHECK: }
 
-// TODO: ConstIteratorOutOfLineEq should also conform to UnsafeCxxInputIterator.
-// CHECK: struct ConstIteratorOutOfLineEq {
+// CHECK: struct ConstIteratorOutOfLineEq : UnsafeCxxInputIterator {
 // CHECK:   var pointee: Int32 { get }
 // CHECK:   func successor() -> ConstIteratorOutOfLineEq
 // CHECK: }
@@ -54,6 +53,7 @@
 // CHECK-NOT: struct HasNoIteratorCategory : UnsafeCxxInputIterator
 // CHECK-NOT: struct HasInvalidIteratorCategory : UnsafeCxxInputIterator
 // CHECK-NOT: struct HasNoEqualEqual : UnsafeCxxInputIterator
+// CHECK-NOT: struct HasInvalidEqualEqual : UnsafeCxxInputIterator
 // CHECK-NOT: struct HasNoIncrementOperator : UnsafeCxxInputIterator
 // CHECK-NOT: struct HasNoPreIncrementOperator : UnsafeCxxInputIterator
 // CHECK-NOT: struct HasNoDereferenceOperator : UnsafeCxxInputIterator

test/Interop/Cxx/stdlib/overlay/custom-sequence-typechecker.swift

@@ -18,8 +18,7 @@ func checkSimpleSequence() {
 }
 
 // === SimpleSequenceWithOutOfLineEqualEqual ===
-// TODO: synthesize conformance to UnsafeCxxInputIterator.
-//extension SimpleSequenceWithOutOfLineEqualEqual : CxxSequence {}
+extension SimpleSequenceWithOutOfLineEqualEqual : CxxSequence {}
 
 // === SimpleArrayWrapper ===
 // No UnsafeCxxInputIterator conformance required, since the iterators are actually UnsafePointers here.