AST: Introduce namelookup::getAllImports() to replace forAllVisibleModules()

slavapestov · slavapestov · commit 2256b1fcc230 · 2019-08-29T15:57:38.000-04:00
diff --git a/include/swift/AST/ImportCache.h b/include/swift/AST/ImportCache.h
@@ -51,7 +51,7 @@ class ImportSet final :
     public llvm::FoldingSetNode,
     private llvm::TrailingObjects<ImportSet, ModuleDecl::ImportedModule> {
   friend TrailingObjects;
-  friend ImportCache;
+  friend class ImportCache;
 
   unsigned HasHeaderImportModule : 1;
   unsigned NumTopLevelImports : 31;
@@ -193,6 +193,8 @@ class alignas(ModuleDecl::ImportedModule) ImportCache {
   }
 };
 
+ArrayRef<ModuleDecl::ImportedModule> getAllImports(const DeclContext *dc);
+
 }  // namespace namelookup
 
 }  // namespace swift
diff --git a/include/swift/AST/Module.h b/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();
diff --git a/include/swift/AST/ModuleNameLookup.h b/include/swift/AST/ModuleNameLookup.h
@@ -59,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
diff --git a/lib/AST/ImportCache.cpp b/lib/AST/ImportCache.cpp
@@ -291,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();
+}
diff --git a/lib/AST/Module.cpp b/lib/AST/Module.cpp
@@ -1431,32 +1431,6 @@ forAllImportedModules(const ModuleDecl *topLevel,
   return true;
 }
 
