Merge pull request swiftlang#26865 from slavapestov/faster-module-name-lookup

Faster module name lookup

Author: slavapestov

include/swift/AST/ImportCache.h

@@ -51,12 +51,14 @@ class ImportSet final :
     public llvm::FoldingSetNode,
     private llvm::TrailingObjects<ImportSet, ModuleDecl::ImportedModule> {
   friend TrailingObjects;
-  friend ImportCache;
+  friend class ImportCache;
 
-  unsigned NumTopLevelImports;
+  unsigned HasHeaderImportModule : 1;
+  unsigned NumTopLevelImports : 31;
   unsigned NumTransitiveImports;
 
-  ImportSet(ArrayRef<ModuleDecl::ImportedModule> topLevelImports,
+  ImportSet(bool hasHeaderImportModule,
+            ArrayRef<ModuleDecl::ImportedModule> topLevelImports,
             ArrayRef<ModuleDecl::ImportedModule> transitiveImports);
 
   ImportSet(const ImportSet &) = delete;
@@ -74,6 +76,13 @@ class ImportSet final :
     return NumTopLevelImports + NumTransitiveImports;
   }
 
+  /// This is a bit of a hack to make module name lookup work properly.
+  /// If our import set includes the ClangImporter's special header import
+  /// module, we have to check it first, before any other imported module.
+  bool hasHeaderImportModule() const {
+    return HasHeaderImportModule;
+  }
+
   ArrayRef<ModuleDecl::ImportedModule> getTopLevelImports() const {
     return {getTrailingObjects<ModuleDecl::ImportedModule>(),
             NumTopLevelImports};
@@ -184,6 +193,8 @@ class alignas(ModuleDecl::ImportedModule) ImportCache {
   }
 };
 
+ArrayRef<ModuleDecl::ImportedModule> getAllImports(const DeclContext *dc);
+
 }  // namespace namelookup
 
 }  // namespace swift

include/swift/AST/Module.h

