[Macro][Dependencies] Properly model macro dependencies in the scanner

Add function to handle all macro dependencies kinds in the scanner,
including taking care of the macro definitions in the module interface
for its client to use. The change involves:
  * Encode the macro definition inside the binary module
  * Resolve macro modules in the dependencies scanners, including those
    declared inside the dependency modules.
  * Propagate the macro defined from the direct dependencies to track
    all the potentially available modules inside a module compilation.

Author: cachemeifyoucan

include/swift/AST/FileUnit.h

@@ -272,6 +272,10 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
   getImportedModules(SmallVectorImpl<ImportedModule> &imports,
                      ModuleDecl::ImportFilter filter) const {}
 
+  /// Looks up which external macros are defined by this file.
+  virtual void
+  getExternalMacros(SmallVectorImpl<ExternalMacroPlugin> &macros) const {}
+
   /// Lists modules that are not imported from this file and used in API.
   virtual void getImplicitImportsForModuleInterface(
       SmallVectorImpl<ImportedModule> &imports) const {}

include/swift/AST/Module.h

@@ -145,6 +145,19 @@ struct SourceFilePathInfo {
   }
 };
 
+/// This is used to idenfity an external macro definition.
+struct ExternalMacroPlugin {
+  std::string ModuleName;
+
+  enum Access {
+    Internal = 0,
+    Package,
+    Public,
+  };
+
+  Access MacroAccess;
+};
+
 /// Discriminator for resilience strategy.
 enum class ResilienceStrategy : unsigned {
   /// Public nominal types: fragile
@@ -969,6 +982,9 @@ class ModuleDecl
   void getImportedModules(SmallVectorImpl<ImportedModule> &imports,
                           ImportFilter filter = ImportFilterKind::Exported) const;
 
+  /// Looks up which external macros are defined by this file.
+  void getExternalMacros(SmallVectorImpl<ExternalMacroPlugin> &macros) const;
+
   /// Lists modules that are not imported from a file and used in API.
   void getImplicitImportsForModuleInterface(
       SmallVectorImpl<ImportedModule> &imports) const;

include/swift/AST/ModuleDependencies.h

@@ -224,6 +224,9 @@ class ModuleDependencyInfoStorageBase {
   /// command-line), no need to be saved to reconstruct from cache.
   std::vector<std::string> auxiliaryFiles;
 
+  /// The macro dependencies.
+  std::map<std::string, MacroPluginDependency> macroDependencies;
+
   /// The direct dependency of the module is resolved by scanner.
   bool resolved;
   /// ModuleDependencyInfo is finalized (with all transitive dependencies
@@ -253,9 +256,6 @@ struct CommonSwiftTextualModuleDependencyDetails {
   /// (Clang) modules on which the bridging header depends.
   std::vector<std::string> bridgingModuleDependencies;
 
-  /// The macro dependencies.
-  std::map<std::string, MacroPluginDependency> macroDependencies;
-
   /// The Swift frontend invocation arguments to build the Swift module from the
   /// interface.
   std::vector<std::string> buildCommandLine;
@@ -326,12 +326,6 @@ class SwiftInterfaceModuleDependenciesStorage
   void updateCommandLine(const std::vector<std::string> &newCommandLine) {
     textualModuleDetails.buildCommandLine = newCommandLine;
   }
-
-  void addMacroDependency(StringRef macroModuleName, StringRef libraryPath,
-                          StringRef executablePath) {
-    textualModuleDetails.macroDependencies.insert(
-        {macroModuleName.str(), {libraryPath.str(), executablePath.str()}});
-  }
 };
 
 /// Describes the dependencies of a Swift module
@@ -382,12 +376,6 @@ class SwiftSourceModuleDependenciesStorage
   void addTestableImport(ImportPath::Module module) {
     testableImports.insert(module.front().Item.str());
   }
-
-  void addMacroDependency(StringRef macroModuleName, StringRef libraryPath,
-                          StringRef executablePath) {
-    textualModuleDetails.macroDependencies.insert(
-        {macroModuleName.str(), {libraryPath.str(), executablePath.str()}});
-  }
 };
 
 /// Describes the dependencies of a pre-built Swift module (with no
@@ -786,6 +774,16 @@ class ModuleDependencyInfo {
     storage->auxiliaryFiles.emplace_back(file);
   }
 
+  void addMacroDependency(StringRef macroModuleName, StringRef libraryPath,
+                          StringRef executablePath) {
+    storage->macroDependencies.insert(
+        {macroModuleName.str(), {libraryPath.str(), executablePath.str()}});
+  }
+
+  std::map<std::string, MacroPluginDependency> &getMacroDependencies() const {
+    return storage->macroDependencies;
+  }
+
   void updateCASFileSystemRootID(const std::string &rootID) {
     if (isSwiftInterfaceModule())
       cast<SwiftInterfaceModuleDependenciesStorage>(storage.get())
@@ -809,13 +807,6 @@ class ModuleDependencyInfo {
   /// For a Source dependency, register a `Testable` import
   void addTestableImport(ImportPath::Module module);
 
-  /// For a Source/Textual dependency, register a macro dependency.
-  void addMacroDependency(StringRef macroModuleName, StringRef libraryPath,
-                          StringRef executablePath);
-
-  /// For a Source/Textual dependency, if it Has macro dependency.
-  bool hasMacroDependencies() const;
-
   /// Whether or not a queried module name is a `@Testable` import dependency
   /// of this module. Can only return `true` for Swift source modules.
   bool isTestableImport(StringRef moduleName) const;

include/swift/AST/SourceFile.h

@@ -576,6 +576,9 @@ class SourceFile final : public FileUnit {
   getImportedModules(SmallVectorImpl<ImportedModule> &imports,
                      ModuleDecl::ImportFilter filter) const override;
 
+  virtual void getExternalMacros(
+      SmallVectorImpl<ExternalMacroPlugin> &macros) const override;
+
   virtual void
   collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const override;
 

include/swift/DependencyScan/DependencyScanImpl.h

@@ -167,6 +167,9 @@ typedef struct {
   /// A flag that indicates this dependency is associated with a static archive
   bool is_static;
 
+  /// Macro dependecies.
+  swiftscan_macro_dependency_set_t *macro_dependencies;
+
   /// ModuleCacheKey
   swiftscan_string_ref_t module_cache_key;
 

include/swift/Serialization/SerializedModuleLoader.h

@@ -15,6 +15,7 @@
 
 #include "swift/AST/FileUnit.h"
 #include "swift/AST/Module.h"
+#include "swift/AST/ModuleDependencies.h"
 #include "swift/AST/ModuleLoader.h"
 #include "swift/AST/SearchPathOptions.h"
 #include "llvm/Support/MemoryBuffer.h"
@@ -191,6 +192,9 @@ class SerializedModuleLoaderBase : public ModuleLoader {
                                         llvm::vfs::FileSystem *fileSystem,
                                         PathObfuscator &recoverer);
 
+  std::optional<MacroPluginDependency>
+  resolveMacroPlugin(const ExternalMacroPlugin &macro, StringRef packageName);
+
 public:
   virtual ~SerializedModuleLoaderBase();
   SerializedModuleLoaderBase(const SerializedModuleLoaderBase &) = delete;
@@ -514,6 +518,9 @@ class SerializedASTFile final : public LoadedFile {
   getImportedModules(SmallVectorImpl<ImportedModule> &imports,
                      ModuleDecl::ImportFilter filter) const override;
 
+  virtual void getExternalMacros(
+      SmallVectorImpl<ExternalMacroPlugin> &macros) const override;
+
   virtual void
   collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const override;
 

lib/AST/Module.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/ASTPrinter.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/AccessScope.h"
+#include "swift/AST/AttrKind.h"
 #include "swift/AST/Builtins.h"
 #include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/DiagnosticsSema.h"
@@ -30,6 +31,7 @@
 #include "swift/AST/ImportCache.h"
 #include "swift/AST/LazyResolver.h"
 #include "swift/AST/LinkLibrary.h"
+#include "swift/AST/MacroDefinition.h"
 #include "swift/AST/ModuleLoader.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/NameLookupRequests.h"
@@ -1614,6 +1616,11 @@ void ModuleDecl::getImportedModules(SmallVectorImpl<ImportedModule> &modules,
   FORWARD(getImportedModules, (modules, filter));
 }
 
+void ModuleDecl::getExternalMacros(
+    SmallVectorImpl<ExternalMacroPlugin> &macros) const {
+  FORWARD(getExternalMacros, (macros));
+}
+
 void ModuleDecl::getImplicitImportsForModuleInterface(
     SmallVectorImpl<ImportedModule> &imports) const {
   FORWARD(getImplicitImportsForModuleInterface, (imports));
@@ -1705,6 +1712,29 @@ SourceFile::getImportedModules(SmallVectorImpl<ImportedModule> &modules,
   }
 }
 
+void SourceFile::getExternalMacros(
+    SmallVectorImpl<ExternalMacroPlugin> &macros) const {
+  for (auto *D : getTopLevelDecls()) {
+    auto macroDecl = dyn_cast<MacroDecl>(D);
+    if (!macroDecl)
+      continue;
+
+    auto macroDef = macroDecl->getDefinition();
+    if (macroDef.kind != MacroDefinition::Kind::External)
+      continue;
+
+    auto formalAccess = macroDecl->getFormalAccess();
+    ExternalMacroPlugin::Access access = ExternalMacroPlugin::Access::Internal;
+    if (formalAccess >= AccessLevel::Public)
+      access = ExternalMacroPlugin::Access::Public;
+    else if (formalAccess >= AccessLevel::Package)
+      access = ExternalMacroPlugin::Access::Package;
+
+    auto external = macroDef.getExternalMacro();
+    macros.push_back({external.moduleName.str().str(), access});
+  }
+}
+
 void SourceFile::getImplicitImportsForModuleInterface(
     SmallVectorImpl<ImportedModule> &modules) const {
   for (auto module : ImplicitImportsForModuleInterface)

lib/AST/ModuleDependencies.cpp

@@ -18,6 +18,7 @@
 #include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/AST/DiagnosticsSema.h"
 #include "swift/AST/MacroDefinition.h"
+#include "swift/AST/Module.h"
 #include "swift/AST/PluginLoader.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/Frontend/Frontend.h"
@@ -105,33 +106,6 @@ void ModuleDependencyInfo::addTestableImport(ImportPath::Module module) {
   dyn_cast<SwiftSourceModuleDependenciesStorage>(storage.get())->addTestableImport(module);
 }
 
-void ModuleDependencyInfo::addMacroDependency(StringRef macroModuleName,
-                                              StringRef libraryPath,
-                                              StringRef executablePath) {
-  if (auto swiftSourceStorage =
-          dyn_cast<SwiftSourceModuleDependenciesStorage>(storage.get()))
-    swiftSourceStorage->addMacroDependency(macroModuleName, libraryPath,
-                                           executablePath);
-  else if (auto swiftInterfaceStorage =
-               dyn_cast<SwiftInterfaceModuleDependenciesStorage>(storage.get()))
-    swiftInterfaceStorage->addMacroDependency(macroModuleName, libraryPath,
-                                              executablePath);
-  else
-    llvm_unreachable("Unexpected dependency kind");
-}
-
-bool ModuleDependencyInfo::hasMacroDependencies() const {
-  if (auto sourceModule =
-          dyn_cast<SwiftSourceModuleDependenciesStorage>(storage.get()))
-    return !sourceModule->textualModuleDetails.macroDependencies.empty();
-
-  if (auto interfaceModule =
-          dyn_cast<SwiftInterfaceModuleDependenciesStorage>(storage.get()))
-    return !interfaceModule->textualModuleDetails.macroDependencies.empty();
-
-  llvm_unreachable("Unexpected dependency kind");
-}
-
 bool ModuleDependencyInfo::isTestableImport(StringRef moduleName) const {
   if (auto swiftSourceDepStorage = getAsSwiftSourceModule())
     return swiftSourceDepStorage->testableImports.contains(moduleName);

lib/DependencyScan/DependencyScanJSON.cpp

@@ -601,6 +601,8 @@ void writeJSON(llvm::raw_ostream &out,
                           /*trailingComma=*/true);
       }
 
+      writeMacroDependencies(out, swiftBinaryDeps->macro_dependencies, 5,
+                             /*trailingComma=*/true);
       writeJSONSingleField(out, "isFramework", swiftBinaryDeps->is_framework,
                            5, /*trailingComma=*/false);
     } else {

lib/DependencyScan/ModuleDependencyScanner.cpp

@@ -17,6 +17,7 @@
 #include "swift/AST/DiagnosticsSema.h"
 #include "swift/AST/ModuleDependencies.h"
 #include "swift/AST/ModuleLoader.h"
+#include "swift/AST/PluginLoader.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Assertions.h"
@@ -180,6 +181,10 @@ ModuleDependencyScanningWorker::ModuleDependencyScanningWorker(
                       workerCompilerInvocation->getSymbolGraphOptions(),
                       workerCompilerInvocation->getCASOptions(),
                       ScanASTContext.SourceMgr, Diagnostics));
+  auto loader = std::make_unique<PluginLoader>(
+      *workerASTContext, /*DepTracker=*/nullptr,
+      workerCompilerInvocation->getFrontendOptions().DisableSandbox);
+  workerASTContext->setPluginLoader(std::move(loader));
 
   // Configure the interface scanning AST delegate.
   auto ClangModuleCachePath = getModuleCachePathFromClang(