[NameLookup] Add requests for module and AnyObject lookup (#28395)

[NameLookup] Add requests for module and AnyObject lookup

Author: hamishknight

include/swift/AST/DeclContext.h

@@ -530,10 +530,6 @@ class alignas(1 << DeclContextAlignInBits) DeclContext {
   bool lookupQualified(ModuleDecl *module, DeclName member, NLOptions options,
                        SmallVectorImpl<ValueDecl *> &decls) const;
 
-  /// Perform \c AnyObject lookup for the given member.
-  bool lookupAnyObject(DeclName member, NLOptions options,
-                       SmallVectorImpl<ValueDecl *> &decls) const;
-
   /// Look up all Objective-C methods with the given selector visible
   /// in the enclosing module.
   void lookupAllObjCMethods(

include/swift/AST/LookupKinds.h

@@ -26,6 +26,8 @@ enum class NLKind {
   QualifiedLookup
 };
 
+void simple_display(llvm::raw_ostream &out, NLKind kind);
+
 /// Constants used to customize name lookup.
 enum NLOptions : unsigned {
   /// Consider declarations within protocols to which the context type conforms.
@@ -103,6 +105,8 @@ static inline NLOptions operator~(NLOptions value) {
   return NLOptions(~(unsigned)value);
 }
 
+void simple_display(llvm::raw_ostream &out, NLOptions options);
+
 } // end namespace swift
 
 #endif

include/swift/AST/ModuleNameLookup.h

@@ -39,6 +39,8 @@ enum class ResolutionKind {
   TypesOnly
 };
 
+void simple_display(llvm::raw_ostream &out, ResolutionKind kind);
+
 /// Performs a lookup into the given module and it's imports.
 ///
 /// If 'moduleOrFile' is a ModuleDecl, we search the module and it's

include/swift/AST/NameLookupRequests.h

@@ -31,6 +31,7 @@ class DestructorDecl;
 class GenericContext;
 class GenericParamList;
 class LookupResult;
+enum class NLKind;
 class SourceLoc;
 class TypeAliasDecl;
 class TypeDecl;
@@ -39,6 +40,9 @@ namespace ast_scope {
 class ASTScopeImpl;
 class ScopeCreator;
 } // namespace ast_scope
+namespace namelookup {
+enum class ResolutionKind;
+} // namespace namelookup
 
 /// Display a nominal type or extension thereof.
 void simple_display(
@@ -360,6 +364,46 @@ class UnqualifiedLookupRequest
                                         UnqualifiedLookupDescriptor desc) const;
 };
 
+using QualifiedLookupResult = SmallVector<ValueDecl *, 4>;
+
+/// Performs a lookup into a given module and its imports.
+class LookupInModuleRequest
+    : public SimpleRequest<LookupInModuleRequest,
+                           QualifiedLookupResult(
+                               const DeclContext *, DeclName, NLKind,
+                               namelookup::ResolutionKind, const DeclContext *),
+                           CacheKind::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  llvm::Expected<QualifiedLookupResult>
+  evaluate(Evaluator &evaluator, const DeclContext *moduleOrFile, DeclName name,
+           NLKind lookupKind, namelookup::ResolutionKind resolutionKind,
+           const DeclContext *moduleScopeContext) const;
+};
+
+/// Perform \c AnyObject lookup for a given member.
+class AnyObjectLookupRequest
+    : public SimpleRequest<AnyObjectLookupRequest,
+                           QualifiedLookupResult(const DeclContext *, DeclName,
+                                                 NLOptions),
+                           CacheKind::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  llvm::Expected<QualifiedLookupResult> evaluate(Evaluator &evaluator,
+                                                 const DeclContext *dc,
+                                                 DeclName name,
+                                                 NLOptions options) const;
+};
+
 #define SWIFT_TYPEID_ZONE NameLookup
 #define SWIFT_TYPEID_HEADER "swift/AST/NameLookupTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/NameLookupTypeIDZone.def

