Merge pull request swiftlang#66320 from DougGregor/macro-operators-without-global-operator-lookup

Eliminate macro-generated operators from global operator lookup

Author: DougGregor

lib/AST/Module.cpp

@@ -249,7 +249,7 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
         // Cache the value under both its compound name and its full name.
         TopLevelValues.add(VD);
 
-        if (VD->getAttrs().hasAttribute<CustomAttr>()) {
+        if (!onlyOperators && VD->getAttrs().hasAttribute<CustomAttr>()) {
           MayHaveAuxiliaryDecls.push_back(VD);
         }
       }
@@ -279,8 +279,10 @@ void SourceLookupCache::addToUnqualifiedLookupCache(Range decls,
     else if (auto *PG = dyn_cast<PrecedenceGroupDecl>(D))
       PrecedenceGroups[PG->getName()].push_back(PG);
 
-    else if (auto *MED = dyn_cast<MacroExpansionDecl>(D))
-      MayHaveAuxiliaryDecls.push_back(MED);
+    else if (auto *MED = dyn_cast<MacroExpansionDecl>(D)) {
+      if (!onlyOperators)
+        MayHaveAuxiliaryDecls.push_back(MED);
+    }
   }
 }
 
@@ -403,13 +405,6 @@ void SourceLookupCache::populateAuxiliaryDeclCache() {
     for (auto macroNames : introducedNames) {
       auto macroRef = macroNames.getFirst();
       for (auto name : macroNames.getSecond()) {
-
-        // If this macro isn't in a module-scope context, and the introduced
-        // name isn't an operator, we shouldn't be able to see it.
-        if (!decl->getDeclContext()->isModuleScopeContext() &&
-            !name.getBaseName().isOperator())
-          continue;
-
         auto *placeholder = MissingDecl::forUnexpandedMacro(macroRef, decl);
         name.addToLookupTable(TopLevelAuxiliaryDecls, placeholder);
       }
@@ -492,12 +487,6 @@ void SourceLookupCache::lookupValue(DeclName Name, NLKind LookupKind,
   for (auto *unexpandedDecl : unexpandedDecls) {
     unexpandedDecl->forEachMacroExpandedDecl(
         [&](ValueDecl *decl) {
-          // If the declaration is not a module-scope declaration, and
-          // isn't an operator, ignore it.
-          if (!decl->getDeclContext()->isModuleScopeContext() &&
-              !decl->getName().getBaseName().isOperator())
-            return;
-
           if (decl->getName().matchesRef(Name)) {
             if (macroExpandedDecls.insert(decl).second)
               Result.push_back(decl);

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"
@@ -1316,6 +1318,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*");
@@ -1333,10 +1354,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,18 +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
 ) {
-  _ = (x == true)
-  _ = (y == true)
 #if TEST_DIAGNOSTICS
-  // FIXME: This is technically a bug, because we should be able to find the
-  // == operator introduced through a member operator. However, we might
-  // want to change the rule rather than implement this.
-  _ = (z == true) // expected-error{{binary operator '==' cannot be applied to operands}}
-  // expected-note@-1{{overloads for '==' exist with these partially matching parameter lists}}
-  _ = (w == true) // expected-error{{binary operator '==' cannot be applied to operands}}
-  // expected-note@-1{{overloads for '==' exist with these partially matching parameter lists}}
+  // Global operator lookup doesn't find member operators introduced by macros.
+  _ = (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)
 }