[Macros] Track plugin dependencies

* Factor out ASTContext plugin loading to newly introduced 'PluginLoader'
* Insert 'DependencyTracker' to 'PluginLoader'
* Add dependencies right before loading the plugins

rdar://104938481

Author: rintaro

include/swift/AST/ASTContext.h

@@ -95,7 +95,7 @@ namespace swift {
   class ModuleDependencyInfo;
   class PatternBindingDecl;
   class PatternBindingInitializer;
-  class PluginRegistry;
+  class PluginLoader;
   class SourceFile;
   class SourceLoc;
   class Type;
@@ -1488,6 +1488,12 @@ class ASTContext final {
 
   Type getNamedSwiftType(ModuleDecl *module, StringRef name);
 
+  /// Set the plugin loader.
+  void setPluginLoader(std::unique_ptr<PluginLoader> loader);
+
+  /// Get the plugin loader.
+  PluginLoader &getPluginLoader();
+
   /// Lookup a library plugin that can handle \p moduleName and return the path
   /// to it.
   /// The path is valid within the VFS, use `FS.getRealPath()` for the
@@ -1521,15 +1527,6 @@ class ASTContext final {
   /// instance is simply returned.
   LoadedExecutablePlugin *loadExecutablePlugin(StringRef path);
 
-  /// Get the plugin registry this ASTContext is using.
-  PluginRegistry *getPluginRegistry() const;
-
-  /// Set the plugin registory this ASTContext should use.
-  /// This should be called before any plugin is loaded.
-  void setPluginRegistry(PluginRegistry *newValue);
-
-  const llvm::StringSet<> &getLoadedPluginLibraryPaths() const;
-
   /// Get the output backend. The output backend needs to be initialized via
   /// constructor or `setOutputBackend`.
   llvm::vfs::OutputBackend &getOutputBackend() const {
@@ -1554,8 +1551,6 @@ class ASTContext final {
   Optional<StringRef> getBriefComment(const Decl *D);
   void setBriefComment(const Decl *D, StringRef Comment);
 
-  void createModuleToExecutablePluginMap();
-
   friend TypeBase;
   friend ArchetypeType;
   friend OpaqueTypeDecl;

include/swift/AST/PluginLoader.h

@@ -0,0 +1,86 @@
+//===--- PluginLoader.h -----------------------------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+#ifndef SWIFT_AST_PLUGIN_LOADER_H
+#define SWIFT_AST_PLUGIN_LOADER_H
+
+#include "swift/AST/ModuleLoader.h"
+#include "swift/AST/PluginRegistry.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace swift {
+
+class ASTContext;
+
+/// Compiler plugin loader tied to an ASTContext. This is responsible for
+///  * Find plugins based on the module name
+///  * Load plugins resolving VFS paths
+///  * Track plugin dependencies
+class PluginLoader {
+  PluginRegistry *Registry = nullptr;
+  std::unique_ptr<PluginRegistry> OwnedRegistry = nullptr;
+  ASTContext &Ctx;
+  DependencyTracker *DepTracker;
+
+  /// Map a module name to an executable plugin path that provides the module.
+  llvm::DenseMap<swift::Identifier, llvm::StringRef> ExecutablePluginPaths;
+
+  void createModuleToExecutablePluginMap();
+
+public:
+  PluginLoader(ASTContext &Ctx, DependencyTracker *DepTracker)
+      : Ctx(Ctx), DepTracker(DepTracker) {
+    createModuleToExecutablePluginMap();
+  }
+
+  void setRegistry(PluginRegistry *newValue);
+  PluginRegistry *getRegistry();
+
+  /// Lookup a library plugin that can handle \p moduleName and return the path
+  /// to it from `-plugin-path` or `-load-plugin-library`.
+  /// The path returned can be loaded by 'loadLibraryPlugin' method.
+  llvm::Optional<std::string>
+  lookupLibraryPluginByModuleName(Identifier moduleName);
+
+  /// Lookup an executable plugin that is declared to handle \p moduleName
+  /// module by '-load-plugin-executable'.
+  /// The path returned can be loaded by 'loadExecutablePlugin' method.
+  llvm::Optional<StringRef>
+  lookupExecutablePluginByModuleName(Identifier moduleName);
+
+  /// Look for dynamic libraries in paths from `-external-plugin-path` and
+  /// return a pair of `(library path, plugin server executable)` if found.
+  /// These paths are valid within the VFS, use `FS.getRealPath()` for their
+  /// underlying path.
+  llvm::Optional<std::pair<std::string, std::string>>
+  lookupExternalLibraryPluginByModuleName(Identifier moduleName);
+
+  /// Load the specified dylib plugin path resolving the path with the
+  /// current VFS. If it fails to load the plugin, a diagnostic is emitted, and
+  /// returns a nullptr.
+  /// NOTE: This method is idempotent. If the plugin is already loaded, the same
+  /// instance is simply returned.
+  LoadedLibraryPlugin *loadLibraryPlugin(llvm::StringRef path);
+
+  /// Launch the specified executable plugin path resolving the path with the
+  /// current VFS. If it fails to load the plugin, a diagnostic is emitted, and
+  /// returns a nullptr.
+  /// NOTE: This method is idempotent. If the plugin is already loaded, the same
+  /// instance is simply returned.
+  LoadedExecutablePlugin *loadExecutablePlugin(llvm::StringRef path);
+};
+
+} // namespace swift
+
+#endif // SWIFT_AST_PLUGIN_LOADER_H

include/swift/Frontend/Frontend.h

@@ -663,6 +663,7 @@ class CompilerInstance {
   void setUpLLVMArguments();
   void setUpDiagnosticOptions();
   bool setUpModuleLoaders();
+  bool setUpPluginLoader();
   bool setUpInputs();
   bool setUpASTContextIfNeeded();
   void setupStatsReporter();

lib/AST/ASTContext.cpp

@@ -41,7 +41,7 @@
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/PackConformance.h"
 #include "swift/AST/ParameterList.h"
-#include "swift/AST/PluginRegistry.h"
+#include "swift/AST/PluginLoader.h"
 #include "swift/AST/PrettyStackTrace.h"
 #include "swift/AST/PropertyWrappers.h"
 #include "swift/AST/ProtocolConformance.h"
@@ -526,17 +526,8 @@ struct ASTContext::Implementation {
 
   llvm::StringMap<OptionSet<SearchPathKind>> SearchPathsSet;
 
-  /// Plugin registry. Lazily populated by get/setPluginRegistry().
-  /// NOTE: Do not reference this directly. Use ASTContext::getPluginRegistry().
-  PluginRegistry *Plugins = nullptr;
-
-  /// `Plugins` storage if this ASTContext owns it.
-  std::unique_ptr<PluginRegistry> OwnedPluginRegistry = nullptr;
-
-  /// Map a module name to an executable plugin path that provides the module.
-  llvm::DenseMap<Identifier, StringRef> ExecutablePluginPaths;
-
-  llvm::StringSet<> LoadedPluginLibraryPaths;
+  /// Plugin loader.
+  std::unique_ptr<swift::PluginLoader> Plugins;
 
   /// The permanent arena.
   Arena Permanent;
@@ -712,8 +703,6 @@ ASTContext::ASTContext(
   registerAccessRequestFunctions(evaluator);
   registerNameLookupRequestFunctions(evaluator);
 
-  createModuleToExecutablePluginMap();
-  
   // Provide a default OnDiskOutputBackend if user didn't supply one.
   if (!OutputBackend)
     OutputBackend = llvm::makeIntrusiveRefCnt<llvm::vfs::OnDiskOutputBackend>();
@@ -6280,34 +6269,33 @@ BuiltinTupleType *ASTContext::getBuiltinTupleType() {
   return result;
 }
 
-void ASTContext::setPluginRegistry(PluginRegistry *newValue) {
-  assert(getImpl().Plugins == nullptr &&
-         "Too late to set a new plugin registry");
-  getImpl().Plugins = newValue;
+void ASTContext::setPluginLoader(std::unique_ptr<PluginLoader> loader) {
+  getImpl().Plugins = std::move(loader);
 }
 
-PluginRegistry *ASTContext::getPluginRegistry() const {
-  PluginRegistry *&registry = getImpl().Plugins;
+PluginLoader &ASTContext::getPluginLoader() { return *getImpl().Plugins; }
 
-  // Create a new one if it hasn't been set.
-  if (!registry) {
-    registry = new PluginRegistry();
-    getImpl().OwnedPluginRegistry.reset(registry);
-  }
+Optional<std::string>
+ASTContext::lookupLibraryPluginByModuleName(Identifier moduleName) {
+  return getImpl().Plugins->lookupLibraryPluginByModuleName(moduleName);
+}
 
-  assert(registry != nullptr);
-  return registry;
+Optional<StringRef>
+ASTContext::lookupExecutablePluginByModuleName(Identifier moduleName) {
+  return getImpl().Plugins->lookupExecutablePluginByModuleName(moduleName);
 }
 
-void ASTContext::createModuleToExecutablePluginMap() {
-  for (auto &arg : SearchPathOpts.getCompilerPluginExecutablePaths()) {
-    // Create a moduleName -> pluginPath mapping.
-    assert(!arg.ExecutablePath.empty() && "empty plugin path");
-    auto pathStr = AllocateCopy(arg.ExecutablePath);
-    for (auto moduleName : arg.ModuleNames) {
-      getImpl().ExecutablePluginPaths[getIdentifier(moduleName)] = pathStr;
-    }
-  }
+Optional<std::pair<std::string, std::string>>
+ASTContext::lookupExternalLibraryPluginByModuleName(Identifier moduleName) {
+  return getImpl().Plugins->lookupExternalLibraryPluginByModuleName(moduleName);
+}
+
+LoadedLibraryPlugin *ASTContext::loadLibraryPlugin(StringRef path) {
+  return getImpl().Plugins->loadLibraryPlugin(path);
+}
+
+LoadedExecutablePlugin *ASTContext::loadExecutablePlugin(StringRef path) {
+  return getImpl().Plugins->loadExecutablePlugin(path);
 }
 
 Type ASTContext::getNamedSwiftType(ModuleDecl *module, StringRef name) {
@@ -6345,105 +6333,6 @@ Type ASTContext::getNamedSwiftType(ModuleDecl *module, StringRef name) {
   return decl->getDeclaredInterfaceType();
 }
 
-Optional<std::string>
-ASTContext::lookupLibraryPluginByModuleName(Identifier moduleName) {
-  auto fs = SourceMgr.getFileSystem();
-
-  // Look for 'lib${module name}(.dylib|.so)'.
-  SmallString<64> expectedBasename;
-  expectedBasename.append("lib");
-  expectedBasename.append(moduleName.str());
-  expectedBasename.append(LTDL_SHLIB_EXT);
-
-  // Try '-plugin-path'.
-  for (const auto &searchPath : SearchPathOpts.PluginSearchPaths) {
-    SmallString<128> fullPath(searchPath);
-    llvm::sys::path::append(fullPath, expectedBasename);
-    if (fs->exists(fullPath)) {
-      return std::string(fullPath);
-    }
-  }
-
-  // Try '-load-plugin-library'.
-  for (const auto &libPath : SearchPathOpts.getCompilerPluginLibraryPaths()) {
-    if (llvm::sys::path::filename(libPath) == expectedBasename) {
-      return libPath;
-    }
-  }
-
-  return None;
-}
-
-Optional<StringRef>
-ASTContext::lookupExecutablePluginByModuleName(Identifier moduleName) {
-  auto &execPluginPaths = getImpl().ExecutablePluginPaths;
-  auto found = execPluginPaths.find(moduleName);
-  if (found == execPluginPaths.end())
-    return None;
-  return found->second;
-}
-
-Optional<std::pair<std::string, std::string>>
-ASTContext::lookupExternalLibraryPluginByModuleName(Identifier moduleName) {
-  auto fs = this->SourceMgr.getFileSystem();
-  for (auto &pair : SearchPathOpts.ExternalPluginSearchPaths) {
-    SmallString<128> fullPath(pair.SearchPath);
-    llvm::sys::path::append(fullPath, "lib" + moduleName.str() + LTDL_SHLIB_EXT);
-
-    if (fs->exists(fullPath)) {
-      return {{std::string(fullPath), pair.ServerPath}};
-    }
-  }
-  return None;
-}
-
-LoadedExecutablePlugin *ASTContext::loadExecutablePlugin(StringRef path) {
-  SmallString<128> resolvedPath;
-  auto fs = this->SourceMgr.getFileSystem();
-  if (auto err = fs->getRealPath(path, resolvedPath)) {
-    Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
-                   err.message());
-    return nullptr;
-  }
-
-  // Load the plugin.
-  auto plugin = getPluginRegistry()->loadExecutablePlugin(resolvedPath);
-  if (!plugin) {
-    Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
-                   llvm::toString(plugin.takeError()));
-    return nullptr;
-  }
-
-  return plugin.get();
-}
-
-LoadedLibraryPlugin *ASTContext::loadLibraryPlugin(StringRef path) {
-  // Remember the path (even if it fails to load.)
-  getImpl().LoadedPluginLibraryPaths.insert(path);
-
-  SmallString<128> resolvedPath;
-  auto fs = this->SourceMgr.getFileSystem();
-  if (auto err = fs->getRealPath(path, resolvedPath)) {
-    Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
-                   err.message());
-    return nullptr;
-  }
-
-  // Load the plugin.
-  auto plugin = getPluginRegistry()->loadLibraryPlugin(resolvedPath);
-  if (!plugin) {
-    Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
-                   llvm::toString(plugin.takeError()));
-    return nullptr;
-  }
-
-  return plugin.get();
-}
-
-const llvm::StringSet<> &ASTContext::getLoadedPluginLibraryPaths() const {
-  return getImpl().LoadedPluginLibraryPaths;
-}
-
 bool ASTContext::supportsMoveOnlyTypes() const {
   // currently the only thing holding back whether the types can appear is this.
   return SILOpts.LexicalLifetimes != LexicalLifetimesOption::Off;

lib/AST/CMakeLists.txt

@@ -72,6 +72,7 @@ add_swift_host_library(swiftAST STATIC
   Parameter.cpp
   Pattern.cpp
   PlatformKind.cpp
+  PluginLoader.cpp
   PluginRegistry.cpp
   PrettyStackTrace.cpp
   ProtocolConformance.cpp