@@ -48,3 +48,11 @@ SWIFT_REQUEST(NameLookup, UnderlyingTypeDeclsReferencedRequest,
 SWIFT_REQUEST(NameLookup, UnqualifiedLookupRequest,
               LookupResult(UnqualifiedLookupDescriptor), Uncached,
               NoLocationInfo)
+SWIFT_REQUEST(NameLookup, LookupInModuleRequest,
+              QualifiedLookupResult(const DeclContext *, DeclName, NLKind,
+                                    namelookup::ResolutionKind,
+                                    const DeclContext *),
+              Uncached, NoLocationInfo)
+SWIFT_REQUEST(NameLookup, AnyObjectLookupRequest,
+              QualifiedLookupResult(const DeclContext *, DeclName, NLOptions),
+              Uncached, NoLocationInfo)

include/swift/Basic/Statistics.def

@@ -144,12 +144,6 @@ FRONTEND_STATISTIC(AST, NumLookupQualifiedInNominal)
 /// Number of qualified lookups into a module.
 FRONTEND_STATISTIC(AST, NumLookupQualifiedInModule)
 
-/// Number of qualified lookups into AnyObject.
-FRONTEND_STATISTIC(AST, NumLookupQualifiedInAnyObject)
-
-/// Number of lookups into a module and its imports.
-FRONTEND_STATISTIC(AST, NumLookupInModule)
-
 /// Number of local lookups into a module.
 FRONTEND_STATISTIC(AST, NumModuleLookupValue)
 

lib/AST/ModuleNameLookup.cpp

@@ -15,6 +15,7 @@
 #include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/ImportCache.h"
 #include "swift/AST/NameLookup.h"
+#include "swift/AST/NameLookupRequests.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace swift;
@@ -228,23 +229,32 @@ void ModuleNameLookup<LookupStrategy>::lookupInModule(
               decls.end());
 }
 
