[5.9][Macros] Track plugin dependencies

include/swift/AST/ASTContext.h

@@ -80,6 +80,7 @@ namespace swift {
   class ExtensionDecl;
   struct ExternalSourceLocs;
   class LoadedExecutablePlugin;
+  class LoadedLibraryPlugin;
   class ForeignRepresentationInfo;
   class FuncDecl;
   class GenericContext;
@@ -92,7 +93,7 @@ namespace swift {
   class ModuleDependencyInfo;
   class PatternBindingDecl;
   class PatternBindingInitializer;
-  class PluginRegistry;
+  class PluginLoader;
   class SourceFile;
   class SourceLoc;
   class Type;
@@ -1476,6 +1477,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
@@ -1487,7 +1494,7 @@ class ASTContext final {
   /// returns a nullptr.
   /// NOTE: This method is idempotent. If the plugin is already loaded, the same
   /// instance is simply returned.
-  void *loadLibraryPlugin(StringRef path);
+  LoadedLibraryPlugin *loadLibraryPlugin(StringRef path);
 
   /// Lookup an executable plugin that is declared to handle \p moduleName
   /// module by '-load-plugin-executable'.
@@ -1509,15 +1516,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;
-
 private:
   friend Decl;
 
@@ -1530,8 +1528,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,91 @@
+//===--- 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 {
+  /// Plugin registry. Lazily populated by get/setRegistry().
+  /// NOTE: Do not reference this directly. Use getRegistry().
+  PluginRegistry *Registry = nullptr;
+
+  /// `Registry` storage if this owns it.
+  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/AST/PluginRegistry.h

@@ -9,6 +9,8 @@
 // See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 //
 //===----------------------------------------------------------------------===//
+#ifndef SWIFT_PLUGIN_REGISTRY_H
+#define SWIFT_PLUGIN_REGISTRY_H
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringMap.h"
@@ -22,6 +24,21 @@
 
 namespace swift {
 
+/// Represent a 'dlopen'ed plugin library.
+class LoadedLibraryPlugin {
+  // Opaque handle used to interface with OS-specific dynamic library.
+  void *handle;
+
+  /// Cache of loaded symbols.
+  llvm::StringMap<void *> resolvedSymbols;
+
+public:
+  LoadedLibraryPlugin(void *handle) : handle(handle) {}
+
+  /// Finds the address of the given symbol within the library.
+  void *getAddressOfSymbol(const char *symbolName);
+};
+
 /// Represent a "resolved" exectuable plugin.
 ///
 /// Plugin clients usually deal with this object to communicate with the actual
@@ -130,7 +147,7 @@ class LoadedExecutablePlugin {
 
 class PluginRegistry {
   /// Record of loaded plugin library modules.
-  llvm::StringMap<void *> LoadedPluginLibraries;
+  llvm::StringMap<std::unique_ptr<LoadedLibraryPlugin>> LoadedPluginLibraries;
 
   /// Record of loaded plugin executables.
   llvm::StringMap<std::unique_ptr<LoadedExecutablePlugin>>
@@ -146,7 +163,7 @@ class PluginRegistry {
 
   /// Load a dynamic link library specified by \p path.
   /// If \p path plugin is already loaded, this returns the cached object.
-  llvm::Expected<void *> loadLibraryPlugin(llvm::StringRef path);
+  llvm::Expected<LoadedLibraryPlugin *> loadLibraryPlugin(llvm::StringRef path);
 
   /// Load an executable plugin specified by \p path .
   /// If \p path plugin is already loaded, this returns the cached object.
@@ -155,3 +172,5 @@ class PluginRegistry {
 };
 
 } // namespace swift
+
+#endif // SWIFT_PLUGIN_REGISTRY_H

include/swift/AST/TypeCheckRequests.h

@@ -26,6 +26,7 @@
 #include "swift/AST/Type.h"
 #include "swift/AST/Evaluator.h"
 #include "swift/AST/Pattern.h"
+#include "swift/AST/PluginRegistry.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/SimpleRequest.h"
 #include "swift/AST/SourceFile.h"
@@ -49,7 +50,6 @@ struct ExternalMacroDefinition;
 class ClosureExpr;
 class GenericParamList;
 class LabeledStmt;
-class LoadedExecutablePlugin;
 class MacroDefinition;
 class PrecedenceGroupDecl;
 class PropertyWrapperInitializerInfo;
@@ -4017,39 +4017,20 @@ class ExpandSynthesizedMemberMacroRequest
   void noteCycleStep(DiagnosticEngine &diags) const;
 };
 
-/// Load a plugin module with the given name.
-///
-///
+/// Represent a loaded plugin either an in-process library or an executable.
 class LoadedCompilerPlugin {
-  enum class PluginKind : uint8_t {
-    None,
-    InProcess,
-    Executable,
-  };
-  PluginKind kind;
-  void *ptr;
-
-  LoadedCompilerPlugin(PluginKind kind, void *ptr) : kind(kind), ptr(ptr) {
-    assert(ptr != nullptr || kind == PluginKind::None);
-  }
+  llvm::PointerUnion<LoadedLibraryPlugin *, LoadedExecutablePlugin *> ptr;
 
 public:
-  LoadedCompilerPlugin(std::nullptr_t) : kind(PluginKind::None), ptr(nullptr) {}
-
-  static LoadedCompilerPlugin inProcess(void *ptr) {
-    return {PluginKind::InProcess, ptr};
-  }
-  static LoadedCompilerPlugin executable(LoadedExecutablePlugin *ptr) {
-    return {PluginKind::Executable, ptr};
-  }
+  LoadedCompilerPlugin(std::nullptr_t) : ptr(nullptr) {}
+  LoadedCompilerPlugin(LoadedLibraryPlugin *ptr) : ptr(ptr){};
+  LoadedCompilerPlugin(LoadedExecutablePlugin *ptr) : ptr(ptr){};
 
-  void *getAsInProcessPlugin() const {
-    return kind == PluginKind::InProcess ? ptr : nullptr;
+  LoadedLibraryPlugin *getAsLibraryPlugin() const {
+    return ptr.dyn_cast<LoadedLibraryPlugin *>();
   }
   LoadedExecutablePlugin *getAsExecutablePlugin() const {
-    return kind == PluginKind::Executable
-               ? static_cast<LoadedExecutablePlugin *>(ptr)
-               : nullptr;
+    return ptr.dyn_cast<LoadedExecutablePlugin *>();
   }
 };
 

include/swift/Frontend/Frontend.h

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