Fix precedencegroup and operator decl lookup

Re-implement operator and precedencegroup decl
lookup to use `namelookup::getAllImports` and
existing decl shadowing logic. This allows us to
find operator decls through `@_exported` imports,
prefer operator decls defined in the same module
over imported decls, and fixes a couple of other
subtle issues.

Because this new implementation is technically
source breaking, as we can find multiple results
where we used to only find one result, it's placed
behind the new Frontend flag
`-enable-new-operator-lookup` (with the aim of
enabling it by default when we get a new language
mode).

However the new logic will always be used if the
result is unambiguous. This means that e.g
`@_exported` operators will be instantly available
as long as there's only one candidate. If multiple
candidates are found, we fall back to the old
logic.

Resolves SR-12132.
Resolves rdar://59198796.

Author: hamishknight

include/swift/AST/Decl.h

@@ -7273,12 +7273,6 @@ class InfixOperatorDecl : public OperatorDecl {
 
   PrecedenceGroupDecl *getPrecedenceGroup() const;
 
-  /// True if this decl's attributes conflict with those declared by another
-  /// operator.
-  bool conflictsWith(InfixOperatorDecl *other) {
-    return getPrecedenceGroup() != other->getPrecedenceGroup();
-  }
-  
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::InfixOperator;
   }
@@ -7307,12 +7301,6 @@ class PrefixOperatorDecl : public OperatorDecl {
     return { getOperatorLoc(), getNameLoc() };
   }
 
-  /// True if this decl's attributes conflict with those declared by another
-  /// PrefixOperatorDecl.
-  bool conflictsWith(PrefixOperatorDecl *other) {
-    return false;
-  }
-  
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::PrefixOperator;
   }