+llvm::Expected<QualifiedLookupResult> LookupInModuleRequest::evaluate(
+    Evaluator &evaluator, const DeclContext *moduleOrFile, DeclName name,
+    NLKind lookupKind, ResolutionKind resolutionKind,
+    const DeclContext *moduleScopeContext) const {
+  assert(moduleScopeContext->isModuleScopeContext());
+
+  auto &ctx = moduleOrFile->getASTContext();
+  FrontendStatsTracer tracer(ctx.Stats, "lookup-in-module");
+
+  QualifiedLookupResult decls;
+  LookupByName lookup(ctx, resolutionKind, name, lookupKind);
+  lookup.lookupInModule(decls, moduleOrFile, {}, moduleScopeContext);
+  return decls;
+}
+
 void namelookup::lookupInModule(const DeclContext *moduleOrFile,
                                 DeclName name,
                                 SmallVectorImpl<ValueDecl *> &decls,
                                 NLKind lookupKind,
                                 ResolutionKind resolutionKind,
                                 const DeclContext *moduleScopeContext) {
-  assert(moduleScopeContext->isModuleScopeContext());
-
   auto &ctx = moduleOrFile->getASTContext();
-  auto *stats = ctx.Stats;
-  if (stats)
-    stats->getFrontendCounters().NumLookupInModule++;
-
-  FrontendStatsTracer tracer(stats, "lookup-in-module");
-
-  LookupByName lookup(ctx, resolutionKind, name, lookupKind);
-  lookup.lookupInModule(decls, moduleOrFile, {}, moduleScopeContext);
+  LookupInModuleRequest req(moduleOrFile, name, lookupKind, resolutionKind,
+                            moduleScopeContext);
+  auto results = evaluateOrDefault(ctx.evaluator, req, {});
+  decls.append(results.begin(), results.end());
 }
 
 void namelookup::lookupVisibleDeclsInModule(
@@ -260,3 +270,14 @@ void namelookup::lookupVisibleDeclsInModule(
   lookup.lookupInModule(decls, moduleOrFile, accessPath, moduleScopeContext);
 }
 
+void namelookup::simple_display(llvm::raw_ostream &out, ResolutionKind kind) {
+  switch (kind) {
+  case ResolutionKind::Overloadable:
+    out << "Overloadable";
+    return;
+  case ResolutionKind::TypesOnly:
+    out << "TypesOnly";
+    return;
+  }
+  llvm_unreachable("Unhandled case in switch");
+}

lib/AST/NameLookup.cpp

@@ -1565,8 +1565,11 @@ bool DeclContext::lookupQualified(Type type,
   assert(decls.empty() && "additive lookup not supported");
 
   // Handle AnyObject lookup.
-  if (type->isAnyObject())
-    return lookupAnyObject(member, options, decls);
+  if (type->isAnyObject()) {
+    AnyObjectLookupRequest req(this, member, options);
+    decls = evaluateOrDefault(getASTContext().evaluator, req, {});
+    return !decls.empty();
+  }
 
   // Handle lookup in a module.
   if (auto moduleTy = type->getAs<ModuleType>())
@@ -1798,31 +1801,28 @@ bool DeclContext::lookupQualified(ModuleDecl *module, DeclName member,
   return !decls.empty();
 }
 
-bool DeclContext::lookupAnyObject(DeclName member, NLOptions options,
-                                  SmallVectorImpl<ValueDecl *> &decls) const {
+llvm::Expected<QualifiedLookupResult>
+AnyObjectLookupRequest::evaluate(Evaluator &evaluator, const DeclContext *dc,
+                                 DeclName member, NLOptions options) const {
   using namespace namelookup;
-  assert(decls.empty() && "additive lookup not supported");
+  QualifiedLookupResult decls;
 
   // Configure lookup and dig out the tracker.
   ReferencedNameTracker *tracker = nullptr;
   bool isLookupCascading;
-  configureLookup(this, options, tracker, isLookupCascading);
+  configureLookup(dc, options, tracker, isLookupCascading);
 
   // Record this lookup.
   if (tracker)
     tracker->addDynamicLookupName(member.getBaseName(), isLookupCascading);
 
   // Type-only lookup won't find anything on AnyObject.
   if (options & NL_OnlyTypes)
-    return false;
-
-  auto *stats = getASTContext().Stats;
-  if (stats)
-    stats->getFrontendCounters().NumLookupQualifiedInAnyObject++;
+    return decls;
 
   // Collect all of the visible declarations.
   SmallVector<ValueDecl *, 4> allDecls;
-  for (auto import : namelookup::getAllImports(this)) {
+  for (auto import : namelookup::getAllImports(dc)) {
     import.second->lookupClassMember(import.first, member, allDecls);
   }
 
@@ -1840,26 +1840,23 @@ bool DeclContext::lookupAnyObject(DeclName member, NLOptions options,
     if (decl->getOverriddenDecl())
       continue;
 
-    auto dc = decl->getDeclContext();
-    auto nominal = dc->getSelfNominalTypeDecl();
-    assert(nominal && "Couldn't find nominal type?");
-    (void)nominal;
+    assert(decl->getDeclContext()->isTypeContext() &&
+           "Couldn't find nominal type?");
 
     // If we didn't see this declaration before, and it's an acceptable
     // result, add it to the list.
     // declaration to the list.
     if (knownDecls.insert(decl).second &&
-        isAcceptableLookupResult(this, options, decl,
+        isAcceptableLookupResult(dc, options, decl,
                                  /*onlyCompleteObjectInits=*/false))
       decls.push_back(decl);
   }
 
-  pruneLookupResultSet(this, options, decls);
-  if (auto *debugClient = this->getParentModule()->getDebugClient()) {
-    debugClient->finishLookupInAnyObject(this, member, options, decls);
+  pruneLookupResultSet(dc, options, decls);
+  if (auto *debugClient = dc->getParentModule()->getDebugClient()) {
+    debugClient->finishLookupInAnyObject(dc, member, options, decls);
   }
-  // We're done. Report success/failure.
-  return !decls.empty();
+  return decls;
 }
 
 void DeclContext::lookupAllObjCMethods(
@@ -2635,3 +2632,40 @@ void FindLocalVal::visitCatchStmt(CatchStmt *S) {
     checkPattern(S->getErrorPattern(), DeclVisibilityKind::LocalVariable);
   visit(S->getBody());
 }
+
+void swift::simple_display(llvm::raw_ostream &out, NLKind kind) {
+  switch (kind) {
+  case NLKind::QualifiedLookup:
+    out << "QualifiedLookup";
+    return;
+  case NLKind::UnqualifiedLookup:
+    out << "UnqualifiedLookup";
+    return;
+  }
+  llvm_unreachable("Unhandled case in switch");
+}
+
+void swift::simple_display(llvm::raw_ostream &out, NLOptions options) {
+  using Flag = std::pair<NLOptions, StringRef>;
+  Flag possibleFlags[] = {
+#define FLAG(Name) {Name, #Name},
+    FLAG(NL_ProtocolMembers)
+    FLAG(NL_RemoveNonVisible)
+    FLAG(NL_RemoveOverridden)
+    FLAG(NL_IgnoreAccessControl)
+    FLAG(NL_KnownNonCascadingDependency)
+    FLAG(NL_KnownCascadingDependency)
+    FLAG(NL_OnlyTypes)
+    FLAG(NL_IncludeAttributeImplements)
+#undef FLAG
+  };
+
+  auto flagsToPrint = llvm::make_filter_range(
+      possibleFlags, [&](Flag flag) { return options & flag.first; });
+
+  out << "{ ";
+  interleave(
+      flagsToPrint, [&](Flag flag) { out << flag.second; },
+      [&] { out << ", "; });
+  out << " }";
+}