Add direct operator lookup requests (#30584)

Add direct operator lookup requests

include/swift/AST/Decl.h

@@ -7063,6 +7063,28 @@ class PrecedenceGroupDecl : public Decl {
   }
 };
 
+/// The fixity of an OperatorDecl.
+enum class OperatorFixity : uint8_t {
+  Infix,
+  Prefix,
+  Postfix
+};
+
+inline void simple_display(llvm::raw_ostream &out, OperatorFixity fixity) {
+  switch (fixity) {
+  case OperatorFixity::Infix:
+    out << "infix";
+    return;
+  case OperatorFixity::Prefix:
+    out << "prefix";
+    return;
+  case OperatorFixity::Postfix:
+    out << "postfix";
+    return;
+  }
+  llvm_unreachable("Unhandled case in switch");
+}
+
 /// Abstract base class of operator declarations.
 class OperatorDecl : public Decl {
   SourceLoc OperatorLoc, NameLoc;
@@ -7088,6 +7110,21 @@ class OperatorDecl : public Decl {
       : Decl(kind, DC), OperatorLoc(OperatorLoc), NameLoc(NameLoc), name(Name),
         DesignatedNominalTypes(DesignatedNominalTypes) {}
 
+  /// Retrieve the operator's fixity, corresponding to the concrete subclass
+  /// of the OperatorDecl.
+  OperatorFixity getFixity() const {
+    switch (getKind()) {
+#define DECL(Id, Name) case DeclKind::Id: llvm_unreachable("Not an operator!");
+#define OPERATOR_DECL(Id, Name)
+#include "swift/AST/DeclNodes.def"
+    case DeclKind::InfixOperator:
+      return OperatorFixity::Infix;
+    case DeclKind::PrefixOperator:
+      return OperatorFixity::Prefix;
+    case DeclKind::PostfixOperator:
+      return OperatorFixity::Postfix;
+    }
+  }
 
   SourceLoc getOperatorLoc() const { return OperatorLoc; }
   SourceLoc getNameLoc() const { return NameLoc; }

include/swift/AST/FileUnit.h

@@ -30,6 +30,9 @@ class FileUnit : public DeclContext {
 #pragma clang diagnostic pop
   virtual void anchor();
 
+  friend class DirectOperatorLookupRequest;
+  friend class DirectPrecedenceGroupLookupRequest;
+
   // FIXME: Stick this in a PointerIntPair.
   const FileUnitKind Kind;
 
@@ -107,6 +110,25 @@ class FileUnit : public DeclContext {
                                const ModuleDecl *importedModule,
                                SmallVectorImpl<Identifier> &spiGroups) const {};
 
+protected:
+  /// Look up an operator declaration. Do not call directly, use
+  /// \c DirectOperatorLookupRequest instead.
+  ///
+  /// \param name The operator name ("+", ">>", etc.)
+  ///
+  /// \param fixity One of Prefix, Infix, or Postfix.
+  virtual void
+  lookupOperatorDirect(Identifier name, OperatorFixity fixity,
+                       TinyPtrVector<OperatorDecl *> &results) const {}
+
+  /// Look up a precedence group. Do not call directly, use
+  /// \c DirectPrecedenceGroupLookupRequest instead.
+  ///
+  /// \param name The precedence group name.
+  virtual void lookupPrecedenceGroupDirect(
+      Identifier name, TinyPtrVector<PrecedenceGroupDecl *> &results) const {}
+
+public:
   /// Returns the comment attached to the given declaration.
   ///
   /// This function is an implementation detail for comment serialization.
@@ -342,22 +364,6 @@ class LoadedFile : public FileUnit {
     return StringRef();
   }
 
-  /// Look up an operator declaration.
-  ///
-  /// \param name The operator name ("+", ">>", etc.)
-  ///
-  /// \param fixity One of PrefixOperator, InfixOperator, or PostfixOperator.
-  virtual OperatorDecl *lookupOperator(Identifier name, DeclKind fixity) const {
-    return nullptr;
-  }
-
-  /// Look up a precedence group.
-  ///
-  /// \param name The precedence group name.
-  virtual PrecedenceGroupDecl *lookupPrecedenceGroup(Identifier name) const {
-    return nullptr;
-  }
-
   /// Returns the Swift module that overlays a Clang module.
   virtual ModuleDecl *getOverlayModule() const { return nullptr; }
 

include/swift/AST/NameLookupRequests.h

@@ -18,6 +18,7 @@
 
 #include "swift/AST/SimpleRequest.h"
 #include "swift/AST/ASTTypeIDs.h"
+#include "swift/AST/FileUnit.h"
 #include "swift/AST/Identifier.h"
 #include "swift/Basic/Statistic.h"
 #include "llvm/ADT/Hashing.h"
@@ -518,29 +519,53 @@ class DirectLookupRequest
 
 class OperatorLookupDescriptor final {
 public:
-  SourceFile *SF;
+  using Storage = llvm::PointerUnion<FileUnit *, ModuleDecl *>;
+  Storage fileOrModule;
   Identifier name;
   bool isCascading;
   SourceLoc diagLoc;
 
-  OperatorLookupDescriptor(SourceFile *SF, Identifier name, bool isCascading,
-                           SourceLoc diagLoc)
-      : SF(SF), name(name), isCascading(isCascading), diagLoc(diagLoc) {}
+private:
+  OperatorLookupDescriptor(Storage fileOrModule, Identifier name,
+                           bool isCascading, SourceLoc diagLoc)
+      : fileOrModule(fileOrModule), name(name), isCascading(isCascading),
+        diagLoc(diagLoc) {}
+
+public:
+  /// Retrieves the files to perform lookup in.
+  ArrayRef<FileUnit *> getFiles() const;
+
+  /// If this is for a module lookup, returns the module. Otherwise returns
+  /// \c nullptr.
+  ModuleDecl *getModule() const {
+    return fileOrModule.dyn_cast<ModuleDecl *>();
+  }
 
   friend llvm::hash_code hash_value(const OperatorLookupDescriptor &desc) {
-    return llvm::hash_combine(desc.SF, desc.name, desc.isCascading);
+    return llvm::hash_combine(desc.fileOrModule, desc.name, desc.isCascading);
   }
 
   friend bool operator==(const OperatorLookupDescriptor &lhs,
                          const OperatorLookupDescriptor &rhs) {
-    return lhs.SF == rhs.SF && lhs.name == rhs.name &&
+    return lhs.fileOrModule == rhs.fileOrModule && lhs.name == rhs.name &&
            lhs.isCascading == rhs.isCascading;
   }
 
   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 forModule(ModuleDecl *mod, Identifier name,
+                                            bool isCascading,
+                                            SourceLoc diagLoc) {
+    return OperatorLookupDescriptor(mod, name, isCascading, diagLoc);
+  }
 };
 
 void simple_display(llvm::raw_ostream &out,
@@ -572,6 +597,41 @@ 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
+    : public SimpleRequest<DirectOperatorLookupRequest,
+                           TinyPtrVector<OperatorDecl *>(
+                               OperatorLookupDescriptor, OperatorFixity),
+                           CacheKind::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  llvm::Expected<TinyPtrVector<OperatorDecl *>>
+  evaluate(Evaluator &evaluator, OperatorLookupDescriptor descriptor,
+           OperatorFixity fixity) const;
+};
+
+/// Looks up an precedencegroup in a given file or module without looking
+/// through imports.
+class DirectPrecedenceGroupLookupRequest
+    : public SimpleRequest<DirectPrecedenceGroupLookupRequest,
+                           TinyPtrVector<PrecedenceGroupDecl *>(
+                               OperatorLookupDescriptor),
+                           CacheKind::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  llvm::Expected<TinyPtrVector<PrecedenceGroupDecl *>>
+  evaluate(Evaluator &evaluator, OperatorLookupDescriptor descriptor) const;
+};
+
 #define SWIFT_TYPEID_ZONE NameLookup
 #define SWIFT_TYPEID_HEADER "swift/AST/NameLookupTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/NameLookupTypeIDZone.def

@@ -24,6 +24,13 @@ SWIFT_REQUEST(NameLookup, CustomAttrNominalRequest,
 SWIFT_REQUEST(NameLookup, DirectLookupRequest,
               TinyPtrVector<ValueDecl *>(DirectLookupDescriptor), Uncached,
               NoLocationInfo)
+SWIFT_REQUEST(NameLookup, DirectOperatorLookupRequest,
+              TinyPtrVector<OperatorDecl *>(OperatorLookupDescriptor,
+                                            OperatorFixity),
+              Uncached, NoLocationInfo)
+SWIFT_REQUEST(NameLookup, DirectPrecedenceGroupLookupRequest,
+              TinyPtrVector<PrecedenceGroupDecl *>(OperatorLookupDescriptor),
+              Uncached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, ExpandASTScopeRequest,
               ast_scope::ASTScopeImpl* (ast_scope::ASTScopeImpl*, ast_scope::ScopeCreator*),
               SeparatelyCached,

include/swift/AST/SourceFile.h

@@ -436,6 +436,16 @@ class SourceFile final : public FileUnit {
          ObjCSelector selector,
          SmallVectorImpl<AbstractFunctionDecl *> &results) const override;
 
+protected:
+  virtual void
+  lookupOperatorDirect(Identifier name, OperatorFixity fixity,
+                       TinyPtrVector<OperatorDecl *> &results) const override;
+
+  virtual void lookupPrecedenceGroupDirect(
+      Identifier name,
+      TinyPtrVector<PrecedenceGroupDecl *> &results) const override;
+
+public:
   virtual void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const override;
 
   virtual void

include/swift/Serialization/SerializedModuleLoader.h

@@ -328,12 +328,16 @@ class SerializedASTFile final : public LoadedFile {
   lookupNestedType(Identifier name,
                    const NominalTypeDecl *parent) const override;
 
-  virtual OperatorDecl *lookupOperator(Identifier name,
-                                       DeclKind fixity) const override;
+protected:
+  virtual void
+  lookupOperatorDirect(Identifier name, OperatorFixity fixity,
+                       TinyPtrVector<OperatorDecl *> &results) const override;
 
-  virtual PrecedenceGroupDecl *
-  lookupPrecedenceGroup(Identifier name) const override;
+  virtual void lookupPrecedenceGroupDirect(
+      Identifier name,
+      TinyPtrVector<PrecedenceGroupDecl *> &results) const override;
 
+public:
   virtual void lookupVisibleDecls(ModuleDecl::AccessPathTy accessPath,
                                   VisibleDeclConsumer &consumer,
                                   NLKind lookupKind) const override;