@@ -7340,12 +7328,6 @@ class PostfixOperatorDecl : public OperatorDecl {
   SourceRange getSourceRange() const {
     return { getOperatorLoc(), getNameLoc() };
   }
-
-  /// True if this decl's attributes conflict with those declared by another
-  /// PostfixOperatorDecl.
-  bool conflictsWith(PostfixOperatorDecl *other) {
-    return false;
-  }
 
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::PostfixOperator;

include/swift/AST/DeclContext.h

@@ -54,13 +54,19 @@ namespace swift {
   class GenericSignature;
   class GenericTypeParamDecl;
   class GenericTypeParamType;
+  class InfixOperatorDecl;
+  class InfixOperatorLookupResult;
+  class PrecedenceGroupDecl;
   class ProtocolDecl;
   class Requirement;
   class SourceFile;
   class Type;
   class ModuleDecl;
   class GenericTypeDecl;
   class NominalTypeDecl;
+  class PrecedenceGroupLookupResult;
+  class PostfixOperatorDecl;
+  class PrefixOperatorDecl;
   class ProtocolConformance;
   class ValueDecl;
   class Initializer;
@@ -561,6 +567,29 @@ class alignas(1 << DeclContextAlignInBits) DeclContext {
          ObjCSelector selector,
          SmallVectorImpl<AbstractFunctionDecl *> &results) const;
 
+  /// Looks up an infix operator with a given \p name.
+  ///
+  /// This returns a vector of results, as it's possible to find multiple infix
+  /// operators with different precedence groups.
+  InfixOperatorLookupResult lookupInfixOperator(Identifier name) const;
+
+  /// Looks up an prefix operator with a given \p name.
+  ///
+  /// If multiple results are found, one is chosen in a stable manner, as
+  /// prefix operator decls cannot differ other than in name. If no results are
+  /// found, returns \c nullptr.
+  PrefixOperatorDecl *lookupPrefixOperator(Identifier name) const;
+
+  /// Looks up an postfix operator with a given \p name.
+  ///
+  /// If multiple results are found, one is chosen in a stable manner, as
+  /// postfix operator decls cannot differ other than in name. If no results are
+  /// found, returns \c nullptr.
+  PostfixOperatorDecl *lookupPostfixOperator(Identifier name) const;
+
+  /// Looks up a precedence group with a given \p name.
+  PrecedenceGroupLookupResult lookupPrecedenceGroup(Identifier name) const;
+
   /// Return the ASTContext for a specified DeclContext by
   /// walking up to the enclosing module and returning its ASTContext.
   LLVM_READONLY

include/swift/AST/Module.h

@@ -592,21 +592,6 @@ class ModuleDecl : public DeclContext, public TypeDecl {
   /// FIXME: Remove the integrated REPL.
   void clearLookupCache();
 
-  /// @{
-
-  /// Look up the given operator in this module.
-  ///
-  /// If the operator is not found, or if there is an ambiguity, returns null.
-  InfixOperatorDecl *lookupInfixOperator(Identifier name,
-                                         SourceLoc diagLoc = {});
-  PrefixOperatorDecl *lookupPrefixOperator(Identifier name,
-                                           SourceLoc diagLoc = {});
-  PostfixOperatorDecl *lookupPostfixOperator(Identifier name,
-                                             SourceLoc diagLoc = {});
-  PrecedenceGroupDecl *lookupPrecedenceGroup(Identifier name,
-                                             SourceLoc diagLoc = {});
-  /// @}
-
   /// Finds all class members defined in this module.
   ///
   /// This does a simple local lookup, not recursively looking through imports.

include/swift/AST/NameLookupRequests.h

@@ -594,14 +594,10 @@ class OperatorLookupDescriptor final {
   using Storage = llvm::PointerUnion<FileUnit *, ModuleDecl *>;
   Storage fileOrModule;
   Identifier name;
-  bool isCascading;
-  SourceLoc diagLoc;
 
 private:
-  OperatorLookupDescriptor(Storage fileOrModule, Identifier name,
-                           bool isCascading, SourceLoc diagLoc)
-      : fileOrModule(fileOrModule), name(name), isCascading(isCascading),
-        diagLoc(diagLoc) {}
+  OperatorLookupDescriptor(Storage fileOrModule, Identifier name)
+      : fileOrModule(fileOrModule), name(name) {}
 
 public:
   /// Retrieves the files to perform lookup in.
@@ -613,66 +609,43 @@ class OperatorLookupDescriptor final {
     return fileOrModule.dyn_cast<ModuleDecl *>();
   }
 
+  /// Retrieve the file or module for the lookup, as a DeclContext.
+  DeclContext *getDC() const {
+    if (auto *module = getModule())
+      return module;
+    return fileOrModule.get<FileUnit *>();
+  }
+
   friend llvm::hash_code hash_value(const OperatorLookupDescriptor &desc) {
-    return llvm::hash_combine(desc.fileOrModule, desc.name, desc.isCascading);
+    return llvm::hash_combine(desc.fileOrModule, desc.name);
   }
 
   friend bool operator==(const OperatorLookupDescriptor &lhs,
                          const OperatorLookupDescriptor &rhs) {
-    return lhs.fileOrModule == rhs.fileOrModule && lhs.name == rhs.name &&
-           lhs.isCascading == rhs.isCascading;
+    return lhs.fileOrModule == rhs.fileOrModule && lhs.name == rhs.name;
   }
 
   friend bool operator!=(const OperatorLookupDescriptor &lhs,
                          const OperatorLookupDescriptor &rhs) {
     return !(lhs == rhs);
   }
 
-  static OperatorLookupDescriptor forFile(FileUnit *file, Identifier name,
-                                          bool isCascading, SourceLoc diagLoc) {
-    return OperatorLookupDescriptor(file, name, isCascading, diagLoc);
+  static OperatorLookupDescriptor forFile(FileUnit *file, Identifier name) {
+    return OperatorLookupDescriptor(file, name);
   }
 
-  static OperatorLookupDescriptor forModule(ModuleDecl *mod, Identifier name,
-                                            bool isCascading,
-                                            SourceLoc diagLoc) {
-    return OperatorLookupDescriptor(mod, name, isCascading, diagLoc);
+  static OperatorLookupDescriptor forModule(ModuleDecl *mod, Identifier name) {
+    return OperatorLookupDescriptor(mod, name);
   }
+
+  static OperatorLookupDescriptor forDC(const DeclContext *DC, Identifier name);
 };
 
 void simple_display(llvm::raw_ostream &out,
                     const OperatorLookupDescriptor &desc);
 
 SourceLoc extractNearestSourceLoc(const OperatorLookupDescriptor &desc);
 
-template <typename OperatorType>
-class LookupOperatorRequest
-    : public SimpleRequest<LookupOperatorRequest<OperatorType>,
-                           OperatorType *(OperatorLookupDescriptor),
-                           RequestFlags::Cached> {
-  using SimpleRequest<LookupOperatorRequest<OperatorType>,
-                      OperatorType *(OperatorLookupDescriptor),
-                      RequestFlags::Cached>::SimpleRequest;
-
-private:
-  friend SimpleRequest<LookupOperatorRequest<OperatorType>,
-                       OperatorType *(OperatorLookupDescriptor),
-                       RequestFlags::Cached>;
-
-  // Evaluation.
-  OperatorType *evaluate(Evaluator &evaluator,
-                         OperatorLookupDescriptor desc) const;
-
-public:
-  // Cached.
-  bool isCached() const { return true; }
-};
-
-using LookupPrefixOperatorRequest = LookupOperatorRequest<PrefixOperatorDecl>;
-using LookupInfixOperatorRequest = LookupOperatorRequest<InfixOperatorDecl>;
-using LookupPostfixOperatorRequest = LookupOperatorRequest<PostfixOperatorDecl>;
-using LookupPrecedenceGroupRequest = LookupOperatorRequest<PrecedenceGroupDecl>;
-
 /// Looks up an operator in a given file or module without looking through
 /// imports.
 class DirectOperatorLookupRequest
@@ -770,6 +743,84 @@ class LookupConformanceInModuleRequest
                            ProtocolConformanceRef result) const;
 };
 
+/// Look up an 'infix operator' decl by name.
+class LookupInfixOperatorRequest
+    : public SimpleRequest<LookupInfixOperatorRequest,
+                           TinyPtrVector<InfixOperatorDecl *>(
+                               OperatorLookupDescriptor),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  TinyPtrVector<InfixOperatorDecl *>
+  evaluate(Evaluator &evaluator, OperatorLookupDescriptor desc) const;
+
+public:
+  // Cached.
+  bool isCached() const { return true; }
+};
+
+/// Look up an 'prefix operator' decl by name.
+class LookupPrefixOperatorRequest
+    : public SimpleRequest<LookupPrefixOperatorRequest,
+                           PrefixOperatorDecl *(OperatorLookupDescriptor),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  PrefixOperatorDecl *evaluate(Evaluator &evaluator,
+                               OperatorLookupDescriptor desc) const;
+
+public:
+  // Cached.
+  bool isCached() const { return true; }
+};
+
+/// Look up an 'postfix operator' decl by name.
+class LookupPostfixOperatorRequest
+    : public SimpleRequest<LookupPostfixOperatorRequest,
+                           PostfixOperatorDecl *(OperatorLookupDescriptor),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  PostfixOperatorDecl *evaluate(Evaluator &evaluator,
+                                OperatorLookupDescriptor desc) const;
+
+public:
+  // Cached.
+  bool isCached() const { return true; }
+};
+
+/// Look up a precedencegroup decl by name.
+class LookupPrecedenceGroupRequest
+    : public SimpleRequest<LookupPrecedenceGroupRequest,
+                           TinyPtrVector<PrecedenceGroupDecl *>(
+                               OperatorLookupDescriptor),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  TinyPtrVector<PrecedenceGroupDecl *>
+  evaluate(Evaluator &evaluator, OperatorLookupDescriptor descriptor) const;
+
+public:
+  // Cached.
+  bool isCached() const { return true; }
+};
+
 #define SWIFT_TYPEID_ZONE NameLookup
 #define SWIFT_TYPEID_HEADER "swift/AST/NameLookupTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/NameLookupTypeIDZone.def

@@ -89,11 +89,11 @@ SWIFT_REQUEST(NameLookup, LookupPrefixOperatorRequest,
               PrefixOperatorDecl *(OperatorLookupDescriptor),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, LookupInfixOperatorRequest,
-              InfixOperatorDecl *(OperatorLookupDescriptor),
+              TinyPtrVector<InfixOperatorDecl *>(OperatorLookupDescriptor),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, LookupPostfixOperatorRequest,
               PostfixOperatorDecl *(OperatorLookupDescriptor),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, LookupPrecedenceGroupRequest,
-              PrecedenceGroupDecl *(OperatorLookupDescriptor),
+              TinyPtrVector<PrecedenceGroupDecl *>(OperatorLookupDescriptor),
               Cached, NoLocationInfo)