[NameLookup] Add LookupInModuleRequest

Author: hamishknight

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.

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,28 @@ 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;
+};
+
 #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,8 @@ 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)

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,12 +229,10 @@ void ModuleNameLookup<LookupStrategy>::lookupInModule(
               decls.end());
 }
 
-void namelookup::lookupInModule(const DeclContext *moduleOrFile,
-                                DeclName name,
-                                SmallVectorImpl<ValueDecl *> &decls,
-                                NLKind lookupKind,
-                                ResolutionKind resolutionKind,
-                                const DeclContext *moduleScopeContext) {
+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();
@@ -243,8 +242,23 @@ void namelookup::lookupInModule(const DeclContext *moduleOrFile,
 
   FrontendStatsTracer tracer(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) {
+  auto &ctx = moduleOrFile->getASTContext();
+  LookupInModuleRequest req(moduleOrFile, name, lookupKind, resolutionKind,
+                            moduleScopeContext);
+  auto results = evaluateOrDefault(ctx.evaluator, req, {});
+  decls.append(results.begin(), results.end());
 }
 
 void namelookup::lookupVisibleDeclsInModule(
@@ -260,3 +274,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

@@ -2635,3 +2635,15 @@ 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");
+}