[ExplicitModule] Fix canImport lookup for swift explicit module build

Previously, canImport lookup is not completely working with explicit
module due to two issues:
* For clang modules, canImport check still do a full modulemap lookup
  which is repeated work from scanner. For caching builds, this lookup
  cannot be performed because all modulemap and search path are dropped
  after scanning.
* For swift module, if the canImport module was never actually imported
  later, this canImport check will fail during the actual compilation,
  causing different dependencies in the actual compilation.

To fix the problem, first unified the lookup method for clang and swift
module, which will only lookup the module dependencies reported by
scanner to determine if `canImport` succeed or not. Secondly, add all
the successful `canImport` check modules into the dependency of the
current module so this information can be used during actual
compilation.

Note the behavior change here is that if a module is only checked in
`canImport` but never imported still needs to be built. Comparing to
implicit module build, this can bring in additional clang modules if
they are only check inside `canImport` but should not increase work for
swift modules (where binary module needs to be on disk anyway) or the
most common usecase for `canImport` which is to check the same module
before importing.

rdar://121082031

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

@@ -1107,14 +1107,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() {
@@ -1129,7 +1129,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() {
@@ -1144,15 +1144,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