Change ModuleDecl::getImportedModules to take an option set

jrose-apple · jrose-apple · commit 9ed3fe061d5b · 2019-03-28T14:44:41.000-07:00
...in preparation for me adding a third kind of import, making the
existing "All" kind a problem. NFC, except that I did rewrite the
ClangModuleUnit implementation of getImportedModules to be simpler!
diff --git a/include/swift/AST/Module.h b/include/swift/AST/Module.h
@@ -419,18 +419,19 @@ class ModuleDecl : public DeclContext, public TypeDecl {
          SmallVectorImpl<AbstractFunctionDecl *> &results) const;
 
   /// \sa getImportedModules
-  enum class ImportFilter {
-    All,
-    Public,
-    Private
+  enum class ImportFilterKind {
+    Public = 1 << 0,
+    Private = 1 << 1
   };
+  /// \sa getImportedModules
+  using ImportFilter = OptionSet<ImportFilterKind>;
 
   /// Looks up which modules are imported by this module.
   ///
   /// \p filter controls whether public, private, or any imports are included
   /// in this list.
   void getImportedModules(SmallVectorImpl<ImportedModule> &imports,
-                          ImportFilter filter = ImportFilter::Public) const;
+                          ImportFilter filter = ImportFilterKind::Public) const;
 
   /// Looks up which modules are imported by this module, ignoring any that
   /// won't contain top-level decls.
