Merge pull request swiftlang#63558 from DougGregor/plugin-search-paths

DougGregor · web-flow · commit fad0f90f4be6 · 2023-02-10T00:14:22.000-08:00
diff --git a/include/swift/AST/SearchPathOptions.h b/include/swift/AST/SearchPathOptions.h
@@ -339,6 +339,10 @@ class SearchPathOptions {
   /// preference.
   std::vector<std::string> RuntimeLibraryPaths;
 
+  /// Paths that contain compiler plugins loaded on demand for, e.g.,
+  /// macro implementations.
+  std::vector<std::string> PluginSearchPaths;
+
   /// Don't look in for compiler-provided modules.
   bool SkipRuntimeLibraryImportPaths = false;
 
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -3860,6 +3860,30 @@ class ExpandSynthesizedMemberMacroRequest
   bool isCached() const { return true; }
 };
 
+/// Load a plugin module with the given name.
+///
+///
+class CompilerPluginLoadRequest
+  : public SimpleRequest<CompilerPluginLoadRequest,
+                         void *(ASTContext *, Identifier),
+                         RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  void *evaluate(
+      Evaluator &evaluator, ASTContext *ctx, Identifier moduleName
+  ) const;
+
+public:
+  // Source location
+  SourceLoc getNearestLoc() const { return SourceLoc(); }
+
+  bool isCached() const { return true; }
+};
+
 /// Resolve an external macro given its module and type name.
 class ExternalMacroDefinitionRequest
     : public SimpleRequest<ExternalMacroDefinitionRequest,
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -422,6 +422,9 @@ SWIFT_REQUEST(TypeChecker, IsSingleValueStmtRequest,
 SWIFT_REQUEST(TypeChecker, MacroDefinitionRequest,
               MacroDefinition(MacroDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, CompilerPluginLoadRequest,
+              void *(ASTContext *, Identifier),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExternalMacroDefinitionRequest,
               ExternalMacroDefinition(ASTContext *, Identifier, Identifier),
               Cached, NoLocationInfo)
diff --git a/include/swift/Option/Options.td b/include/swift/Option/Options.td
@@ -271,6 +271,10 @@ def I : JoinedOrSeparate<["-"], "I">,
 def I_EQ : Joined<["-"], "I=">, Flags<[FrontendOption, ArgumentIsPath]>,
   Alias<I>;
 
+def plugin_path : Separate<["-"], "plugin-path">,
+  Flags<[FrontendOption, ArgumentIsPath, SwiftAPIExtractOption, SwiftSymbolGraphExtractOption, SwiftAPIDigesterOption]>,
+  HelpText<"Add directory to the plugin search path">;
+
 def import_underlying_module : Flag<["-"], "import-underlying-module">,
   Flags<[FrontendOption, NoInteractiveOption]>,
   HelpText<"Implicitly imports the Objective-C half of a module">;
diff --git a/lib/Driver/ToolChains.cpp b/lib/Driver/ToolChains.cpp
@@ -218,6 +218,7 @@ void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
   inputArgs.AddAllArgs(arguments, options::OPT_I);
   inputArgs.AddAllArgs(arguments, options::OPT_F, options::OPT_Fsystem);
   inputArgs.AddAllArgs(arguments, options::OPT_vfsoverlay);
+  inputArgs.AddAllArgs(arguments, options::OPT_plugin_path);
 
   inputArgs.AddLastArg(arguments, options::OPT_AssertConfig);
   inputArgs.AddLastArg(arguments, options::OPT_autolink_force_load);
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
@@ -1413,6 +1413,10 @@ static bool ParseSearchPathArgs(SearchPathOptions &Opts,
   }
   Opts.setFrameworkSearchPaths(FrameworkSearchPaths);
 
+  for (const Arg *A : Args.filtered(OPT_plugin_path)) {
+    Opts.PluginSearchPaths.push_back(resolveSearchPath(A->getValue()));
+  }
+
   for (const Arg *A : Args.filtered(OPT_L)) {
     Opts.LibrarySearchPaths.push_back(resolveSearchPath(A->getValue()));
   }
diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp
@@ -32,6 +32,15 @@
 #include "swift/Demangling/ManglingMacros.h"
 #include "swift/Parse/Lexer.h"
 #include "swift/Subsystems.h"
+#include "llvm/Config/config.h"
+
+#if defined(_WIN32)
+#define WIN32_LEAN_AND_MEAN
+#define NOMINMAX
+#include <windows.h>
+#else
+#include <dlfcn.h>
+#endif
 
 using namespace swift;
 
@@ -100,7 +109,9 @@ static std::string mangledNameForTypeMetadataAccessor(
 #if SWIFT_SWIFT_PARSER
 /// Look for macro's type metadata given its external module and type name.
 static void const *lookupMacroTypeMetadataByExternalName(
-    ASTContext &ctx, StringRef moduleName, StringRef typeName) {
+    ASTContext &ctx, StringRef moduleName, StringRef typeName,
+    void *libraryHint = nullptr
+) {
   // Look up the type metadata accessor as a struct, enum, or class.
   const Demangle::Node::Kind typeKinds[] = {
     Demangle::Node::Kind::Structure,
@@ -111,8 +122,8 @@ static void const *lookupMacroTypeMetadataByExternalName(
   void *accessorAddr = nullptr;
   for (auto typeKind : typeKinds) {
     auto symbolName = mangledNameForTypeMetadataAccessor(
-                                                         moduleName, typeName, typeKind);
-    accessorAddr = ctx.getAddressOfSymbol(symbolName.c_str());
+        moduleName, typeName, typeKind);
+    accessorAddr = ctx.getAddressOfSymbol(symbolName.c_str(), libraryHint);
     if (accessorAddr)
       break;
   }
@@ -289,20 +300,43 @@ MacroDefinition MacroDefinitionRequest::evaluate(
   return MacroDefinition::forExternal(*moduleName, *typeName);
 }
 
-ExternalMacroDefinition
-ExternalMacroDefinitionRequest::evaluate(
-    Evaluator &evaluator, ASTContext *ctx,
-    Identifier moduleName, Identifier typeName
+/// Load a plugin library based on a module name.
+static void *loadPluginByName(StringRef searchPath, StringRef moduleName) {
+  SmallString<128> fullPath(searchPath);
+  llvm::sys::path::append(fullPath, "lib" + moduleName + LTDL_SHLIB_EXT);
+#if defined(_WIN32)
+  return LoadLibraryA(fullPath.c_str());
+#else
+  return dlopen(fullPath.c_str(), RTLD_LAZY);
+#endif
+}
+
+void *CompilerPluginLoadRequest::evaluate(
+    Evaluator &evaluator, ASTContext *ctx, Identifier moduleName
 ) const {
+  auto &searchPathOpts = ctx->SearchPathOpts;
+  for (const auto &path : searchPathOpts.PluginSearchPaths) {
+    if (auto found = loadPluginByName(path, moduleName.str()))
+      return found;
+  }
+
+  return nullptr;
+}
+
+static Optional<ExternalMacroDefinition>
+resolveInProcessMacro(
+    ASTContext &ctx, Identifier moduleName, Identifier typeName,
+    void *libraryHint = nullptr
+) {
 #if SWIFT_SWIFT_PARSER
   /// Look for the type metadata given the external module and type names.
   auto macroMetatype = lookupMacroTypeMetadataByExternalName(
-      *ctx, moduleName.str(), typeName.str());
+      ctx, moduleName.str(), typeName.str(), libraryHint);
   if (macroMetatype) {
     // Check whether the macro metatype is in-process.
     if (auto inProcess = swift_ASTGen_resolveMacroType(macroMetatype)) {
       // Make sure we clean up after the macro.
-      ctx->addCleanup([inProcess]() {
+      ctx.addCleanup([inProcess]() {
         swift_ASTGen_destroyMacro(inProcess);
       });
 
@@ -311,6 +345,28 @@ ExternalMacroDefinitionRequest::evaluate(
   }
 #endif
 
+  return None;
+}
+
+ExternalMacroDefinition
+ExternalMacroDefinitionRequest::evaluate(
+    Evaluator &evaluator, ASTContext *ctx,
+    Identifier moduleName, Identifier typeName
+) const {
+  // Try to load a plugin module from the plugin search paths. If it
+  // succeeds, resolve in-process from that plugin
+  CompilerPluginLoadRequest loadRequest{ctx, moduleName};
+  if (auto loadedLibrary = evaluateOrDefault(
+          evaluator, loadRequest, nullptr)) {
+    if (auto inProcess = resolveInProcessMacro(
+            *ctx, moduleName, typeName, loadedLibrary))
+      return *inProcess;
+  }
+
+  // Try to resolve in-process.
+  if (auto inProcess = resolveInProcessMacro(*ctx, moduleName, typeName))
+    return *inProcess;
+
   return ExternalMacroDefinition{nullptr};
 }
 
diff --git a/test/Macros/macro_expand.swift b/test/Macros/macro_expand.swift
@@ -5,6 +5,8 @@
 // Diagnostics testing
 // RUN: %target-typecheck-verify-swift -swift-version 5 -enable-experimental-feature Macros -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -module-name MacroUser -DTEST_DIAGNOSTICS
 
+// RUN: %target-typecheck-verify-swift -swift-version 5 -enable-experimental-feature Macros -plugin-path %t -I %swift-host-lib-dir -module-name MacroUser -DTEST_DIAGNOSTICS
+
 // RUN: not %target-swift-frontend -swift-version 5 -typecheck -enable-experimental-feature Macros -load-plugin-library %t/%target-library-name(MacroDefinition) -I %swift-host-lib-dir -module-name MacroUser -DTEST_DIAGNOSTICS -serialize-diagnostics-path %t/macro_expand.dia %s -emit-macro-expansion-files no-diagnostics > %t/macro-printing.txt
 // RUN: c-index-test -read-diagnostics %t/macro_expand.dia 2>&1 | %FileCheck -check-prefix CHECK-DIAGS %s
 

Original file line number	Diff line number	Diff line change
`@@ -1413,6 +1413,10 @@ static bool ParseSearchPathArgs(SearchPathOptions &Opts,`
`1413`	`1413`	`}`
`1414`	`1414`	`Opts.setFrameworkSearchPaths(FrameworkSearchPaths);`
`1415`	`1415`
	`1416`	`+ for (const Arg *A : Args.filtered(OPT_plugin_path)) {`
	`1417`	`+ Opts.PluginSearchPaths.push_back(resolveSearchPath(A->getValue()));`
	`1418`	`+ }`
	`1419`	`+`
`1416`	`1420`	`for (const Arg *A : Args.filtered(OPT_L)) {`
`1417`	`1421`	`Opts.LibrarySearchPaths.push_back(resolveSearchPath(A->getValue()));`
`1418`	`1422`	`}`