Add direct operator lookup requests

hamishknight · hamishknight · commit cc9c85185645 · 2020-03-23T09:17:58.000-07:00
Introduce DirectOperatorLookupRequest &amp;
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.
diff --git a/include/swift/AST/FileUnit.h b/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,23 +364,6 @@ class LoadedFile : public FileUnit {
     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; }
 
diff --git a/include/swift/AST/NameLookupRequests.h b/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"
diff --git a/include/swift/AST/NameLookupTypeIDZone.def b/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,
diff --git a/include/swift/AST/SourceFile.h b/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
diff --git a/include/swift/Serialization/SerializedModuleLoader.h b/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,
-                                       OperatorFixity 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;
diff --git a/lib/AST/Module.cpp b/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(T &container, Identifier name) {
-      return cast_or_null<PrefixOperatorDecl>(
-               container.lookupOperator(name, OperatorFixity::Prefix));
+    static PrefixOperatorDecl *lookup(Evaluator &eval,
+                                      const OperatorLookupDescriptor &desc) {
+      // We can return the first prefix operator. All prefix operators of the
+      // 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(T &container, Identifier name) {
-      return cast_or_null<InfixOperatorDecl>(
-               container.lookupOperator(name, OperatorFixity::Infix));
+    static InfixOperatorDecl *lookup(Evaluator &eval,
+                                     const OperatorLookupDescriptor &desc) {
+      // Return the first result if it exists.
+      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(T &container, Identifier name) {
-      return cast_or_null<PostfixOperatorDecl>(
-               container.lookupOperator(name, OperatorFixity::Postfix));
+    static PostfixOperatorDecl *lookup(Evaluator &eval,
+                                       const OperatorLookupDescriptor &desc) {
+      // We can return the first postfix operator. All postfix operators of the
+      // 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(T &container, Identifier name) {
-      return container.lookupPrecedenceGroup(name);
+    static PrecedenceGroupDecl *lookup(Evaluator &eval,
+                                       const OperatorLookupDescriptor &desc) {
+      // 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:
-    return OperatorLookup<OP_DECL>::lookup(cast<LoadedFile>(File), Name);
+  case FileUnitKind::DWARFModule: {
+    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,
-                                                        desc.name,
-                                                        /*includePrivate*/ true);
+  auto result =
+      lookupOperatorDeclForName<OperatorType>(*file, desc.diagLoc, desc.name,
+                                              /*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));
diff --git a/lib/Serialization/SerializedModuleLoader.cpp b/lib/Serialization/SerializedModuleLoader.cpp