@@ -758,7 +759,7 @@ class FileUnit : public DeclContext {
   /// \see ModuleDecl::getImportedModulesForLookup
   virtual void getImportedModulesForLookup(
       SmallVectorImpl<ModuleDecl::ImportedModule> &imports) const {
-    return getImportedModules(imports, ModuleDecl::ImportFilter::Public);
+    return getImportedModules(imports, ModuleDecl::ImportFilterKind::Public);
   }
 
   /// Generates the list of libraries needed to link this file, based on its
diff --git a/lib/AST/Module.cpp b/lib/AST/Module.cpp
@@ -1021,17 +1021,12 @@ SourceFile::getImportedModules(SmallVectorImpl<ModuleDecl::ImportedModule> &modu
                                ModuleDecl::ImportFilter filter) const {
   assert(ASTStage >= Parsed || Kind == SourceFileKind::SIL);
   for (auto desc : Imports) {
-    switch (filter) {
-    case ModuleDecl::ImportFilter::All:
-      break;
-    case ModuleDecl::ImportFilter::Public:
-      if (!desc.importOptions.contains(ImportFlags::Exported))
+    if (desc.importOptions.contains(ImportFlags::Exported)) {
+      if (!filter.contains(ModuleDecl::ImportFilterKind::Public))
         continue;
-      break;
-    case ModuleDecl::ImportFilter::Private:
-      if (desc.importOptions.contains(ImportFlags::Exported))
+    } else {
+      if (!filter.contains(ModuleDecl::ImportFilterKind::Private))
         continue;
-      break;
     }
 
     modules.push_back(desc.module);
@@ -1277,8 +1272,9 @@ forAllImportedModules(ModuleDecl *topLevel, ModuleDecl::AccessPathTy thisPath,
   llvm::SmallSet<ImportedModule, 32, ModuleDecl::OrderImportedModules> visited;
   SmallVector<ImportedModule, 32> stack;
 
-  auto filter = respectVisibility ? ModuleDecl::ImportFilter::Public
-                                  : ModuleDecl::ImportFilter::All;
+  ModuleDecl::ImportFilter filter = ModuleDecl::ImportFilterKind::Public;
+  if (!respectVisibility)
+    filter |= ModuleDecl::ImportFilterKind::Private;
   topLevel->getImportedModules(stack, filter);
 
   // Make sure the top-level module is first; we want pre-order-ish traversal.
@@ -1332,7 +1328,7 @@ bool FileUnit::forAllVisibleModules(
     // Handle privately visible modules as well.
     // FIXME: Should this apply to all FileUnits?
     SmallVector<ModuleDecl::ImportedModule, 4> imports;
-    SF->getImportedModules(imports, ModuleDecl::ImportFilter::Private);
+    SF->getImportedModules(imports, ModuleDecl::ImportFilterKind::Private);
     for (auto importPair : imports)
       if (!importPair.second->forAllVisibleModules(importPair.first, fn))
         return false;
diff --git a/lib/AST/UnqualifiedLookup.cpp b/lib/AST/UnqualifiedLookup.cpp
@@ -1045,7 +1045,7 @@ void UnqualifiedLookupFactory::addImportedResults(DeclContext *const dc) {
   // Add private imports to the extra search list.
   SmallVector<ModuleDecl::ImportedModule, 8> extraImports;
   if (auto FU = dyn_cast<FileUnit>(dc))
-    FU->getImportedModules(extraImports, ModuleDecl::ImportFilter::Private);
+    FU->getImportedModules(extraImports, ModuleDecl::ImportFilterKind::Private);
 
   using namespace namelookup;
   SmallVector<ValueDecl *, 8> CurModuleResults;
diff --git a/lib/ClangImporter/ClangImporter.cpp b/lib/ClangImporter/ClangImporter.cpp
@@ -3203,71 +3203,38 @@ ModuleDecl *ClangModuleUnit::getAdapterModule() const {
 void ClangModuleUnit::getImportedModules(
     SmallVectorImpl<ModuleDecl::ImportedModule> &imports,
     ModuleDecl::ImportFilter filter) const {
-  switch (filter) {
-  case ModuleDecl::ImportFilter::All:
-  case ModuleDecl::ImportFilter::Private:
+  if (filter.contains(ModuleDecl::ImportFilterKind::Private))
     if (auto stdlib = owner.getStdlibModule())
       imports.push_back({ModuleDecl::AccessPathTy(), stdlib});
-    break;
-  case ModuleDecl::ImportFilter::Public:
-    break;
-  }
 
   SmallVector<clang::Module *, 8> imported;
   if (!clangModule) {
     // This is the special "imported headers" module.
-    switch (filter) {
-    case ModuleDecl::ImportFilter::All:
-    case ModuleDecl::ImportFilter::Public:
+    if (filter.contains(ModuleDecl::ImportFilterKind::Public)) {
       imported.append(owner.ImportedHeaderExports.begin(),
                       owner.ImportedHeaderExports.end());
-      break;
-
-    case ModuleDecl::ImportFilter::Private:
-      break;
     }
 
   } else {
     clangModule->getExportedModules(imported);
 
-    switch (filter) {
-    case ModuleDecl::ImportFilter::All: {
-      llvm::SmallPtrSet<clang::Module *, 8> knownModules;
-      imported.append(clangModule->Imports.begin(), clangModule->Imports.end());
-      imported.erase(std::remove_if(imported.begin(), imported.end(),
-                                    [&](clang::Module *mod) -> bool {
-                                      return !knownModules.insert(mod).second;
-                                    }),
-                     imported.end());
-
-      // FIXME: The parent module isn't exactly a private import, but it is
-      // needed for link dependencies.
-      if (clangModule->Parent)
-        imported.push_back(clangModule->Parent);
-
-      break;
-    }
-
-    case ModuleDecl::ImportFilter::Private: {
+    if (filter.contains(ModuleDecl::ImportFilterKind::Private)) {
+      // Copy in any modules that are imported but not exported.
       llvm::SmallPtrSet<clang::Module *, 8> knownModules(imported.begin(),
                                                          imported.end());
-      SmallVector<clang::Module *, 8> privateImports;
-      std::copy_if(clangModule->Imports.begin(), clangModule->Imports.end(),
-                   std::back_inserter(privateImports), [&](clang::Module *mod) {
-                     return knownModules.count(mod) == 0;
-                   });
-      imported.swap(privateImports);
+      if (!filter.contains(ModuleDecl::ImportFilterKind::Public)) {
+        // Remove the exported ones now that we're done with them.
+        imported.clear();
+      }
+      llvm::copy_if(clangModule->Imports, std::back_inserter(imported),
+                    [&](clang::Module *mod) {
+                     return !knownModules.insert(mod).second;
+                    });
 
       // FIXME: The parent module isn't exactly a private import, but it is
       // needed for link dependencies.
       if (clangModule->Parent)
         imported.push_back(clangModule->Parent);
-
-      break;
-    }
-
-    case ModuleDecl::ImportFilter::Public:
-      break;
     }
   }
 
diff --git a/lib/Frontend/ParseableInterfaceSupport.cpp b/lib/Frontend/ParseableInterfaceSupport.cpp
@@ -86,15 +86,19 @@ llvm::Regex swift::getSwiftInterfaceModuleFlagsRegex() {
 static void printImports(raw_ostream &out, ModuleDecl *M) {
   // FIXME: This is very similar to what's in Serializer::writeInputBlock, but
   // it's not obvious what higher-level optimization would be factored out here.
+  ModuleDecl::ImportFilter allImportFilter;
+  allImportFilter |= ModuleDecl::ImportFilterKind::Public;
+  allImportFilter |= ModuleDecl::ImportFilterKind::Private;
+
   SmallVector<ModuleDecl::ImportedModule, 8> allImports;
-  M->getImportedModules(allImports, ModuleDecl::ImportFilter::All);
+  M->getImportedModules(allImports, allImportFilter);
   ModuleDecl::removeDuplicateImports(allImports);
   diagnoseScopedImports(M->getASTContext().Diags, allImports);
 
   // Collect the public imports as a subset so that we can mark them with
   // '@_exported'.
   SmallVector<ModuleDecl::ImportedModule, 8> publicImports;
-  M->getImportedModules(publicImports, ModuleDecl::ImportFilter::Public);
+  M->getImportedModules(publicImports, ModuleDecl::ImportFilterKind::Public);
   llvm::SmallSet<ModuleDecl::ImportedModule, 8,
                  ModuleDecl::OrderImportedModules> publicImportSet;
   publicImportSet.insert(publicImports.begin(), publicImports.end());
diff --git a/lib/FrontendTool/ImportedModules.cpp b/lib/FrontendTool/ImportedModules.cpp
@@ -79,9 +79,13 @@ bool swift::emitImportedModules(ASTContext &Context, ModuleDecl *mainModule,
   StringRef implicitHeaderPath = opts.ImplicitObjCHeaderPath;
   if (!implicitHeaderPath.empty()) {
     if (!clangImporter->importBridgingHeader(implicitHeaderPath, mainModule)) {
+      ModuleDecl::ImportFilter importFilter;
+      importFilter |= ModuleDecl::ImportFilterKind::Public;
+      importFilter |= ModuleDecl::ImportFilterKind::Private;
+
       SmallVector<ModuleDecl::ImportedModule, 16> imported;
       clangImporter->getImportedHeaderModule()->getImportedModules(
-          imported, ModuleDecl::ImportFilter::All);
+          imported, importFilter);
 
       for (auto IM : imported) {
         if (auto clangModule = IM.second->findUnderlyingClangModule())
diff --git a/lib/IDE/CodeCompletion.cpp b/lib/IDE/CodeCompletion.cpp
@@ -1707,17 +1707,22 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   }
 
   void collectImportedModules(llvm::StringSet<> &ImportedModules) {
+    ModuleDecl::ImportFilter ImportFilter;
+    ImportFilter |= ModuleDecl::ImportFilterKind::Public;
+    ImportFilter |= ModuleDecl::ImportFilterKind::Private;
+
     SmallVector<ModuleDecl::ImportedModule, 16> Imported;
     SmallVector<ModuleDecl::ImportedModule, 16> FurtherImported;
     CurrDeclContext->getParentSourceFile()->getImportedModules(Imported,
-      ModuleDecl::ImportFilter::All);
+                                                               ImportFilter);
     while (!Imported.empty()) {
       ModuleDecl *MD = Imported.back().second;
       Imported.pop_back();
       if (!ImportedModules.insert(MD->getNameStr()).second)
         continue;
       FurtherImported.clear();
-      MD->getImportedModules(FurtherImported, ModuleDecl::ImportFilter::Public);
+      MD->getImportedModules(FurtherImported,
+                             ModuleDecl::ImportFilterKind::Public);
       Imported.append(FurtherImported.begin(), FurtherImported.end());
       for (auto SubMod : FurtherImported) {
         Imported.push_back(SubMod);
@@ -5331,9 +5336,12 @@ void CodeCompletionCallbacksImpl::doneParsing() {
       Lookup.getToplevelCompletions(Request.OnlyTypes);
 
       // Add results for all imported modules.
+      ModuleDecl::ImportFilter ImportFilter;
+      ImportFilter |= ModuleDecl::ImportFilterKind::Public;
+      ImportFilter |= ModuleDecl::ImportFilterKind::Private;
       SmallVector<ModuleDecl::ImportedModule, 4> Imports;
       auto *SF = CurDeclContext->getParentSourceFile();
-      SF->getImportedModules(Imports, ModuleDecl::ImportFilter::All);
+      SF->getImportedModules(Imports, ImportFilter);
 
       for (auto Imported : Imports) {
         ModuleDecl *TheModule = Imported.second;
diff --git a/lib/IRGen/IRGenDebugInfo.cpp b/lib/IRGen/IRGenDebugInfo.cpp
@@ -1741,9 +1741,11 @@ void IRGenDebugInfoImpl::finalize() {
 
   // Get the list of imported modules (which may actually be different
   // from all ImportDecls).
+  ModuleDecl::ImportFilter ImportFilter;
+  ImportFilter |= ModuleDecl::ImportFilterKind::Public;
+  ImportFilter |= ModuleDecl::ImportFilterKind::Private;
   SmallVector<ModuleDecl::ImportedModule, 8> ModuleWideImports;
-  IGM.getSwiftModule()->getImportedModules(ModuleWideImports,
-                                           ModuleDecl::ImportFilter::All);
+  IGM.getSwiftModule()->getImportedModules(ModuleWideImports, ImportFilter);
   for (auto M : ModuleWideImports)
     if (!ImportedModules.count(M.second))
       DBuilder.createImportedModule(MainFile, getOrCreateModule(M), MainFile,
diff --git a/lib/Immediate/REPL.cpp b/lib/Immediate/REPL.cpp
@@ -176,7 +176,7 @@ typeCheckREPLInput(ModuleDecl *MostRecentModule, StringRef Name,
 
   SmallVector<ModuleDecl::ImportedModule, 8> Imports;
   MostRecentModule->getImportedModules(Imports,
-                                       ModuleDecl::ImportFilter::Private);
+                                       ModuleDecl::ImportFilterKind::Private);
   if (!Imports.empty()) {
     SmallVector<SourceFile::ImportedModuleDesc, 8> ImportsWithOptions;
     for (auto Import : Imports) {
diff --git a/lib/Index/Index.cpp b/lib/Index/Index.cpp
@@ -120,11 +120,14 @@ class SourceFileOrModule {
 
   void
   getImportedModules(SmallVectorImpl<ModuleDecl::ImportedModule> &Modules) const {
+    ModuleDecl::ImportFilter ImportFilter;
+    ImportFilter |= ModuleDecl::ImportFilterKind::Public;
+    ImportFilter |= ModuleDecl::ImportFilterKind::Private;
+
     if (auto *SF = SFOrMod.dyn_cast<SourceFile *>()) {
-      SF->getImportedModules(Modules, ModuleDecl::ImportFilter::All);
+      SF->getImportedModules(Modules, ImportFilter);
     } else {
-      SFOrMod.get<ModuleDecl *>()->getImportedModules(Modules,
-                                                  ModuleDecl::ImportFilter::All);
+      SFOrMod.get<ModuleDecl *>()->getImportedModules(Modules, ImportFilter);
     }
   }
 };
@@ -1536,8 +1539,11 @@ void IndexSwiftASTWalker::collectRecursiveModuleImports(
     return;
   }
 
+  ModuleDecl::ImportFilter ImportFilter;
+  ImportFilter |= ModuleDecl::ImportFilterKind::Public;
+  ImportFilter |= ModuleDecl::ImportFilterKind::Private;
   SmallVector<ModuleDecl::ImportedModule, 8> Imports;
-  TopMod.getImportedModules(Imports, ModuleDecl::ImportFilter::All);
+  TopMod.getImportedModules(Imports);
 
   for (auto Import : Imports) {
     collectRecursiveModuleImports(*Import.second, Visited);
diff --git a/lib/Index/IndexRecord.cpp b/lib/Index/IndexRecord.cpp
@@ -562,8 +562,11 @@ emitDataForSwiftSerializedModule(ModuleDecl *module,
     unitWriter.addRecordFile(recordFile, FE, isSystemModule, mod);
   }
 
+  ModuleDecl::ImportFilter importFilter;
+  importFilter |= ModuleDecl::ImportFilterKind::Public;
+  importFilter |= ModuleDecl::ImportFilterKind::Private;
   SmallVector<ModuleDecl::ImportedModule, 8> imports;
-  module->getImportedModules(imports, ModuleDecl::ImportFilter::All);
+  module->getImportedModules(imports, importFilter);
   StringScratchSpace moduleNameScratch;
   addModuleDependencies(imports, indexStorePath, indexSystemModules,
                         targetTriple, clangCI, diags, unitWriter, moduleNameScratch);
@@ -597,8 +600,11 @@ recordSourceFileUnit(SourceFile *primarySourceFile, StringRef indexUnitToken,
     targetTriple, sysrootPath, getModuleInfoFromOpaqueModule);
 
   // Module dependencies.
+  ModuleDecl::ImportFilter importFilter;
+  importFilter |= ModuleDecl::ImportFilterKind::Public;
+  importFilter |= ModuleDecl::ImportFilterKind::Private;
   SmallVector<ModuleDecl::ImportedModule, 8> imports;
-  primarySourceFile->getImportedModules(imports, ModuleDecl::ImportFilter::All);
+  primarySourceFile->getImportedModules(imports, importFilter);
   StringScratchSpace moduleNameScratch;
   addModuleDependencies(imports, indexStorePath, indexSystemModules,
                         targetTriple, clangCI, diags, unitWriter, moduleNameScratch);
diff --git a/lib/Sema/LookupVisibleDecls.cpp b/lib/Sema/LookupVisibleDecls.cpp
@@ -968,7 +968,8 @@ static void lookupVisibleDeclsImpl(VisibleDeclConsumer &Consumer,
         return;
       }
 
-      SF->getImportedModules(extraImports, ModuleDecl::ImportFilter::Private);
+      SF->getImportedModules(extraImports,
+                             ModuleDecl::ImportFilterKind::Private);
     }
   }
 
diff --git a/lib/Serialization/ModuleFile.cpp b/lib/Serialization/ModuleFile.cpp
@@ -1695,24 +1695,12 @@ void ModuleFile::getImportedModules(
   PrettyStackTraceModuleFile stackEntry(*this);
 
   for (auto &dep : Dependencies) {
-    switch (filter) {
-    case ModuleDecl::ImportFilter::All:
-      // We're including all imports.
-      break;
-
-    case ModuleDecl::ImportFilter::Private:
-      // Skip @_exported imports.
-      if (dep.isExported())
+    if (dep.isExported()) {
+      if (!filter.contains(ModuleDecl::ImportFilterKind::Public))
         continue;
-
-      break;
-
-    case ModuleDecl::ImportFilter::Public:
-      // Only include @_exported imports.
-      if (!dep.isExported())
+    } else {
+      if (!filter.contains(ModuleDecl::ImportFilterKind::Private))
         continue;
-
-      break;
     }
 
     assert(dep.isLoaded());
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -1068,14 +1068,17 @@ void Serializer::writeInputBlock(const SerializationOptions &options) {
                         dep.getPath());
   }
 
+  ModuleDecl::ImportFilter allImportFilter;
+  allImportFilter |= ModuleDecl::ImportFilterKind::Public;
+  allImportFilter |= ModuleDecl::ImportFilterKind::Private;
   SmallVector<ModuleDecl::ImportedModule, 8> allImports;
-  M->getImportedModules(allImports, ModuleDecl::ImportFilter::All);
+  M->getImportedModules(allImports, allImportFilter);
   ModuleDecl::removeDuplicateImports(allImports);
 
   // Collect the public imports as a subset so that we can mark them with an
   // extra flag.
   SmallVector<ModuleDecl::ImportedModule, 8> publicImports;
-  M->getImportedModules(publicImports, ModuleDecl::ImportFilter::Public);
+  M->getImportedModules(publicImports, ModuleDecl::ImportFilterKind::Public);
   llvm::SmallSet<ModuleDecl::ImportedModule, 8,
                  ModuleDecl::OrderImportedModules> publicImportSet;
   publicImportSet.insert(publicImports.begin(), publicImports.end());
diff --git a/lib/Serialization/SerializedModuleLoader.cpp b/lib/Serialization/SerializedModuleLoader.cpp
@@ -703,8 +703,12 @@ void SerializedASTFile::getImportedModules(
 
 void SerializedASTFile::collectLinkLibrariesFromImports(
     ModuleDecl::LinkLibraryCallback callback) const {
+  ModuleDecl::ImportFilter ImportFilter;
+  ImportFilter |= ModuleDecl::ImportFilterKind::Public;
+  ImportFilter |= ModuleDecl::ImportFilterKind::Private;
+
   llvm::SmallVector<ModuleDecl::ImportedModule, 8> Imports;
-  File.getImportedModules(Imports, ModuleDecl::ImportFilter::All);
+  File.getImportedModules(Imports, ImportFilter);
 
   for (auto Import : Imports)
     Import.second->collectLinkLibraries(callback);
diff --git a/tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp b/tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp b/tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp
diff --git a/tools/swift-ide-test/swift-ide-test.cpp b/tools/swift-ide-test/swift-ide-test.cpp
