Merge pull request #70974 from cachemeifyoucan/eng/PR-121082031

[ExplicitModule] Fix `canImport` lookup for swift explicit module build

Author: cachemeifyoucan

include/swift/AST/ASTContext.h

@@ -395,6 +395,10 @@ class ASTContext final {
   /// Cache of module names that fail the 'canImport' test in this context.
   mutable llvm::StringSet<> FailedModuleImportNames;
 
+  /// 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;
+
   /// 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.
@@ -1202,9 +1206,19 @@ class ASTContext final {
   bool canImportModule(ImportPath::Module ModulePath,
                        llvm::VersionTuple version = llvm::VersionTuple(),
                        bool underlyingVersion = false);
-  bool canImportModule(ImportPath::Module ModulePath,
-                       llvm::VersionTuple version = llvm::VersionTuple(),
-                       bool underlyingVersion = false) const;
+
+  /// Check whether the module with a given name can be imported without
+  /// importing it. This is a const method that won't remember the outcome so
+  /// repeated check of the same module will induce full cost and won't count
+  /// as the dependency for current module.
+  bool testImportModule(ImportPath::Module ModulePath,
+                        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;
+  }
 
   /// \returns a module with a given name that was already loaded.  If the
   /// module was not loaded, returns nullptr.

lib/AST/ASTContext.cpp

@@ -2318,6 +2318,10 @@ bool ASTContext::canImportModuleImpl(ImportPath::Module ModuleName,
     return false;
 
   if (version.empty()) {
+    // If this module has already been checked successfully, it is importable.
+    if (SucceededModuleImportNames.count(ModuleNameStr))
+      return true;
+
     // If this module has already been successfully imported, it is importable.
     if (getLoadedModule(ModuleName) != nullptr)
       return true;
@@ -2375,12 +2379,19 @@ bool ASTContext::canImportModuleImpl(ImportPath::Module ModuleName,
 bool ASTContext::canImportModule(ImportPath::Module ModuleName,
                                  llvm::VersionTuple version,
                                  bool underlyingVersion) {
-  return canImportModuleImpl(ModuleName, version, underlyingVersion, true);
+  if (!canImportModuleImpl(ModuleName, version, underlyingVersion, true))
+    return false;
+
+  // If inserted successfully, add that to success list as dependency.
+  SmallString<64> FullModuleName;
+  ModuleName.getString(FullModuleName);
+  SucceededModuleImportNames.insert(FullModuleName.str());
+  return true;
 }
 
-bool ASTContext::canImportModule(ImportPath::Module ModuleName,
-                                 llvm::VersionTuple version,
-                                 bool underlyingVersion) const {
+bool ASTContext::testImportModule(ImportPath::Module ModuleName,
+                                  llvm::VersionTuple version,
+                                  bool underlyingVersion) const {
   return canImportModuleImpl(ModuleName, version, underlyingVersion, false);
 }
 

lib/ClangImporter/ClangImporter.cpp

@@ -2016,11 +2016,8 @@ bool ClangImporter::canImportModule(ImportPath::Module modulePath,
                                     ModuleVersionInfo *versionInfo,
                                     bool isTestableDependencyLookup) {
   // Look up the top-level module to see if it exists.
-  auto &clangHeaderSearch = Impl.getClangPreprocessor().getHeaderSearchInfo();
   auto topModule = modulePath.front();
-  clang::Module *clangModule = clangHeaderSearch.lookupModule(
-      topModule.Item.str(), /*ImportLoc=*/clang::SourceLocation(),
-      /*AllowSearch=*/true, /*AllowExtraModuleMapSearch=*/true);
+  clang::Module *clangModule = Impl.lookupModule(topModule.Item.str());
   if (!clangModule) {
     return false;
   }
@@ -2045,11 +2042,8 @@ bool ClangImporter::canImportModule(ImportPath::Module modulePath,
       // 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);
+        clangModule =
+            Impl.lookupModule((topModule.Item.str() + "_Private").str());
       }
       if (!clangModule || !clangModule->isAvailable(lo, ti, r, mh, m)) {
         return false;
@@ -2073,6 +2067,39 @@ bool ClangImporter::canImportModule(ImportPath::Module modulePath,
   return true;
 }
 
+clang::Module *
+ClangImporter::Implementation::lookupModule(StringRef moduleName) {
+  auto &clangHeaderSearch = getClangPreprocessor().getHeaderSearchInfo();
+  if (getClangASTContext().getLangOpts().ImplicitModules)
+    return clangHeaderSearch.lookupModule(
+        moduleName, /*ImportLoc=*/clang::SourceLocation(),
+        /*AllowSearch=*/true, /*AllowExtraModuleMapSearch=*/true);
+
+  // Explicit module. Try load from modulemap.
+  auto &PP = Instance->getPreprocessor();
+  auto &MM = PP.getHeaderSearchInfo().getModuleMap();
+  auto loadFromMM = [&]() -> clang::Module * {
+    auto *II = PP.getIdentifierInfo(moduleName);
+    if (auto clangModule = MM.getCachedModuleLoad(*II))
+      return *clangModule;
+    return nullptr;
+  };
+  // Check if it is already loaded.
+  if (auto *clangModule = loadFromMM())
+    return clangModule;
+
+  // If not, try load it.
+  auto &PrebuiltModules = Instance->getHeaderSearchOpts().PrebuiltModuleFiles;
+  auto moduleFile = PrebuiltModules.find(moduleName);
+  if (moduleFile == PrebuiltModules.end())
+    return nullptr;
+
+  if (!Instance->loadModuleFile(moduleFile->second))
+    return nullptr; // error loading, return not found.
+  // Lookup again.
+  return loadFromMM();
+}
+
 ModuleDecl *ClangImporter::Implementation::loadModuleClang(
     SourceLoc importLoc, ImportPath::Module path) {
   auto &clangHeaderSearch = getClangPreprocessor().getHeaderSearchInfo();

lib/ClangImporter/ImporterImpl.h

@@ -860,6 +860,9 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   ModuleDecl *loadModuleDWARF(SourceLoc importLoc,
                               ImportPath::Module path);
 
+  /// Lookup a clang module.
+  clang::Module *lookupModule(StringRef moduleName);
+
 public:
   /// Load a module using either method.
   ModuleDecl *loadModule(SourceLoc importLoc,

lib/DependencyScan/ModuleDependencyScanner.cpp

@@ -428,10 +428,18 @@ ModuleDependencyScanner::getMainModuleDependencyInfo(
     for (auto fileUnit : mainModule->getFiles()) {
       auto sf = dyn_cast<SourceFile>(fileUnit);
       if (!sf)
-	continue;
+        continue;
 
       mainDependencies.addModuleImport(*sf, alreadyAddedModules);
     }
+
+    // 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
+    // SourceFiles.
+    for (auto &Module :
+         mainModule->getASTContext().getSuccessfulCanImportCheckNames())
+      mainDependencies.addModuleImport(Module.first(), &alreadyAddedModules);
   }    
 
   return mainDependencies;

lib/Frontend/Frontend.cpp

@@ -1119,14 +1119,14 @@ bool CompilerInstance::canImportSwiftConcurrency() const {
   ImportPath::Module::Builder builder(
       getASTContext().getIdentifier(SWIFT_CONCURRENCY_NAME));
   auto modulePath = builder.get();
-  return getASTContext().canImportModule(modulePath);
+  return getASTContext().testImportModule(modulePath);
 }
 
 bool CompilerInstance::canImportSwiftConcurrencyShims() const {
   ImportPath::Module::Builder builder(
       getASTContext().getIdentifier(SWIFT_CONCURRENCY_SHIMS_NAME));
   auto modulePath = builder.get();
-  return getASTContext().canImportModule(modulePath);
+  return getASTContext().testImportModule(modulePath);
 }
 
 void CompilerInstance::verifyImplicitStringProcessingImport() {
@@ -1141,7 +1141,7 @@ bool CompilerInstance::canImportSwiftStringProcessing() const {
   ImportPath::Module::Builder builder(
       getASTContext().getIdentifier(SWIFT_STRING_PROCESSING_NAME));
   auto modulePath = builder.get();
-  return getASTContext().canImportModule(modulePath);
+  return getASTContext().testImportModule(modulePath);
 }
 
 void CompilerInstance::verifyImplicitBacktracingImport() {
@@ -1156,15 +1156,16 @@ bool CompilerInstance::canImportSwiftBacktracing() const {
   ImportPath::Module::Builder builder(
       getASTContext().getIdentifier(SWIFT_BACKTRACING_NAME));
   auto modulePath = builder.get();
-  return getASTContext().canImportModule(modulePath);
+  return getASTContext().testImportModule(modulePath);
 }
 
 bool CompilerInstance::canImportCxxShim() const {
   ImportPath::Module::Builder builder(
       getASTContext().getIdentifier(CXX_SHIM_NAME));
   auto modulePath = builder.get();
-  return getASTContext().canImportModule(modulePath) &&
-         !Invocation.getFrontendOptions().InputsAndOutputs.hasModuleInterfaceOutputPath();
+  return getASTContext().testImportModule(modulePath) &&
+         !Invocation.getFrontendOptions()
+              .InputsAndOutputs.hasModuleInterfaceOutputPath();
 }
 
 bool CompilerInstance::supportCaching() const {

test/CAS/can-import.swift

@@ -0,0 +1,105 @@
+// rdar://119964830 Temporarily disabling in Linux
+// UNSUPPORTED: OS=linux-gnu
+
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// RUN: %target-swift-frontend -scan-dependencies -module-name Test -module-cache-path %t/clang-module-cache -O \
+// RUN:   -disable-implicit-string-processing-module-import -disable-implicit-concurrency-module-import \
+// RUN:   %t/main.swift -o %t/deps.json -swift-version 5 -cache-compile-job -cas-path %t/cas -I %t/include
+
+// RUN: %S/Inputs/BuildCommandExtractor.py %t/deps.json clang:A > %t/A.cmd
+// RUN: %swift_frontend_plain @%t/A.cmd
+
+// RUN: %S/Inputs/BuildCommandExtractor.py %t/deps.json clang:B > %t/B.cmd
+// RUN: %swift_frontend_plain @%t/B.cmd
+
+// RUN: %S/Inputs/BuildCommandExtractor.py %t/deps.json clang:SwiftShims > %t/SwiftShims.cmd
+// RUN: %swift_frontend_plain @%t/SwiftShims.cmd
+
+// RUN: %S/Inputs/BuildCommandExtractor.py %t/deps.json Swift > %t/Swift.cmd
+// RUN: %swift_frontend_plain @%t/Swift.cmd
+
+// RUN: %S/Inputs/SwiftDepsExtractor.py %t/deps.json Swift moduleCacheKey | tr -d '\n' > %t/Swift.key
+// RUN: %S/Inputs/SwiftDepsExtractor.py %t/deps.json clang:SwiftShims moduleCacheKey | tr -d '\n' > %t/Shims.key
+// RUN: %S/Inputs/SwiftDepsExtractor.py %t/deps.json clang:A moduleCacheKey | tr -d '\n' > %t/A.key
+// RUN: %S/Inputs/SwiftDepsExtractor.py %t/deps.json clang:B moduleCacheKey | tr -d '\n' > %t/B.key
+// RUN: %S/Inputs/BuildCommandExtractor.py %t/deps.json MyApp > %t/MyApp.cmd
+
+// RUN: echo "[{" > %/t/map.json
+// RUN: echo "\"moduleName\": \"Swift\"," >> %/t/map.json
+// RUN: echo "\"modulePath\": \"Swift.swiftmodule\"," >> %/t/map.json
+// RUN: echo -n "\"moduleCacheKey\": " >> %/t/map.json
+// RUN: cat %t/Swift.key >> %/t/map.json
+// RUN: echo "," >> %/t/map.json
+// RUN: echo "\"isFramework\": false" >> %/t/map.json
+// RUN: echo "}," >> %/t/map.json
+// RUN: echo "{" >> %/t/map.json
+// RUN: echo "\"moduleName\": \"SwiftShims\"," >> %/t/map.json
+// RUN: echo "\"clangModulePath\": \"SwiftShims.pcm\"," >> %/t/map.json
+// RUN: echo -n "\"clangModuleCacheKey\": " >> %/t/map.json
+// RUN: cat %t/Shims.key >> %/t/map.json
+// RUN: echo "," >> %/t/map.json
+// RUN: echo "\"isFramework\": false" >> %/t/map.json
+// RUN: echo "}," >> %/t/map.json
+// RUN: echo "{" >> %/t/map.json
+// RUN: echo "\"moduleName\": \"A\"," >> %/t/map.json
+// RUN: echo "\"clangModulePath\": \"A.pcm\"," >> %/t/map.json
+// RUN: echo -n "\"clangModuleCacheKey\": " >> %/t/map.json
+// RUN: cat %t/A.key >> %/t/map.json
+// RUN: echo "," >> %/t/map.json
+// RUN: echo "\"isFramework\": false" >> %/t/map.json
+// RUN: echo "}," >> %/t/map.json
+// RUN: echo "{" >> %/t/map.json
+// RUN: echo "\"moduleName\": \"B\"," >> %/t/map.json
+// RUN: echo "\"clangModulePath\": \"B.pcm\"," >> %/t/map.json
+// RUN: echo -n "\"clangModuleCacheKey\": " >> %/t/map.json
+// RUN: cat %t/B.key >> %/t/map.json
+// RUN: echo "," >> %/t/map.json
+// RUN: echo "\"isFramework\": false" >> %/t/map.json
+// RUN: echo "}]" >> %/t/map.json
+// RUN: llvm-cas --cas %t/cas --make-blob --data %t/map.json > %t/map.casid
+
+// RUN: %target-swift-frontend \
+// RUN:   -typecheck -cache-compile-job -cas-path %t/cas \
+// RUN:   -swift-version 5 -disable-implicit-swift-modules \
+// RUN:   -disable-implicit-string-processing-module-import -disable-implicit-concurrency-module-import \
+// RUN:   -module-name MyApp -explicit-swift-module-map-file @%t/map.casid \
+// RUN:   %t/main.swift @%t/MyApp.cmd
+
+//--- main.swift
+#if canImport(A.Sub)
+import A.Sub
+#endif
+
+#if canImport(A.Missing)
+import A.Missing
+#endif
+
+#if canImport(B) // never actually import B
+func b() {}
+#endif
+
+func useA() {
+  a()
+  b()
+}
+
+//--- include/module.modulemap
+module A {
+  module Sub {
+    header "sub.h"
+    export *
+  }
+}
+
+module B {
+  header "B.h"
+  export *
+}
+
+//--- include/sub.h
+void a(void);
+
+//--- include/B.h
+void notused(void);

test/ScanDependencies/module_deps_can_import.swift

@@ -0,0 +1,34 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -scan-dependencies -module-cache-path %t/clang-module-cache %s -module-name Test -o %t/deps.json -I %S/Inputs/CHeaders -I %S/Inputs/Swift -swift-version 4
+
+// RUN: %{python} %S/../CAS/Inputs/SwiftDepsExtractor.py %t/deps.json Test directDependencies | %FileCheck %s
+
+// CHECK-DAG: "clang": "C"
+// CHECK-DAG: "swift": "A"
+// CHECK-DAG: "swift": "F"
+// CHECK-NOT: "swift": "G"
+// CHECK-NOT: "swift": "B"
+// CHECK-NOT: "Missing"
+
+#if canImport(Missing)
+import G
+#endif
+
+#if canImport(C)
+import A
+#else
+import B
+#endif
+
+#if !canImport(F)
+import Missing
+#endif
+
+// B is not dependency
+#if false && canImport(B)
+import Missing
+#endif
+
+// B is short circuited
+#if canImport(C) || canImport(B)
+#endif