[ScanDependency][canImport] Improve canImport handling in explicit build

include/swift/AST/ASTContext.h

@@ -47,6 +47,7 @@
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/DataTypes.h"
+#include "llvm/Support/VersionTuple.h"
 #include "llvm/Support/VirtualOutputBackend.h"
 #include <functional>
 #include <memory>
@@ -409,9 +410,17 @@ class ASTContext final {
   /// Cache of module names that fail the 'canImport' test in this context.
   mutable llvm::StringSet<> FailedModuleImportNames;
 
+  /// Versions of the modules found during versioned canImport checks. 
+  struct ImportedModuleVersionInfo {
+    llvm::VersionTuple Version;
+    llvm::VersionTuple UnderlyingVersion;
+  };
+
   /// Cache of module names that passed the 'canImport' test. This cannot be
-  /// mutable since it needs to be queried for dependency discovery.
-  llvm::StringSet<> SucceededModuleImportNames;
+  /// mutable since it needs to be queried for dependency discovery. Keep sorted
+  /// so caller of `forEachCanImportVersionCheck` can expect deterministic
+  /// ordering.
+  std::map<std::string, ImportedModuleVersionInfo> CanImportModuleVersions;
 
   /// 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
@@ -1102,7 +1111,12 @@ class ASTContext final {
   /// module is loaded in full.
   bool canImportModuleImpl(ImportPath::Module ModulePath,
                            llvm::VersionTuple version, bool underlyingVersion,
-                           bool updateFailingList) const;
+                           bool updateFailingList,
+                           llvm::VersionTuple &foundVersion) const;
+
+  /// Add successful canImport modules.
+  void addSucceededCanImportModule(StringRef moduleName, bool underlyingVersion,
+                                   const llvm::VersionTuple &versionInfo);
 
 public:
   namelookup::ImportCache &getImportCache() const;
@@ -1144,10 +1158,10 @@ class ASTContext final {
                         llvm::VersionTuple version = llvm::VersionTuple(),
                         bool underlyingVersion = false) const;
 
-  /// \returns a set of names from all successfully canImport module checks.
-  const llvm::StringSet<> &getSuccessfulCanImportCheckNames() const {
-    return SucceededModuleImportNames;
-  }
+  /// Callback on each successful imported.
+  void forEachCanImportVersionCheck(
+      std::function<void(StringRef, const llvm::VersionTuple &,
+                         const llvm::VersionTuple &)>) const;
 
   /// \returns a module with a given name that was already loaded.  If the
   /// module was not loaded, returns nullptr.

include/swift/AST/DiagnosticsFrontend.def

@@ -481,6 +481,8 @@ ERROR(unknown_forced_module_loading_mode,none,
       (StringRef))
 ERROR(error_creating_remark_serializer,none,
       "error while creating remark serializer: '%0'", (StringRef))
+ERROR(invalid_can_import_module_version,none,
+      "invalid version passed to -module-can-import-version: '%0'", (StringRef))
 
 REMARK(interface_file_lock_failure,none,
       "building module interface without lock file", ())

include/swift/AST/SearchPathOptions.h

@@ -20,6 +20,7 @@
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/VersionTuple.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include <optional>
 
@@ -509,6 +510,14 @@ class SearchPathOptions {
   using CrossImportMap = llvm::StringMap<std::vector<std::string>>;
   CrossImportMap CrossImportInfo;
 
+  /// CanImport information passed from scanning.
+  struct CanImportInfo {
+    std::string ModuleName;
+    llvm::VersionTuple Version;
+    llvm::VersionTuple UnderlyingVersion;
+  };
+  std::vector<CanImportInfo> CanImportModuleInfo;
+
   /// Whether to search for cross import overlay on file system.
   bool DisableCrossImportOverlaySearch = false;
 

include/swift/Option/FrontendOptions.td

@@ -224,6 +224,14 @@ def swift_module_cross_import: MultiArg<["-"], "swift-module-cross-import", 2>,
   MetaVarName<"<moduleName> <crossImport.swiftoverlay>">,
   HelpText<"Specify the cross import module">;
 
+def module_can_import: Separate<["-"], "module-can-import">,
+  MetaVarName<"<moduleName>">,
+  HelpText<"Specify canImport module name">;
+
+def module_can_import_version: MultiArg<["-"], "module-can-import-version", 3>,
+  MetaVarName<"<moduleName> <version> <underlyingVersion>">,
+  HelpText<"Specify canImport module and versions">;
+
 def disable_cross_import_overlay_search: Flag<["-"], "disable-cross-import-overlay-search">,
   HelpText<"Disable searching for cross import overlay file">;
 

lib/AST/ASTContext.cpp

@@ -74,6 +74,7 @@
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/VersionTuple.h"
 #include "llvm/Support/VirtualOutputBackend.h"
 #include "llvm/Support/VirtualOutputBackends.h"
 #include <algorithm>
@@ -786,6 +787,12 @@ ASTContext::ASTContext(
   registerAccessRequestFunctions(evaluator);
   registerNameLookupRequestFunctions(evaluator);
 
+  // Register canImport module info.
+  for (auto &info: SearchPathOpts.CanImportModuleInfo) {
+    addSucceededCanImportModule(info.ModuleName, false, info.Version);
+    addSucceededCanImportModule(info.ModuleName, true, info.UnderlyingVersion);
+  }
+
   // Provide a default OnDiskOutputBackend if user didn't supply one.
   if (!OutputBackend)
     OutputBackend = llvm::makeIntrusiveRefCnt<llvm::vfs::OnDiskOutputBackend>();
@@ -2385,23 +2392,56 @@ getModuleVersionKindString(const ModuleLoader::ModuleVersionInfo &info) {
   }
 }
 
-bool ASTContext::canImportModuleImpl(ImportPath::Module ModuleName,
-                                     llvm::VersionTuple version,
-                                     bool underlyingVersion,
-                                     bool updateFailingList) const {
+void ASTContext::addSucceededCanImportModule(
+    StringRef moduleName, bool underlyingVersion,
+    const llvm::VersionTuple &versionInfo) {
+  auto &entry = CanImportModuleVersions[moduleName.str()];
+  if (!versionInfo.empty()) {
+    if (underlyingVersion)
+      entry.UnderlyingVersion = versionInfo;
+    else
+      entry.Version = versionInfo;
+  }
+}
+
+bool ASTContext::canImportModuleImpl(
+    ImportPath::Module ModuleName, llvm::VersionTuple version,
+    bool underlyingVersion, bool updateFailingList,
+    llvm::VersionTuple &foundVersion) const {
   SmallString<64> FullModuleName;
   ModuleName.getString(FullModuleName);
-  auto ModuleNameStr = FullModuleName.str();
+  auto ModuleNameStr = FullModuleName.str().str();
 
   // If we've failed loading this module before, don't look for it again.
   if (FailedModuleImportNames.count(ModuleNameStr))
     return false;
 
-  if (version.empty()) {
-    // If this module has already been checked successfully, it is importable.
-    if (SucceededModuleImportNames.count(ModuleNameStr))
+  // If this module has already been checked or there is information for the
+  // module from commandline, use that information instead of loading the
+  // module.
+  auto Found = CanImportModuleVersions.find(ModuleNameStr);
+  if (Found != CanImportModuleVersions.end()) {
+    if (version.empty())
       return true;
 
+    if (underlyingVersion) {
+      if (!Found->second.UnderlyingVersion.empty())
+        return version <= Found->second.UnderlyingVersion;
+    } else {
+      if (!Found->second.Version.empty())
+        return version <= Found->second.Version;
+    }
+
+    // If the canImport information is coming from the command-line, then no
+    // need to continue the search, return false. For checking modules that are
+    // not passed from command-line, allow fallback to the module loading since
+    // this is not in a canImport request context that has already been resolved
+    // by scanner.
+    if (!SearchPathOpts.CanImportModuleInfo.empty())
+      return false;
+  }
+
+  if (version.empty()) {
     // If this module has already been successfully imported, it is importable.
     if (getLoadedModule(ModuleName) != nullptr)
       return true;
@@ -2453,28 +2493,38 @@ bool ASTContext::canImportModuleImpl(ImportPath::Module ModuleName,
     return true;
   }
 
+  foundVersion = bestVersionInfo.getVersion();
   return version <= bestVersionInfo.getVersion();
 }
 
-bool ASTContext::canImportModule(ImportPath::Module ModuleName,
+void ASTContext::forEachCanImportVersionCheck(
+    std::function<void(StringRef, const llvm::VersionTuple &,
+                       const llvm::VersionTuple &)>
+        Callback) const {
+  for (auto &entry : CanImportModuleVersions)
+    Callback(entry.first, entry.second.Version, entry.second.UnderlyingVersion);
+}
+
+bool ASTContext::canImportModule(ImportPath::Module moduleName,
                                  llvm::VersionTuple version,
                                  bool underlyingVersion) {
-  if (!canImportModuleImpl(ModuleName, version, underlyingVersion, true))
+  llvm::VersionTuple versionInfo;
+  if (!canImportModuleImpl(moduleName, version, underlyingVersion, true,
+                           versionInfo))
     return false;
 
-  // If checked successfully, add the top level name to success list as
-  // dependency to handle clang submodule correctly. Swift does not have
-  // submodule so the name should be the same.
-  SmallString<64> TopModuleName;
-  ModuleName.getTopLevelPath().getString(TopModuleName);
-  SucceededModuleImportNames.insert(TopModuleName.str());
+  SmallString<64> fullModuleName;
+  moduleName.getString(fullModuleName);
+  addSucceededCanImportModule(fullModuleName, underlyingVersion, versionInfo);
   return true;
 }
 
 bool ASTContext::testImportModule(ImportPath::Module ModuleName,
                                   llvm::VersionTuple version,
                                   bool underlyingVersion) const {
-  return canImportModuleImpl(ModuleName, version, underlyingVersion, false);
+  llvm::VersionTuple versionInfo;
+  return canImportModuleImpl(ModuleName, version, underlyingVersion, false,
+                             versionInfo);
 }
 
 ModuleDecl *

lib/ClangImporter/ClangImporter.cpp

@@ -2120,7 +2120,8 @@ bool ClangImporter::isModuleImported(const clang::Module *M) {
   return M->NameVisibility == clang::Module::NameVisibilityKind::AllVisible;
 }
 
-static llvm::VersionTuple getCurrentVersionFromTBD(StringRef path,
+static llvm::VersionTuple getCurrentVersionFromTBD(llvm::vfs::FileSystem &FS,
+                                                   StringRef path,
                                                    StringRef moduleName) {
   std::string fwName = (moduleName + ".framework").str();
   auto pos = path.find(fwName);
@@ -2130,7 +2131,7 @@ static llvm::VersionTuple getCurrentVersionFromTBD(StringRef path,
   llvm::sys::path::append(buffer, moduleName + ".tbd");
   auto tbdPath = buffer.str();
   llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> tbdBufOrErr =
-      llvm::MemoryBuffer::getFile(tbdPath);
+      FS.getBufferForFile(tbdPath);
   // .tbd file doesn't exist, exit.
   if (!tbdBufOrErr)
     return {};
@@ -2193,8 +2194,8 @@ bool ClangImporter::canImportModule(ImportPath::Module modulePath,
 
   // Look for the .tbd file inside .framework dir to get the project version
   // number.
-  llvm::VersionTuple currentVersion =
-      getCurrentVersionFromTBD(path, topModule.Item.str());
+  llvm::VersionTuple currentVersion = getCurrentVersionFromTBD(
+      Impl.Instance->getVirtualFileSystem(), path, topModule.Item.str());
   versionInfo->setVersion(currentVersion,
                           ModuleVersionSourceKind::ClangModuleTBD);
   return true;

lib/DependencyScan/ModuleDependencyScanner.cpp

@@ -31,6 +31,7 @@
 #include "llvm/CAS/CachingOnDiskFileSystem.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/Threading.h"
+#include "llvm/Support/VersionTuple.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include <algorithm>
 
@@ -464,15 +465,26 @@ ModuleDependencyScanner::getMainModuleDependencyInfo(ModuleDecl *mainModule) {
                                         &ScanASTContext.SourceMgr);
     }
 
-    // Add all the successful canImport checks from the ASTContext as part of
-    // the dependency since only mainModule can have `canImport` check. This
-    // needs to happen after visiting all the top-level decls from all
+    // Pass all the successful canImport checks from the ASTContext as part of
+    // build command to main module to ensure frontend gets the same result.
+    // This needs to happen after visiting all the top-level decls from all
     // SourceFiles.
-    for (auto &Module :
-         mainModule->getASTContext().getSuccessfulCanImportCheckNames())
-      mainDependencies.addModuleImport(Module.first(),
-                                       &alreadyAddedModules);
-  }
+    auto buildArgs = mainDependencies.getCommandline();
+    mainModule->getASTContext().forEachCanImportVersionCheck(
+        [&](StringRef moduleName, const llvm::VersionTuple &Version,
+            const llvm::VersionTuple &UnderlyingVersion) {
+          if (Version.empty() && UnderlyingVersion.empty()) {
+            buildArgs.push_back("-module-can-import");
+            buildArgs.push_back(moduleName.str());
+          } else {
+            buildArgs.push_back("-module-can-import-version");
+            buildArgs.push_back(moduleName.str());
+            buildArgs.push_back(Version.getAsString());
+            buildArgs.push_back(UnderlyingVersion.getAsString());
+          }
+        });
+    mainDependencies.updateCommandLine(buildArgs);
+  }    
 
   return mainDependencies;
 }

lib/Frontend/CompilerInvocation.cpp

@@ -25,6 +25,7 @@
 #include "swift/Strings.h"
 #include "swift/SymbolGraphGen/SymbolGraphOptions.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/VersionTuple.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Option/Arg.h"
 #include "llvm/Option/ArgList.h"
@@ -2144,6 +2145,21 @@ static bool ParseSearchPathArgs(SearchPathOptions &Opts, ArgList &Args,
   for (auto *A : Args.filtered(OPT_swift_module_cross_import))
     Opts.CrossImportInfo[A->getValue(0)].push_back(A->getValue(1));
 
+  for (auto &Name : Args.getAllArgValues(OPT_module_can_import))
+    Opts.CanImportModuleInfo.push_back({Name, {}, {}});
+
+  for (auto *A: Args.filtered(OPT_module_can_import_version)) {
+    llvm::VersionTuple Version, UnderlyingVersion;
+    if (Version.tryParse(A->getValue(1)))
+      Diags.diagnose(SourceLoc(), diag::invalid_can_import_module_version,
+                     A->getValue(1));
+    if (UnderlyingVersion.tryParse(A->getValue(2)))
+      Diags.diagnose(SourceLoc(), diag::invalid_can_import_module_version,
+                     A->getValue(2));
+    Opts.CanImportModuleInfo.push_back(
+        {A->getValue(0), Version, UnderlyingVersion});
+  }
+
   Opts.DisableCrossImportOverlaySearch |=
       Args.hasArg(OPT_disable_cross_import_overlay_search);