-bool
-ModuleDecl::forAllVisibleModules(AccessPathTy thisPath,
-                                 llvm::function_ref<bool(ImportedModule)> fn) const {
-  return forAllImportedModules<true>(this, thisPath, fn);
-}
-
-bool FileUnit::forAllVisibleModules(
-    llvm::function_ref<bool(ModuleDecl::ImportedModule)> fn) const {
-  if (!getParentModule()->forAllVisibleModules(ModuleDecl::AccessPathTy(), fn))
-    return false;
-
-  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;
-  }
-
-  return true;
-}
-
 void ModuleDecl::collectLinkLibraries(LinkLibraryCallback callback) const {
   // FIXME: The proper way to do this depends on the decls used.
   FORWARD(collectLinkLibraries, (callback));
diff --git a/lib/AST/NameLookup.cpp b/lib/AST/NameLookup.cpp
@@ -1706,9 +1706,9 @@ bool DeclContext::lookupAnyObject(DeclName member, NLOptions options,
 
   // Collect all of the visible declarations.
   SmallVector<ValueDecl *, 4> allDecls;
-  forAllVisibleModules(this, [&](ModuleDecl::ImportedModule import) {
+  for (auto import : namelookup::getAllImports(this)) {
     import.second->lookupClassMember(import.first, member, allDecls);
-  });
+  }
 
   // For each declaration whose context is not something we've
   // already visited above, add it to the list of declarations.
@@ -1750,9 +1750,9 @@ void DeclContext::lookupAllObjCMethods(
        ObjCSelector selector,
        SmallVectorImpl<AbstractFunctionDecl *> &results) const {
   // Collect all of the methods with this selector.
-  forAllVisibleModules(this, [&](ModuleDecl::ImportedModule import) {
+  for (auto import : namelookup::getAllImports(this)) {
     import.second->lookupObjCMethods(selector, results);
-  });
+  }
 
   // Filter out duplicates.
   llvm::SmallPtrSet<AbstractFunctionDecl *, 8> visited;
diff --git a/lib/ClangImporter/ClangImporter.cpp b/lib/ClangImporter/ClangImporter.cpp
@@ -22,6 +22,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsClangImporter.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/AST/IRGenOptions.h"
 #include "swift/AST/LinkLibrary.h"
 #include "swift/AST/Module.h"
@@ -1684,7 +1685,7 @@ ModuleDecl *ClangImporter::Implementation::finishLoadingClangModule(
       // FIXME: This forces the creation of wrapper modules for all imports as
       // well, and may do unnecessary work.
       cacheEntry.setInt(true);
-      result->forAllVisibleModules({}, [&](ModuleDecl::ImportedModule import) {});
+      (void) namelookup::getAllImports(result);
     }
   } else {
     // Build the representation of the Clang module in Swift.
@@ -1704,7 +1705,7 @@ ModuleDecl *ClangImporter::Implementation::finishLoadingClangModule(
     // Force load overlays for all imported modules.
     // FIXME: This forces the creation of wrapper modules for all imports as
     // well, and may do unnecessary work.
-    result->forAllVisibleModules({}, [](ModuleDecl::ImportedModule import) {});
+    (void) namelookup::getAllImports(result);
   }
 
   if (clangModule->isSubModule()) {
diff --git a/lib/ClangImporter/DWARFImporter.cpp b/lib/ClangImporter/DWARFImporter.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "ImporterImpl.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/AST/Module.h"
 
 using namespace swift;
@@ -118,7 +119,7 @@ ModuleDecl *ClangImporter::Implementation::loadModuleDWARF(
   decl->addFile(*wrapperUnit);
 
   // Force load overlay modules for all imported modules.
-  decl->forAllVisibleModules({}, [](ModuleDecl::ImportedModule import) {});
+  (void) namelookup::getAllImports(decl);
 
   // Register the module with the ASTContext so it is available for lookups.
   ModuleDecl *&loaded = SwiftContext.LoadedModules[name];
diff --git a/lib/IDE/CodeCompletion.cpp b/lib/IDE/CodeCompletion.cpp
@@ -16,6 +16,7 @@
 #include "swift/AST/ASTPrinter.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Comment.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/LazyResolver.h"
@@ -3357,8 +3358,7 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   std::vector<OperatorDecl *> collectOperators() {
     std::vector<OperatorDecl *> results;
     assert(CurrDeclContext);
-    CurrDeclContext->getParentSourceFile()->forAllVisibleModules(
-    [&](ModuleDecl::ImportedModule import) {
+    for (auto import : namelookup::getAllImports(CurrDeclContext)) {
       for (auto fileUnit : import.second->getFiles()) {
         switch (fileUnit->getKind()) {
         case FileUnitKind::Builtin:
@@ -3373,7 +3373,7 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
           break;
         }
       }
-    });
+    }
     return results;
   }
 
@@ -3960,15 +3960,14 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
             addPrecedenceGroupRef(PG);
       }
     }
-    CurrDeclContext->getParentSourceFile()
-      ->forAllVisibleModules([&](ModuleDecl::ImportedModule Import) {
+    for (auto Import : namelookup::getAllImports(CurrDeclContext)) {
       auto Module = Import.second;
       if (Module == CurrModule)
-        return;
+        continue;
 
       RequestedCachedResults.push_back(
         RequestedResultsTy::fromModule(Module).onlyPrecedenceGroups());
-    });
+    }
   }
 
   void getPrecedenceGroupCompletions(SyntaxKind SK) {
@@ -5536,8 +5535,9 @@ void CodeCompletionCallbacksImpl::doneParsing() {
 
     if (Request.TheModule) {
       // FIXME: actually check imports.
-      const_cast<ModuleDecl*>(Request.TheModule)
-          ->forAllVisibleModules({}, handleImport);
+      for (auto Import : namelookup::getAllImports(Request.TheModule)) {
+        handleImport(Import);
+      }
     } else {
       // Add results from current module.
       Lookup.getToplevelCompletions(Request.OnlyTypes);
@@ -5552,9 +5552,8 @@ void CodeCompletionCallbacksImpl::doneParsing() {
       SF->getImportedModules(Imports, ImportFilter);
 
       for (auto Imported : Imports) {
-        ModuleDecl *TheModule = Imported.second;
-        ModuleDecl::AccessPathTy AccessPath = Imported.first;
-        TheModule->forAllVisibleModules(AccessPath, handleImport);
+        for (auto Import : namelookup::getAllImports(Imported.second))
+          handleImport(Import);
       }
     }
   }
diff --git a/lib/Sema/LookupVisibleDecls.cpp b/lib/Sema/LookupVisibleDecls.cpp
@@ -18,6 +18,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/GenericSignatureBuilder.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/LazyResolver.h"
 #include "swift/AST/ModuleNameLookup.h"
@@ -364,10 +365,9 @@ static void doDynamicLookup(VisibleDeclConsumer &Consumer,
 
   DynamicLookupConsumer ConsumerWrapper(Consumer, LS, CurrDC);
 
-  CurrDC->getParentSourceFile()->forAllVisibleModules(
-      [&](ModuleDecl::ImportedModule Import) {
-        Import.second->lookupClassMembers(Import.first, ConsumerWrapper);
-      });
+  for (auto Import : namelookup::getAllImports(CurrDC)) {
+    Import.second->lookupClassMembers(Import.first, ConsumerWrapper);
+  }
 }
 
 namespace {
diff --git a/lib/Sema/TypeCheckDeclObjC.cpp b/lib/Sema/TypeCheckDeclObjC.cpp
@@ -21,6 +21,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/ForeignErrorConvention.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/StringExtras.h"
@@ -1795,10 +1796,12 @@ void swift::diagnoseAttrsRequiringFoundation(SourceFile &SF) {
       return;
   }
 
-  SF.forAllVisibleModules([&](ModuleDecl::ImportedModule import) {
-    if (import.second->getName() == Ctx.Id_Foundation)
+  for (auto import : namelookup::getAllImports(&SF)) {
+    if (import.second->getName() == Ctx.Id_Foundation) {
       ImportsFoundationModule = true;
-  });
+      break;
+    }
+  }
 
   if (ImportsFoundationModule)
     return;
diff --git a/lib/Sema/TypeChecker.cpp b/lib/Sema/TypeChecker.cpp
@@ -27,6 +27,7 @@
 #include "swift/AST/DiagnosticSuppression.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/Identifier.h"
+#include "swift/AST/ImportCache.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/ModuleLoader.h"
 #include "swift/AST/NameLookup.h"
@@ -249,7 +250,7 @@ static void bindExtensions(SourceFile &SF) {
 
   // FIXME: The current source file needs to be handled specially, because of
   // private extensions.
-  SF.forAllVisibleModules([&](ModuleDecl::ImportedModule import) {
+  for (auto import : namelookup::getAllImports(&SF)) {
     // FIXME: Respect the access path?
     for (auto file : import.second->getFiles()) {
       auto SF = dyn_cast<SourceFile>(file);
@@ -262,7 +263,7 @@ static void bindExtensions(SourceFile &SF) {
             worklist.push_back(ED);
       }
     }
-  });
+  }
 
   // Phase 2 - repeatedly go through the worklist and attempt to bind each
   // extension there, removing it from the worklist if we succeed.
diff --git a/tools/swift-ide-test/swift-ide-test.cpp b/tools/swift-ide-test/swift-ide-test.cpp
