Add direct operator lookup requests

Introduce DirectOperatorLookupRequest &
DirectPrecedenceGroupLookupRequest that lookup
operator and precedence groups within a given
file or module without looking through imports.

These will eventually be used as the basis for the
new operator lookup implementation, but for now
just use them when querying lookup results from
serialized module files.

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 @@ public:
                                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,23 +364,6 @@ public:
     return StringRef();
   }
 
-  /// Look up an operator declaration.
-  ///
-  /// \param name The operator name ("+", ">>", etc.)
-  ///
-  /// \param fixity One of Prefix, Infix, or Postfix.
-  virtual OperatorDecl *lookupOperator(Identifier name,
-                                       OperatorFixity 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

@@ -597,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 @@ public:
          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 @@ public:
   lookupNestedType(Identifier name,
                    const NominalTypeDecl *parent) const override;
 
-  virtual OperatorDecl *lookupOperator(Identifier name,
+protected:
-                                       OperatorFixity fixity) const override;
+  virtual void
+  lookupOperatorDirect(Identifier name, OperatorFixity fixity,
+                       TinyPtrVector<OperatorDecl *> &results) const override;
 
-  virtual PrecedenceGroupDecl *
+  virtual void lookupPrecedenceGroupDirect(
-  lookupPrecedenceGroup(Identifier name) const override;
+      Identifier name,
+      TinyPtrVector<PrecedenceGroupDecl *> &results) const override;
 
+public:
   virtual void lookupVisibleDecls(ModuleDecl::AccessPathTy accessPath,
                                   VisibleDeclConsumer &consumer,
                                   NLKind lookupKind) const override;

lib/AST/Module.cpp

@@ -945,39 +945,50 @@ namespace {
   template <>
   struct OperatorLookup<PrefixOperatorDecl> {
     constexpr static auto map_ptr = &SourceFile::PrefixOperators;
-    template <typename T>
+    static PrefixOperatorDecl *lookup(Evaluator &eval,
-    static PrefixOperatorDecl *lookup(T &container, Identifier name) {
+                                      const OperatorLookupDescriptor &desc) {
-      return cast_or_null<PrefixOperatorDecl>(
+      // We can return the first prefix operator. All prefix operators of the
-               container.lookupOperator(name, OperatorFixity::Prefix));
+      // same name are equivalent.
+      DirectOperatorLookupRequest req{desc, OperatorFixity::Prefix};
+      auto results = evaluateOrDefault(eval, req, {});
+      return results.empty() ? nullptr : cast<PrefixOperatorDecl>(results[0]);
     }
   };
 
   template <>
   struct OperatorLookup<InfixOperatorDecl> {
     constexpr static auto map_ptr = &SourceFile::InfixOperators;
-    template <typename T>
+    static InfixOperatorDecl *lookup(Evaluator &eval,
-    static InfixOperatorDecl *lookup(T &container, Identifier name) {
+                                     const OperatorLookupDescriptor &desc) {
-      return cast_or_null<InfixOperatorDecl>(
+      // Return the first result if it exists.
-               container.lookupOperator(name, OperatorFixity::Infix));
+      DirectOperatorLookupRequest req{desc, OperatorFixity::Infix};
+      auto results = evaluateOrDefault(eval, req, {});
+      return results.empty() ? nullptr : cast<InfixOperatorDecl>(results[0]);
     }
   };
 
   template <>
   struct OperatorLookup<PostfixOperatorDecl> {
     constexpr static auto map_ptr = &SourceFile::PostfixOperators;
-    template <typename T>
+    static PostfixOperatorDecl *lookup(Evaluator &eval,
-    static PostfixOperatorDecl *lookup(T &container, Identifier name) {
+                                       const OperatorLookupDescriptor &desc) {
-      return cast_or_null<PostfixOperatorDecl>(
+      // We can return the first postfix operator. All postfix operators of the
-               container.lookupOperator(name, OperatorFixity::Postfix));
+      // same name are equivalent.
+      DirectOperatorLookupRequest req{desc, OperatorFixity::Postfix};
+      auto results = evaluateOrDefault(eval, req, {});
+      return results.empty() ? nullptr : cast<PostfixOperatorDecl>(results[0]);
     }
   };
 
   template <>
   struct OperatorLookup<PrecedenceGroupDecl> {
     constexpr static auto map_ptr = &SourceFile::PrecedenceGroups;
-    template <typename T>
+    static PrecedenceGroupDecl *lookup(Evaluator &eval,
-    static PrecedenceGroupDecl *lookup(T &container, Identifier name) {
+                                       const OperatorLookupDescriptor &desc) {
-      return container.lookupPrecedenceGroup(name);
+      // Return the first result if it exists.
+      auto results =
+          evaluateOrDefault(eval, DirectPrecedenceGroupLookupRequest{desc}, {});
+      return results.empty() ? nullptr : results[0];
     }
   };
 } // end anonymous namespace
@@ -1014,7 +1025,8 @@ void SourceFile::setSyntaxRoot(syntax::SourceFileSyntax &&Root) {
 
 template<typename OP_DECL>
 static Optional<OP_DECL *>
-lookupOperatorDeclForName(ModuleDecl *M, SourceLoc Loc, Identifier Name);
+lookupOperatorDeclForName(ModuleDecl *M, SourceLoc Loc, Identifier Name,
+                          bool isCascading);
 
 template<typename OP_DECL>
 using ImportedOperatorsMap = llvm::SmallDenseMap<OP_DECL*, bool, 16>;
@@ -1063,7 +1075,8 @@ checkOperatorConflicts(const SourceFile &SF, SourceLoc loc,
 template <typename OP_DECL>
 static Optional<OP_DECL *>
 lookupOperatorDeclForName(const FileUnit &File, SourceLoc Loc,
-                          Identifier Name, bool includePrivate) {
+                          Identifier Name, bool includePrivate,
+                          bool isCascading) {
   switch (File.getKind()) {
   case FileUnitKind::Builtin:
     // The Builtin module declares no operators.
@@ -1072,8 +1085,13 @@ lookupOperatorDeclForName(const FileUnit &File, SourceLoc Loc,
     break;
   case FileUnitKind::SerializedAST:
   case FileUnitKind::ClangModule:
-  case FileUnitKind::DWARFModule:
+  case FileUnitKind::DWARFModule: {
-    return OperatorLookup<OP_DECL>::lookup(cast<LoadedFile>(File), Name);
+    auto &eval = File.getASTContext().evaluator;
+    auto desc = OperatorLookupDescriptor::forFile(const_cast<FileUnit *>(&File),
+                                                  Name, isCascading,
+                                                  /*diagLoc*/ SourceLoc());
+    return OperatorLookup<OP_DECL>::lookup(eval, desc);
+  }
   }
 
   auto &SF = cast<SourceFile>(File);
@@ -1102,7 +1120,8 @@ lookupOperatorDeclForName(const FileUnit &File, SourceLoc Loc,
       continue;
 
     Optional<OP_DECL *> maybeOp =
-        lookupOperatorDeclForName<OP_DECL>(imported.module.second, Loc, Name);
+        lookupOperatorDeclForName<OP_DECL>(imported.module.second, Loc, Name,
+                                           isCascading);
     if (!maybeOp)
       return None;
     
@@ -1135,10 +1154,12 @@ lookupOperatorDeclForName(const FileUnit &File, SourceLoc Loc,
 
 template<typename OP_DECL>
 static Optional<OP_DECL *>
-lookupOperatorDeclForName(ModuleDecl *M, SourceLoc Loc, Identifier Name) {
+lookupOperatorDeclForName(ModuleDecl *M, SourceLoc Loc, Identifier Name,
+                          bool isCascading) {
   OP_DECL *result = nullptr;
   for (const FileUnit *File : M->getFiles()) {
-    auto next = lookupOperatorDeclForName<OP_DECL>(*File, Loc, Name, false);
+    auto next = lookupOperatorDeclForName<OP_DECL>(*File, Loc, Name, false,
+                                                   isCascading);
     if (!next.hasValue())
       return next;
 
@@ -1155,9 +1176,10 @@ template <typename OperatorType>
 llvm::Expected<OperatorType *> LookupOperatorRequest<OperatorType>::evaluate(
     Evaluator &evaluator, OperatorLookupDescriptor desc) const {
   auto *file = desc.fileOrModule.get<FileUnit *>();
-  auto result = lookupOperatorDeclForName<OperatorType>(*file, desc.diagLoc,
+  auto result =
-                                                        desc.name,
+      lookupOperatorDeclForName<OperatorType>(*file, desc.diagLoc, desc.name,
-                                                        /*includePrivate*/ true);
+                                              /*includePrivate*/ true,
+                                              desc.isCascading);
   if (!result.hasValue())
     return nullptr;
 
@@ -1167,16 +1189,17 @@ llvm::Expected<OperatorType *> LookupOperatorRequest<OperatorType>::evaluate(
   }
   if (!result.getValue()) {
     result = lookupOperatorDeclForName<OperatorType>(file->getParentModule(),
-                                                     desc.diagLoc, desc.name);
+                                                     desc.diagLoc, desc.name,
+                                                     desc.isCascading);
   }
   return result.hasValue() ? result.getValue() : nullptr;
 }
 
-
 #define LOOKUP_OPERATOR(Kind)                                                  \
   Kind##Decl *ModuleDecl::lookup##Kind(Identifier name, SourceLoc loc) {       \
     auto result =                                                              \
-        lookupOperatorDeclForName<Kind##Decl>(this, loc, name);                \
+        lookupOperatorDeclForName<Kind##Decl>(this, loc, name,                 \
+                                              /*isCascading*/ false);          \
     return result ? *result : nullptr;                                         \
   }                                                                            \
   template llvm::Expected<Kind##Decl *>                                        \
@@ -1189,6 +1212,75 @@ LOOKUP_OPERATOR(PostfixOperator)
 LOOKUP_OPERATOR(PrecedenceGroup)
 #undef LOOKUP_OPERATOR
 
+llvm::Expected<TinyPtrVector<OperatorDecl *>>
+DirectOperatorLookupRequest::evaluate(Evaluator &evaluator,
+                                      OperatorLookupDescriptor descriptor,
+                                      OperatorFixity fixity) const {
+  // Query each file.
+  // TODO: Module-level caching.
+  TinyPtrVector<OperatorDecl *> results;
+  for (auto *file : descriptor.getFiles())
+    file->lookupOperatorDirect(descriptor.name, fixity, results);
+
+  return std::move(results);
+}
+
+void SourceFile::lookupOperatorDirect(
+    Identifier name, OperatorFixity fixity,
+    TinyPtrVector<OperatorDecl *> &results) const {
+  OperatorDecl *op = nullptr;
+  switch (fixity) {
+  case OperatorFixity::Infix: {
+    auto result = InfixOperators.find(name);
+    if (result != InfixOperators.end())
+      op = result->second.getPointer();
+    break;
+  }
+  case OperatorFixity::Postfix: {
+    auto result = PostfixOperators.find(name);
+    if (result != PostfixOperators.end())
+      op = result->second.getPointer();
+    break;
+  }
+  case OperatorFixity::Prefix: {
+    auto result = PrefixOperators.find(name);
+    if (result != PrefixOperators.end())
+      op = result->second.getPointer();
+    break;
+  }
+  }
+
+  // We currently can use the operator maps to cache lookup results from other
+  // modules. Make sure we only return results from the source file.
+  if (op && op->getDeclContext()->getParentSourceFile() == this)
+    results.push_back(op);
+}
+
+llvm::Expected<TinyPtrVector<PrecedenceGroupDecl *>>
+DirectPrecedenceGroupLookupRequest::evaluate(
+    Evaluator &evaluator, OperatorLookupDescriptor descriptor) const {
+  // Query each file.
+  // TODO: Module-level caching.
+  TinyPtrVector<PrecedenceGroupDecl *> results;
+  for (auto *file : descriptor.getFiles())
+    file->lookupPrecedenceGroupDirect(descriptor.name, results);
+
+  return std::move(results);
+}
+
+void SourceFile::lookupPrecedenceGroupDirect(
+    Identifier name, TinyPtrVector<PrecedenceGroupDecl *> &results) const {
+  auto result = PrecedenceGroups.find(name);
+  if (result == PrecedenceGroups.end())
+    return;
+
+  // We currently can use the operator maps to cache lookup results from other
+  // modules. Make sure we only return results from the source file.
+  auto *group = result->second.getPointer();
+  if (group->getDeclContext()->getParentSourceFile() == this)
+    results.push_back(group);
+}
+
 void ModuleDecl::getImportedModules(SmallVectorImpl<ImportedModule> &modules,
                                     ModuleDecl::ImportFilter filter) const {
   FORWARD(getImportedModules, (modules, filter));

lib/Serialization/SerializedModuleLoader.cpp

@@ -1084,14 +1084,17 @@ SerializedASTFile::lookupNestedType(Identifier name,
   return File.lookupNestedType(name, parent);
 }
 
-OperatorDecl *SerializedASTFile::lookupOperator(Identifier name,
+void SerializedASTFile::lookupOperatorDirect(
-                                                OperatorFixity fixity) const {
+    Identifier name, OperatorFixity fixity,
-  return File.lookupOperator(name, fixity);
+    TinyPtrVector<OperatorDecl *> &results) const {
+  if (auto *op = File.lookupOperator(name, fixity))
+    results.push_back(op);
 }
 
-PrecedenceGroupDecl *
+void SerializedASTFile::lookupPrecedenceGroupDirect(
-SerializedASTFile::lookupPrecedenceGroup(Identifier name) const {
+    Identifier name, TinyPtrVector<PrecedenceGroupDecl *> &results) const {
-  return File.lookupPrecedenceGroup(name);
+  if (auto *group = File.lookupPrecedenceGroup(name))
+    results.push_back(group);
 }
 
 void SerializedASTFile::lookupVisibleDecls(ModuleDecl::AccessPathTy accessPath,