Ensure that macro-provided operators can be found by witness lookup

When looking for an operator to satisfy a protocol requirement, we
currently depend on global operator lookup for everything except local
types. However, macro-provided operators aren't found by global
operator lookup, so perform a member lookup for such cases in addition
to global operator lookup. This makes macro-provided operators visible
through protocol requirements they witness.

Author: DougGregor

lib/Sema/TypeCheckProtocol.cpp

@@ -35,8 +35,10 @@
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/NameLookup.h"
+#include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/PrettyStackTrace.h"
+#include "swift/AST/PotentialMacroExpansions.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/TypeDeclFinder.h"
@@ -1318,6 +1320,25 @@ WitnessChecker::lookupValueWitnessesViaImplementsAttr(
   removeShadowedDecls(witnesses, DC);
 }
 
+/// Determine whether the given context may expand an operator with the given name.
+static bool contextMayExpandOperator(
+    DeclContext *dc, DeclBaseName operatorName
+) {
+  TypeOrExtensionDecl decl;
+  if (auto nominal = dyn_cast<NominalTypeDecl>(dc))
+    decl = nominal;
+  else if (auto ext = dyn_cast<ExtensionDecl>(dc))
+    decl = ext;
+  else
+    return false;
+
+  ASTContext &ctx = dc->getASTContext();
+  auto potentialExpansions = evaluateOrDefault(
+      ctx.evaluator, PotentialMacroExpansionsInContextRequest{decl},
+      PotentialMacroExpansions());
+  return potentialExpansions.shouldExpandForName(operatorName);
+}
+
 SmallVector<ValueDecl *, 4>
 WitnessChecker::lookupValueWitnesses(ValueDecl *req, bool *ignoringNames) {
   assert(!isa<AssociatedTypeDecl>(req) && "Not for lookup for type witnesses*");
@@ -1335,10 +1356,12 @@ WitnessChecker::lookupValueWitnesses(ValueDecl *req, bool *ignoringNames) {
   // An operator function is the only kind of witness that requires global
   // lookup. However, because global lookup doesn't enter local contexts,
   // an additional, qualified lookup is warranted when the conforming type
-  // is declared in a local context.
+  // is declared in a local context or when the operator could come from a
+  // macro expansion.
   const bool doUnqualifiedLookup = req->isOperator();
   const bool doQualifiedLookup =
-      !req->isOperator() || DC->getParent()->getLocalContext();
+      !req->isOperator() || DC->getParent()->getLocalContext() ||
+      contextMayExpandOperator(DC, req->getName().getBaseName());
 
   if (doUnqualifiedLookup) {
     auto lookup = TypeChecker::lookupUnqualified(DC->getModuleScopeContext(),

test/Macros/macro_expand.swift

@@ -465,14 +465,42 @@ struct HasEqualsSelf3 {
 struct HasEqualsSelf4 {
 }
 
+protocol SelfEqualsBoolProto { // expected-note 4{{where 'Self' =}}
+  static func ==(lhs: Self, rhs: Bool) -> Bool
+}
+
+struct HasEqualsSelfP: SelfEqualsBoolProto {
+  #addSelfEqualsOperator
+}
+
+struct HasEqualsSelf2P: SelfEqualsBoolProto {
+  #addSelfEqualsOperatorArbitrary
+}
+
+@AddSelfEqualsMemberOperator
+struct HasEqualsSelf3P: SelfEqualsBoolProto {
+}
+
+@AddSelfEqualsMemberOperatorArbitrary
+struct HasEqualsSelf4P: SelfEqualsBoolProto {
+}
+
 func testHasEqualsSelf(
-  x: HasEqualsSelf, y: HasEqualsSelf2, z: HasEqualsSelf3, w: HasEqualsSelf4
+  x: HasEqualsSelf, y: HasEqualsSelf2, z: HasEqualsSelf3, w: HasEqualsSelf4,
+  xP: HasEqualsSelfP, yP: HasEqualsSelf2P, zP: HasEqualsSelf3P,
+  wP: HasEqualsSelf4P
 ) {
 #if TEST_DIAGNOSTICS
   // Global operator lookup doesn't find member operators introduced by macros.
-  _ = (x == true) // expected-error{{cannot convert value of type}}
-  _ = (y == true) // expected-error{{cannot convert value of type}}
-  _ = (z == true) // expected-error{{cannot convert value of type}}
-  _ = (w == true) // expected-error{{cannot convert value of type}}
+  _ = (x == true) // expected-error{{referencing operator function '=='}}
+  _ = (y == true) // expected-error{{referencing operator function '=='}}
+  _ = (z == true) // expected-error{{referencing operator function '=='}}
+  _ = (w == true) // expected-error{{referencing operator function '=='}}
   #endif
+
+  // These should be found through the protocol.
+  _ = (xP == true)
+  _ = (yP == true)
+  _ = (zP == true)
+  _ = (wP == true)
 }