Extends canImport to check for submodule availability

Author: ApolloZhu

include/swift/AST/ASTContext.h

@@ -357,8 +357,8 @@ class ASTContext final {
     DelayedPatternContexts;
 
   /// Cache of module names that fail the 'canImport' test in this context.
-  mutable llvm::SmallPtrSet<Identifier, 8> FailedModuleImportNames;
-  
+  mutable llvm::StringSet<> FailedModuleImportNames;
+
   /// Set if a `-module-alias` was passed. Used to store mapping between module aliases and
   /// their corresponding real names, and vice versa for a reverse lookup, which is needed to check
   /// if the module names appearing in source files are aliases or real names.
@@ -980,10 +980,10 @@ class ASTContext final {
   ///
   /// Note that even if this check succeeds, errors may still occur if the
   /// module is loaded in full.
-  bool canImportModuleImpl(ImportPath::Element ModulePath,
-                           llvm::VersionTuple version,
-                           bool underlyingVersion,
+  bool canImportModuleImpl(ImportPath::Module ModulePath,
+                           llvm::VersionTuple version, bool underlyingVersion,
                            bool updateFailingList) const;
+
 public:
   namelookup::ImportCache &getImportCache() const;
 
@@ -1012,10 +1012,10 @@ class ASTContext final {
   ///
   /// Note that even if this check succeeds, errors may still occur if the
   /// module is loaded in full.
-  bool canImportModule(ImportPath::Element ModulePath,
+  bool canImportModule(ImportPath::Module ModulePath,
                        llvm::VersionTuple version = llvm::VersionTuple(),
                        bool underlyingVersion = false);
-  bool canImportModule(ImportPath::Element ModulePath,
+  bool canImportModule(ImportPath::Module ModulePath,
                        llvm::VersionTuple version = llvm::VersionTuple(),
                        bool underlyingVersion = false) const;
 

include/swift/AST/Import.h

@@ -225,6 +225,22 @@ namespace detail {
       }
     }
 
+    /// Parses \p text into elements separated by \p separator, with identifiers
+    /// from \p ctx starting at \p loc.
+    ///
+    /// \warning This is not very robust; for instance, it doesn't check the
+    /// validity of the identifiers.
+    ImportPathBuilder(ASTContext &ctx, StringRef text, char separator,
+                      SourceLoc loc)
+        : scratch() {
+      while (!text.empty()) {
+        StringRef next;
+        std::tie(next, text) = text.split(separator);
+        push_back({ImportPathBuilder_getIdentifierImpl(ctx, next), loc});
+        loc = loc.getAdvancedLocOrInvalid(next.size() + 1);
+      }
+    }
+
     void push_back(const ImportPathElement &elem) { scratch.push_back(elem); }
     void push_back(Identifier name, SourceLoc loc = SourceLoc()) {
       scratch.push_back({ name, loc });

include/swift/AST/ModuleLoader.h

@@ -200,7 +200,7 @@ class ModuleLoader {
   ///
   /// Note that even if this check succeeds, errors may still occur if the
   /// module is loaded in full.
-  virtual bool canImportModule(ImportPath::Element named,
+  virtual bool canImportModule(ImportPath::Module named,
                                llvm::VersionTuple version,
                                bool underlyingVersion) = 0;
 

include/swift/ClangImporter/ClangImporter.h

@@ -192,7 +192,8 @@ class ClangImporter final : public ClangModuleLoader {
   ///
   /// Note that even if this check succeeds, errors may still occur if the
   /// module is loaded in full.
-  virtual bool canImportModule(ImportPath::Element named, llvm::VersionTuple version,
+  virtual bool canImportModule(ImportPath::Module named,
+                               llvm::VersionTuple version,
                                bool underlyingVersion) override;
 
   /// Import a module with the given module path.

include/swift/Frontend/ModuleInterfaceLoader.h

@@ -151,7 +151,7 @@ class ExplicitSwiftModuleLoader: public SerializedModuleLoaderBase {
                   std::unique_ptr<llvm::MemoryBuffer> *ModuleSourceInfoBuffer,
                   bool skipBuildingInterface, bool IsFramework) override;
 
-  bool canImportModule(ImportPath::Element mID, llvm::VersionTuple version,
+  bool canImportModule(ImportPath::Module named, llvm::VersionTuple version,
                        bool underlyingVersion) override;
 
   bool isCached(StringRef DepPath) override { return false; };

include/swift/Sema/SourceLoader.h

@@ -57,7 +57,7 @@ class SourceLoader : public ModuleLoader {
   ///
   /// Note that even if this check succeeds, errors may still occur if the
   /// module is loaded in full.
-  virtual bool canImportModule(ImportPath::Element named,
+  virtual bool canImportModule(ImportPath::Module named,
                                llvm::VersionTuple version,
                                bool underlyingVersion) override;
 

include/swift/Serialization/SerializedModuleLoader.h

@@ -168,7 +168,7 @@ class SerializedModuleLoaderBase : public ModuleLoader {
   ///
   /// Note that even if this check succeeds, errors may still occur if the
   /// module is loaded in full.
-  virtual bool canImportModule(ImportPath::Element named,
+  virtual bool canImportModule(ImportPath::Module named,
                                llvm::VersionTuple version,
                                bool underlyingVersion) override;
 
@@ -287,7 +287,7 @@ class MemoryBufferSerializedModuleLoader : public SerializedModuleLoaderBase {
 public:
   virtual ~MemoryBufferSerializedModuleLoader();
 
-  bool canImportModule(ImportPath::Element named, llvm::VersionTuple version,
+  bool canImportModule(ImportPath::Module named, llvm::VersionTuple version,
                        bool underlyingVersion) override;
   ModuleDecl *
   loadModule(SourceLoc importLoc,

lib/AST/ASTContext.cpp

@@ -2168,17 +2168,21 @@ bool ASTContext::shouldPerformTypoCorrection() {
   return NumTypoCorrections <= LangOpts.TypoCorrectionLimit;
 }
 
-bool ASTContext::canImportModuleImpl(ImportPath::Element ModuleName,
+bool ASTContext::canImportModuleImpl(ImportPath::Module ModuleName,
                                      llvm::VersionTuple version,
                                      bool underlyingVersion,
                                      bool updateFailingList) const {
+  SmallString<64> FullModuleName;
+  ModuleName.getString(FullModuleName);
+  auto ModuleNameStr = FullModuleName.str();
+
   // If we've failed loading this module before, don't look for it again.
-  if (FailedModuleImportNames.count(ModuleName.Item))
+  if (FailedModuleImportNames.count(ModuleNameStr))
     return false;
   // If no specific version, the module is importable if it has already been imported.
   if (version.empty()) {
     // If this module has already been successfully imported, it is importable.
-    if (getLoadedModule(ImportPath::Module::Builder(ModuleName).get()) != nullptr)
+    if (getLoadedModule(ModuleName) != nullptr)
       return true;
   }
   // Otherwise, ask the module loaders.
@@ -2188,18 +2192,18 @@ bool ASTContext::canImportModuleImpl(ImportPath::Element ModuleName,
     }
   }
   if (updateFailingList && version.empty()) {
-    FailedModuleImportNames.insert(ModuleName.Item);
+    FailedModuleImportNames.insert(ModuleNameStr);
   }
   return false;
 }
 
-bool ASTContext::canImportModule(ImportPath::Element ModuleName,
+bool ASTContext::canImportModule(ImportPath::Module ModuleName,
                                  llvm::VersionTuple version,
                                  bool underlyingVersion) {
   return canImportModuleImpl(ModuleName, version, underlyingVersion, true);
 }
 
-bool ASTContext::canImportModule(ImportPath::Element ModuleName,
+bool ASTContext::canImportModule(ImportPath::Module ModuleName,
                                  llvm::VersionTuple version,
                                  bool underlyingVersion) const {
   return canImportModuleImpl(ModuleName, version, underlyingVersion, false);

lib/ClangImporter/ClangImporter.cpp

@@ -1808,14 +1808,14 @@ static std::string getScalaNodeText(llvm::yaml::Node *N) {
   return cast<llvm::yaml::ScalarNode>(N)->getValue(Buffer).str();
 }
 
-bool ClangImporter::canImportModule(ImportPath::Element moduleID,
+bool ClangImporter::canImportModule(ImportPath::Module modulePath,
                                     llvm::VersionTuple version,
                                     bool underlyingVersion) {
   // Look up the top-level module to see if it exists.
-  // FIXME: This only works with top-level modules.
   auto &clangHeaderSearch = Impl.getClangPreprocessor().getHeaderSearchInfo();
+  auto topModule = modulePath.front();
   clang::Module *clangModule = clangHeaderSearch.lookupModule(
-      moduleID.Item.str(), /*ImportLoc=*/clang::SourceLocation(),
+      topModule.Item.str(), /*ImportLoc=*/clang::SourceLocation(),
       /*AllowSearch=*/true, /*AllowExtraModuleMapSearch=*/true);
   if (!clangModule) {
     return false;
@@ -1825,10 +1825,34 @@ bool ClangImporter::canImportModule(ImportPath::Element moduleID,
   clang::Module::UnresolvedHeaderDirective mh;
   clang::Module *m;
   auto &ctx = Impl.getClangASTContext();
-  auto available = clangModule->isAvailable(ctx.getLangOpts(), getTargetInfo(),
-                                            r, mh, m);
+  auto &lo = ctx.getLangOpts();
+  auto &ti = getTargetInfo();
+
+  auto available = clangModule->isAvailable(lo, ti, r, mh, m);
   if (!available)
     return false;
+
+  if (modulePath.hasSubmodule()) {
+    for (auto &component : modulePath.getSubmodulePath()) {
+      clangModule = clangModule->findSubmodule(component.Item.str());
+
+      // Special case: a submodule named "Foo.Private" can be moved to a
+      // top-level module named "Foo_Private". Clang has special support for
+      // this.
+      if (!clangModule && component.Item.str() == "Private" &&
+          (&component) == (&modulePath.getRaw()[1])) {
+        clangModule = clangHeaderSearch.lookupModule(
+            (topModule.Item.str() + "_Private").str(),
+            /*ImportLoc=*/clang::SourceLocation(),
+            /*AllowSearch=*/true,
+            /*AllowExtraModuleMapSearch=*/true);
+      }
+      if (!clangModule || !clangModule->isAvailable(lo, ti, r, mh, m)) {
+        return false;
+      }
+    }
+  }
+
   if (version.empty())
     return true;
   assert(available);
@@ -1838,11 +1862,11 @@ bool ClangImporter::canImportModule(ImportPath::Element moduleID,
     .getFilename(clangModule->DefinitionLoc);
   // Look for the .tbd file inside .framework dir to get the project version
   // number.
-  std::string fwName = (llvm::Twine(moduleID.Item.str()) + ".framework").str();
+  std::string fwName = (llvm::Twine(topModule.Item.str()) + ".framework").str();
   auto pos = path.find(fwName);
   while (pos != StringRef::npos) {
     llvm::SmallString<256> buffer(path.substr(0, pos + fwName.size()));
-    llvm::sys::path::append(buffer, llvm::Twine(moduleID.Item.str()) + ".tbd");
+    llvm::sys::path::append(buffer, llvm::Twine(topModule.Item.str()) + ".tbd");
     auto tbdPath = buffer.str();
     llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> tbdBufOrErr =
       llvm::MemoryBuffer::getFile(tbdPath);
@@ -1877,8 +1901,8 @@ bool ClangImporter::canImportModule(ImportPath::Element moduleID,
   }
   // Diagnose unable to checking the current version.
   if (currentVersion.empty()) {
-    Impl.diagnose(moduleID.Loc, diag::cannot_find_project_version, "Clang",
-                  moduleID.Item.str());
+    Impl.diagnose(topModule.Loc, diag::cannot_find_project_version, "Clang",
+                  topModule.Item.str());
     return true;
   }
   assert(!currentVersion.empty());

lib/ClangImporter/ImportType.cpp

@@ -2726,7 +2726,9 @@ ModuleDecl *ClangImporter::Implementation::tryLoadFoundationModule() {
 }
 
 bool ClangImporter::Implementation::canImportFoundationModule() {
-  return SwiftContext.canImportModule({SwiftContext.Id_Foundation, SourceLoc()});
+  ImportPath::Module::Builder builder(SwiftContext.Id_Foundation);
+  auto modulePath = builder.get();
+  return SwiftContext.canImportModule(modulePath);
 }
 
 Type ClangImporter::Implementation::getNamedSwiftType(ModuleDecl *module,