AST: Fold forAllImportedModules() into collectLinkLibraries()

Author: slavapestov

lib/AST/Module.cpp

@@ -1375,84 +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);
-  }
+    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;
 
-  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);
+    }
 
-    next->collectLinkLibraries(callback);
-    return true;
-  });
+    next->getImportedModules(stack, filter);
+  }
 }
 
 bool ModuleDecl::walk(ASTWalker &Walker) {