DependencyScanner: honor additional compiler flags in interfaces files when collecting imports

Additional flags in interface files may change parsing behavior like #if
statements. We should use a fresh ASTContext with these additional
flags when parsing interface files to collect imports.

rdar://62612027

Author: nkcsgexi

include/swift/AST/ASTContext.h

@@ -121,6 +121,7 @@ namespace swift {
   class UnifiedStatsReporter;
   class IndexSubset;
   struct SILAutoDiffDerivativeFunctionKey;
+  struct SubASTContextDelegate;
 
   enum class KnownProtocolKind : uint8_t;
 
@@ -719,7 +720,8 @@ class ASTContext final {
   Optional<ModuleDependencies> getModuleDependencies(
       StringRef moduleName,
       bool isUnderlyingClangModule,
-      ModuleDependenciesCache &cache);
+      ModuleDependenciesCache &cache,
+      SubASTContextDelegate &delegate);
 
   /// Load extensions to the given nominal type from the external
   /// module loaders.

include/swift/AST/ModuleLoader.h

@@ -25,6 +25,7 @@
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TinyPtrVector.h"
+#include "swift/AST/ModuleDependencies.h"
 
 namespace llvm {
 class FileCollector;
@@ -84,6 +85,15 @@ class DependencyTracker {
   std::shared_ptr<clang::DependencyCollector> getClangCollector();
 };
 
+/// Abstract interface to run an action in a sub ASTContext.
+struct SubASTContextDelegate {
+  virtual bool runInSubContext(ASTContext &ctx, StringRef interfacePath,
+    llvm::function_ref<bool(ASTContext&)> action) {
+    llvm_unreachable("function should be overriden");
+  }
+  virtual ~SubASTContextDelegate() = default;
+};
+
 /// Abstract interface that loads named modules into the AST.
 class ModuleLoader {
   virtual void anchor();
@@ -186,7 +196,7 @@ class ModuleLoader {
   /// if no such module exists.
   virtual Optional<ModuleDependencies> getModuleDependencies(
       StringRef moduleName,
-      ModuleDependenciesCache &cache) = 0;
+      ModuleDependenciesCache &cache, SubASTContextDelegate &delegate) = 0;
 };
 
 } // namespace swift

include/swift/ClangImporter/ClangImporter.h

@@ -370,7 +370,8 @@ class ClangImporter final : public ClangModuleLoader {
   void verifyAllModules() override;
 
   Optional<ModuleDependencies> getModuleDependencies(
-      StringRef moduleName, ModuleDependenciesCache &cache) override;
+      StringRef moduleName, ModuleDependenciesCache &cache,
+      SubASTContextDelegate &delegate) override;
 
   /// Add dependency information for the bridging header.
   ///

include/swift/Frontend/ModuleInterfaceLoader.h

@@ -110,6 +110,7 @@
 #include "swift/Basic/LLVM.h"
 #include "swift/Frontend/ModuleInterfaceSupport.h"
 #include "swift/Serialization/SerializedModuleLoader.h"
+#include "llvm/Support/StringSaver.h"
 
 namespace clang {
 class CompilerInstance;
@@ -201,6 +202,14 @@ class ModuleInterfaceLoader : public SerializedModuleLoaderBase {
 std::string
 getModuleCachePathFromClang(const clang::CompilerInstance &Instance);
 
+bool extractSwiftInterfaceVersionAndArgs(SourceManager &SM,
+                                         DiagnosticEngine &Diags,
+                                         StringRef InterfacePath,
+                                         version::Version &Vers,
+                                         StringRef &CompilerVersion,
+                                         llvm::StringSaver &SubArgSaver,
+                                         SmallVectorImpl<const char *> &SubArgs,
+                                         SourceLoc diagnosticLoc = SourceLoc());
 }
 
 #endif

include/swift/Sema/SourceLoader.h

@@ -93,7 +93,8 @@ class SourceLoader : public ModuleLoader {
   }
 
   Optional<ModuleDependencies> getModuleDependencies(
-      StringRef moduleName, ModuleDependenciesCache &cache) override {
+      StringRef moduleName, ModuleDependenciesCache &cache,
+      SubASTContextDelegate &delegate) override {
     // FIXME: Implement?
     return None;
   }

include/swift/Serialization/SerializedModuleLoader.h

@@ -194,7 +194,8 @@ class SerializedModuleLoaderBase : public ModuleLoader {
   virtual void verifyAllModules() override;
 
   virtual Optional<ModuleDependencies> getModuleDependencies(
-      StringRef moduleName, ModuleDependenciesCache &cache) override;
+      StringRef moduleName, ModuleDependenciesCache &cache,
+      SubASTContextDelegate &delegate) override;
 };
 
 /// Imports serialized Swift modules into an ASTContext.

lib/AST/ASTContext.cpp

@@ -1462,13 +1462,14 @@ void ASTContext::addModuleLoader(std::unique_ptr<ModuleLoader> loader,
 
 Optional<ModuleDependencies> ASTContext::getModuleDependencies(
     StringRef moduleName, bool isUnderlyingClangModule,
-    ModuleDependenciesCache &cache) {
+    ModuleDependenciesCache &cache, SubASTContextDelegate &delegate) {
   for (auto &loader : getImpl().ModuleLoaders) {
     if (isUnderlyingClangModule &&
         loader.get() != getImpl().TheClangModuleLoader)
       continue;
 
-    if (auto dependencies = loader->getModuleDependencies(moduleName, cache))
+    if (auto dependencies = loader->getModuleDependencies(moduleName, cache,
+                                                          delegate))
       return dependencies;
   }
 

lib/ClangImporter/ClangModuleDependencyScanner.cpp

@@ -224,7 +224,8 @@ static void recordModuleDependencies(
 }
 
 Optional<ModuleDependencies> ClangImporter::getModuleDependencies(
-    StringRef moduleName, ModuleDependenciesCache &cache) {
+    StringRef moduleName, ModuleDependenciesCache &cache,
+    SubASTContextDelegate &delegate) {
   // Check whether there is already a cached result.
   if (auto found = cache.findDependencies(
           moduleName, ModuleDependenciesKind::Clang))

lib/Frontend/ModuleInterfaceBuilder.cpp

@@ -12,6 +12,7 @@
 
 #define DEBUG_TYPE "textual-module-interface"
 
+#include "swift/Frontend/ModuleInterfaceLoader.h"
 #include "ModuleInterfaceBuilder.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/DiagnosticsFrontend.h"
@@ -134,35 +135,44 @@ void ModuleInterfaceBuilder::configureSubInvocation(
   remarkOnRebuildFromInterface;
 }
 
-bool ModuleInterfaceBuilder::extractSwiftInterfaceVersionAndArgs(
-    swift::version::Version &Vers, StringRef &CompilerVersion,
-    llvm::StringSaver &SubArgSaver, SmallVectorImpl<const char *> &SubArgs) {
-  llvm::vfs::FileSystem &fs = *sourceMgr.getFileSystem();
-  auto FileOrError = swift::vfs::getFileOrSTDIN(fs, interfacePath);
+bool swift::extractSwiftInterfaceVersionAndArgs(
+    SourceManager &SM,
+    DiagnosticEngine &Diags,
+    StringRef InterfacePath,
+    version::Version &Vers,
+    StringRef &CompilerVersion,
+    llvm::StringSaver &SubArgSaver,
+    SmallVectorImpl<const char *> &SubArgs,
+    SourceLoc diagnosticLoc) {
+  llvm::vfs::FileSystem &fs = *SM.getFileSystem();
+  auto FileOrError = swift::vfs::getFileOrSTDIN(fs, InterfacePath);
   if (!FileOrError) {
     // Don't use this->diagnose() because it'll just try to re-open
     // interfacePath.
-    diags.diagnose(diagnosticLoc, diag::error_open_input_file,
-                   interfacePath, FileOrError.getError().message());
+    Diags.diagnose(diagnosticLoc, diag::error_open_input_file,
+                   InterfacePath, FileOrError.getError().message());
     return true;
   }
   auto SB = FileOrError.get()->getBuffer();
   auto VersRe = getSwiftInterfaceFormatVersionRegex();
   auto CompRe = getSwiftInterfaceCompilerVersionRegex();
   auto FlagRe = getSwiftInterfaceModuleFlagsRegex();
   SmallVector<StringRef, 1> VersMatches, FlagMatches, CompMatches;
+
   if (!VersRe.match(SB, &VersMatches)) {
-    diagnose(diag::error_extracting_version_from_module_interface);
+    ModuleInterfaceBuilder::diagnose(Diags, SM, InterfacePath, diagnosticLoc,
+      diag::error_extracting_version_from_module_interface);
     return true;
   }
   if (!FlagRe.match(SB, &FlagMatches)) {
-    diagnose(diag::error_extracting_flags_from_module_interface);
+    ModuleInterfaceBuilder::diagnose(Diags, SM, InterfacePath, diagnosticLoc,
+      diag::error_extracting_version_from_module_interface);
     return true;
   }
   assert(VersMatches.size() == 2);
   assert(FlagMatches.size() == 2);
   // FIXME We should diagnose this at a location that makes sense:
-  Vers = swift::version::Version(VersMatches[1], SourceLoc(), &diags);
+  Vers = swift::version::Version(VersMatches[1], SourceLoc(), &Diags);
   llvm::cl::TokenizeGNUCommandLine(FlagMatches[1], SubArgSaver, SubArgs);
 
   if (CompRe.match(SB, &CompMatches)) {
@@ -177,6 +187,13 @@ bool ModuleInterfaceBuilder::extractSwiftInterfaceVersionAndArgs(
   return false;
 }
 
+bool ModuleInterfaceBuilder::extractSwiftInterfaceVersionAndArgs(
+    swift::version::Version &Vers, StringRef &CompilerVersion,
+    llvm::StringSaver &SubArgSaver, SmallVectorImpl<const char *> &SubArgs) {
+  return swift::extractSwiftInterfaceVersionAndArgs(sourceMgr, diags,
+    interfacePath, Vers, CompilerVersion, SubArgSaver, SubArgs, diagnosticLoc);
+}
+
 bool ModuleInterfaceBuilder::collectDepsForSerialization(
     CompilerInstance &SubInstance, SmallVectorImpl<FileDependency> &Deps,
     bool IsHashBased) {

lib/Frontend/ModuleInterfaceBuilder.h

@@ -48,17 +48,31 @@ class ModuleInterfaceBuilder {
   CompilerInvocation subInvocation;
   SmallVector<StringRef, 3> extraDependencies;
 
+public:
   /// Emit a diagnostic tied to this declaration.
   template<typename ...ArgTypes>
-  InFlightDiagnostic diagnose(
+  static InFlightDiagnostic diagnose(
+      DiagnosticEngine &Diags,
+      SourceManager &SM,
+      StringRef InterfacePath,
+      SourceLoc Loc,
       Diag<ArgTypes...> ID,
-      typename detail::PassArgument<ArgTypes>::type... Args) const {
-    SourceLoc loc = diagnosticLoc;
-    if (loc.isInvalid()) {
+      typename detail::PassArgument<ArgTypes>::type... Args) {
+    if (Loc.isInvalid()) {
       // Diagnose this inside the interface file, if possible.
-      loc = sourceMgr.getLocFromExternalSource(interfacePath, 1, 1);
+      Loc = SM.getLocFromExternalSource(InterfacePath, 1, 1);
     }
-    return diags.diagnose(loc, ID, std::move(Args)...);
+    return Diags.diagnose(Loc, ID, std::move(Args)...);
+  }
+
+private:
+  /// Emit a diagnostic tied to this declaration.
+  template<typename ...ArgTypes>
+  InFlightDiagnostic diagnose(
+      Diag<ArgTypes...> ID,
+      typename detail::PassArgument<ArgTypes>::type... Args) const {
+    return diagnose(diags, sourceMgr, interfacePath, diagnosticLoc,
+                    ID, std::move(Args)...);
   }
 
   void configureSubInvocationInputsAndOutputs(StringRef OutPath);

lib/FrontendTool/ScanDependencies.cpp

@@ -24,6 +24,7 @@
 #include "swift/Basic/STLExtras.h"
 #include "swift/Frontend/Frontend.h"
 #include "swift/Frontend/FrontendOptions.h"
+#include "swift/Frontend/ModuleInterfaceLoader.h"
 #include "swift/Strings.h"
 #include "clang/Basic/Module.h"
 #include "llvm/ADT/SetVector.h"
@@ -58,13 +59,48 @@ static void findAllImportedClangModules(ASTContext &ctx, StringRef moduleName,
   }
 }
 
+struct InterfaceSubASTContextDelegate: SubASTContextDelegate {
+  bool runInSubContext(ASTContext &ctx, StringRef interfacePath,
+                       llvm::function_ref<bool(ASTContext&)> action) override {
+    // Parse the interface file using the current context to get the additional
+    // compiler arguments we should use when creating a sub-ASTContext.
+    // These arguments are in "swift-module-flags:"
+    version::Version Vers;
+    StringRef CompilerVersion;
+    llvm::BumpPtrAllocator Allocator;
+    llvm::StringSaver SubArgSaver(Allocator);
+    SmallVector<const char *, 64> SubArgs;
+    if (extractSwiftInterfaceVersionAndArgs(ctx.SourceMgr, ctx.Diags,
+                                            interfacePath, Vers, CompilerVersion,
+                                            SubArgSaver, SubArgs)) {
+      return true;
+    }
+    CompilerInvocation invok;
+    CompilerInstance inst;
+    // Use the additional flags to setup the compiler instance.
+    if (invok.parseArgs(SubArgs, ctx.Diags)) {
+      return true;
+    }
+    if (inst.setup(invok)) {
+      return true;
+    }
+    // Add the diag consumers to the sub context to make sure we don't lose
+    // diagnostics.
+    for (auto *consumer: ctx.Diags.getConsumers()) {
+      inst.getDiags().addConsumer(*consumer);
+    }
+    // Run the action under the sub-ASTContext.
+    return action(inst.getASTContext());
+  }
+};
+
 /// Resolve the direct dependencies of the given module.
 static std::vector<ModuleDependencyID> resolveDirectDependencies(
     ASTContext &ctx, ModuleDependencyID module,
     ModuleDependenciesCache &cache) {
   auto knownDependencies = *cache.findDependencies(module.first, module.second);
   auto isSwift = knownDependencies.isSwiftModule();
-
+  InterfaceSubASTContextDelegate ASTDelegate;
   // Find the dependencies of every module this module directly depends on.
   std::vector<ModuleDependencyID> result;
   for (auto dependsOn : knownDependencies.getModuleDependencies()) {
@@ -73,7 +109,7 @@ static std::vector<ModuleDependencyID> resolveDirectDependencies(
 
     // Retrieve the dependencies for this module.
     if (auto found = ctx.getModuleDependencies(
-            dependsOn, onlyClangModule, cache)) {
+            dependsOn, onlyClangModule, cache, ASTDelegate)) {
       result.push_back({dependsOn, found->getKind()});
     }
   }
@@ -115,7 +151,7 @@ static std::vector<ModuleDependencyID> resolveDirectDependencies(
     // directly depends on these.
     for (const auto &clangDep : allClangModules) {
       if (auto found = ctx.getModuleDependencies(
-              clangDep, /*onlyClangModule=*/false, cache)) {
+              clangDep, /*onlyClangModule=*/false, cache, ASTDelegate)) {
         if (found->getKind() == ModuleDependenciesKind::Swift)
           result.push_back({clangDep, found->getKind()});
       }