@@ -519,45 +519,6 @@ class ModuleDecl : public DeclContext, public TypeDecl {
   /// shadowed clang module.
   void getDisplayDecls(SmallVectorImpl<Decl*> &results) const;
 
-  /// @{
-
-  /// Perform an action for every module visible from this module.
-  ///
-  /// This only includes modules with at least one declaration visible: if two
-  /// import access paths are incompatible, the indirect module will be skipped.
-  /// Modules that can't be used for lookup (including Clang submodules at the
-  /// time this comment was written) are also skipped under certain
-  /// circumstances.
-  ///
-  /// \param topLevelAccessPath If present, include the top-level module in the
-  ///        results, with the given access path.
-  /// \param fn A callback of type bool(ImportedModule) or void(ImportedModule).
-  ///        Return \c false to abort iteration.
-  ///
-  /// \return True if the traversal ran to completion, false if it ended early
-  ///         due to the callback.
-  bool forAllVisibleModules(AccessPathTy topLevelAccessPath,
-                            llvm::function_ref<bool(ImportedModule)> fn) const;
-
-  bool forAllVisibleModules(AccessPathTy topLevelAccessPath,
-                            llvm::function_ref<void(ImportedModule)> fn) const {
-    return forAllVisibleModules(topLevelAccessPath,
-                                [=](const ImportedModule &import) -> bool {
-      fn(import);
-      return true;
-    });
-  }
-
-  template <typename Fn>
-  bool forAllVisibleModules(AccessPathTy topLevelAccessPath,
-                            Fn &&fn) const {
-    using RetTy = typename std::result_of<Fn(ImportedModule)>::type;
-    llvm::function_ref<RetTy(ImportedModule)> wrapped{std::forward<Fn>(fn)};
-    return forAllVisibleModules(topLevelAccessPath, wrapped);
-  }
-
-  /// @}
-
   using LinkLibraryCallback = llvm::function_ref<void(LinkLibrary)>;
 
   /// Generate the list of libraries needed to link this module, based on its
@@ -817,39 +778,6 @@ class FileUnit : public DeclContext {
   virtual void
   collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const {}
 
-  /// @{
-
-  /// Perform an action for every module visible from this file.
-  ///
-  /// \param fn A callback of type bool(ImportedModule) or void(ImportedModule).
-  ///           Return \c false to abort iteration.
-  ///
-  /// \return True if the traversal ran to completion, false if it ended early
-  ///         due to the callback.
-  bool
-  forAllVisibleModules(
-      llvm::function_ref<bool(ModuleDecl::ImportedModule)> fn) const;
-
-  bool
-  forAllVisibleModules(
-      llvm::function_ref<void(ModuleDecl::ImportedModule)> fn) const {
-    return forAllVisibleModules([=](ModuleDecl::ImportedModule import) -> bool {
-      fn(import);
-      return true;
-    });
-  }
-  
-  template <typename Fn>
-  bool forAllVisibleModules(Fn &&fn) const {
-    using RetTy = typename std::result_of<Fn(ModuleDecl::ImportedModule)>::type;
-    llvm::function_ref<RetTy(ModuleDecl::ImportedModule)> wrapped{
-      std::forward<Fn>(fn)
-    };
-    return forAllVisibleModules(wrapped);
-  }
-
-  /// @}
-
   /// True if this file contains the main class for the module.
   bool hasMainClass() const {
     return getMainClass();

include/swift/AST/ModuleNameLookup.h

@@ -35,9 +35,6 @@ enum class ResolutionKind {
   /// If non-overloadable decls are returned, this indicates ambiguous lookup.
   Overloadable,
 
-  /// Lookup should match a single decl.
-  Exact,
-
   /// Lookup should match a single decl that declares a type.
   TypesOnly
 };
@@ -62,16 +59,6 @@ void lookupInModule(const DeclContext *moduleOrFile,
                     NLKind lookupKind, ResolutionKind resolutionKind,
                     const DeclContext *moduleScopeContext);
 
-template <typename Fn>
-void forAllVisibleModules(const DeclContext *DC, const Fn &fn) {
-  DeclContext *moduleScope = DC->getModuleScopeContext();
-  if (auto file = dyn_cast<FileUnit>(moduleScope))
-    file->forAllVisibleModules(fn);
-  else
-    cast<ModuleDecl>(moduleScope)
-        ->forAllVisibleModules(ModuleDecl::AccessPathTy(), fn);
-}
-
 /// Performs a qualified lookup into the given module and, if necessary, its
 /// reexports, observing proper shadowing rules.
 void

lib/AST/ImportCache.cpp

@@ -23,9 +23,11 @@
 using namespace swift;
 using namespace namelookup;
 
-ImportSet::ImportSet(ArrayRef<ModuleDecl::ImportedModule> topLevelImports,
+ImportSet::ImportSet(bool hasHeaderImportModule,
+                     ArrayRef<ModuleDecl::ImportedModule> topLevelImports,
                      ArrayRef<ModuleDecl::ImportedModule> transitiveImports)
-  : NumTopLevelImports(topLevelImports.size()),
+  : HasHeaderImportModule(hasHeaderImportModule),
+    NumTopLevelImports(topLevelImports.size()),
     NumTransitiveImports(transitiveImports.size()) {
   auto buffer = getTrailingObjects<ModuleDecl::ImportedModule>();
   std::uninitialized_copy(topLevelImports.begin(), topLevelImports.end(),
@@ -78,6 +80,11 @@ static void collectExports(ModuleDecl::ImportedModule next,
 ImportSet &
 ImportCache::getImportSet(ASTContext &ctx,
                           ArrayRef<ModuleDecl::ImportedModule> imports) {
+  bool hasHeaderImportModule = false;
+  ModuleDecl *headerImportModule = nullptr;
+  if (auto *loader = ctx.getClangModuleLoader())
+    headerImportModule = loader->getImportedHeaderModule();
+
   SmallVector<ModuleDecl::ImportedModule, 4> topLevelImports;
 
   SmallVector<ModuleDecl::ImportedModule, 4> transitiveImports;
@@ -88,6 +95,8 @@ ImportCache::getImportSet(ASTContext &ctx,
       continue;
 
     topLevelImports.push_back(next);
+    if (next.second == headerImportModule)
+      hasHeaderImportModule = true;
   }
 
   void *InsertPos = nullptr;
@@ -116,6 +125,9 @@ ImportCache::getImportSet(ASTContext &ctx,
       continue;
 
     transitiveImports.push_back(next);
+    if (next.second == headerImportModule)
+      hasHeaderImportModule = true;
+
     collectExports(next, stack);
   }
 
@@ -131,7 +143,9 @@ ImportCache::getImportSet(ASTContext &ctx,
     sizeof(ModuleDecl::ImportedModule) * transitiveImports.size(),
     alignof(ImportSet), AllocationArena::Permanent);
 
-  auto *result = new (mem) ImportSet(topLevelImports, transitiveImports);
+  auto *result = new (mem) ImportSet(hasHeaderImportModule,
+                                     topLevelImports,
+                                     transitiveImports);
   ImportSets.InsertNode(result, InsertPos);
 
   return *result;
@@ -277,4 +291,10 @@ ImportCache::getAllAccessPathsNotShadowedBy(const ModuleDecl *mod,
   auto result = allocateArray(ctx, accessPaths);
   ShadowCache[key] = result;
   return result;
-};
+};
+
+ArrayRef<ModuleDecl::ImportedModule>
+swift::namelookup::getAllImports(const DeclContext *dc) {
+  return dc->getASTContext().getImportCache().getImportSet(dc)
+    .getAllImports();
+}

lib/AST/Module.cpp

@@ -1375,110 +1375,47 @@ bool ModuleDecl::registerEntryPointFile(FileUnit *file, SourceLoc diagLoc,
   return true;
 }
 
-template<bool respectVisibility>
-static bool
-forAllImportedModules(const ModuleDecl *topLevel,
-                      ModuleDecl::AccessPathTy thisPath,
-                      llvm::function_ref<bool(ModuleDecl::ImportedModule)> fn) {
-  using ImportedModule = ModuleDecl::ImportedModule;
-  using AccessPathTy = ModuleDecl::AccessPathTy;
-  
-  llvm::SmallSet<ImportedModule, 32, ModuleDecl::OrderImportedModules> visited;
-  SmallVector<ImportedModule, 32> stack;
+void ModuleDecl::collectLinkLibraries(LinkLibraryCallback callback) const {
+  // FIXME: The proper way to do this depends on the decls used.
+  FORWARD(collectLinkLibraries, (callback));
+}
+
+void
+SourceFile::collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const {
+  llvm::SmallDenseSet<ModuleDecl *, 32> visited;
+  SmallVector<ModuleDecl::ImportedModule, 32> stack;
 
   ModuleDecl::ImportFilter filter = ModuleDecl::ImportFilterKind::Public;
-  if (!respectVisibility)
-    filter |= ModuleDecl::ImportFilterKind::Private;
+  filter |= ModuleDecl::ImportFilterKind::Private;
+
+  auto *topLevel = getParentModule();
 
   ModuleDecl::ImportFilter topLevelFilter = filter;
-  if (!respectVisibility)
-    topLevelFilter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
+  topLevelFilter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
   topLevel->getImportedModules(stack, topLevelFilter);
 
   // Make sure the top-level module is first; we want pre-order-ish traversal.
-  AccessPathTy overridingPath;
-  if (respectVisibility)
-    overridingPath = thisPath;
-  stack.emplace_back(overridingPath, const_cast<ModuleDecl *>(topLevel));
+  stack.emplace_back(ModuleDecl::AccessPathTy(),
+                     const_cast<ModuleDecl *>(topLevel));
 
   while (!stack.empty()) {
-    auto next = stack.pop_back_val();
-
-    // Filter any whole-module imports, and skip specific-decl imports if the
-    // import path doesn't match exactly.
-    if (next.first.empty() || !respectVisibility)
-      next.first = overridingPath;
-    else if (!overridingPath.empty() &&
-             !ModuleDecl::isSameAccessPath(next.first, overridingPath)) {
-      // If we ever allow importing non-top-level decls, it's possible the rule
-      // above isn't what we want.
-      assert(next.first.size() == 1 && "import of non-top-level decl");
-      continue;
-    }
+    auto next = stack.pop_back_val().second;
 
     if (!visited.insert(next).second)
       continue;
 
-    if (!fn(next))
-      return false;
-
-    if (respectVisibility)
-      next.second->getImportedModulesForLookup(stack);
-    else
-      next.second->getImportedModules(stack, filter);
-  }
-
-  return true;
-}
-
-bool
-ModuleDecl::forAllVisibleModules(AccessPathTy thisPath,
-                                 llvm::function_ref<bool(ImportedModule)> fn) const {
-  return forAllImportedModules<true>(this, thisPath, fn);
-}
+    if (next->getName() != getParentModule()->getName()) {
+      // Hack: Assume other REPL files already have their libraries linked.
+      if (!next->getFiles().empty())
+        if (auto *nextSource = dyn_cast<SourceFile>(next->getFiles().front()))
+          if (nextSource->Kind == SourceFileKind::REPL)
+            continue;
 
-bool FileUnit::forAllVisibleModules(
-    llvm::function_ref<bool(ModuleDecl::ImportedModule)> fn) const {
-  if (!getParentModule()->forAllVisibleModules(ModuleDecl::AccessPathTy(), fn))
-    return false;
+      next->collectLinkLibraries(callback);
+    }
 
-  if (auto SF = dyn_cast<SourceFile>(this)) {
-    // Handle privately visible modules as well.
-    ModuleDecl::ImportFilter importFilter;
-    importFilter |= ModuleDecl::ImportFilterKind::Private;
-    importFilter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
-    SmallVector<ModuleDecl::ImportedModule, 4> imports;
-    SF->getImportedModules(imports, importFilter);
-    for (auto importPair : imports)
-      if (!importPair.second->forAllVisibleModules(importPair.first, fn))
-        return false;
+    next->getImportedModules(stack, filter);
   }
-
-  return true;
-}
-
-void ModuleDecl::collectLinkLibraries(LinkLibraryCallback callback) const {
-  // FIXME: The proper way to do this depends on the decls used.
-  FORWARD(collectLinkLibraries, (callback));
-}
-
-void
-SourceFile::collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const {
-  forAllImportedModules<false>(getParentModule(), /*thisPath*/{},
-                               [=](ModuleDecl::ImportedModule import) -> bool {
-    swift::ModuleDecl *next = import.second;
-    if (next->getName() == getParentModule()->getName())
-      return true;
-
-    // Hack: Assume other REPL files already have their libraries linked.
-    if (!next->getFiles().empty())
-      if (auto *nextSource = dyn_cast<SourceFile>(next->getFiles().front()))
-        if (nextSource->Kind == SourceFileKind::REPL)
-          return true;
-
-    next->collectLinkLibraries(callback);
-    return true;
-  });
 }
 
 bool ModuleDecl::walk(ASTWalker &Walker) {