[Macros] Type check user-defined macro plugins (#61861)

Type check user-defined macros plugins with user-provided type signatures.

Also, load plugin libraries with `RTLD_LOCAL` instead of `RTLD_GLOBAL` to prevent symbol collision between plugins. `llvm::sys::DynamicLibrary` only supports `RTLD_GLOBAL` so we use the plain `dlopen` instead. This does not work on Windows and needs to be fixed.

Friend PR: swiftlang/swift-syntax#1042

Author: rxwei

include/swift/AST/ASTContext.h

@@ -44,7 +44,6 @@
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/DataTypes.h"
-#include "llvm/Support/DynamicLibrary.h"
 #include <functional>
 #include <memory>
 #include <utility>
@@ -1456,10 +1455,9 @@ class ASTContext final {
   /// Finds the loaded compiler plugin given its name.
   CompilerPlugin *getLoadedPlugin(StringRef name);
 
-  /// Finds the address of the given symbol. If `libraryHint` is non-null,
+  /// Finds the address of the given symbol. If `libraryHandleHint` is non-null,
   /// search within the library.
-  void *getAddressOfSymbol(StringRef name,
-                           llvm::sys::DynamicLibrary *libraryHint = nullptr);
+  void *getAddressOfSymbol(const char *name, void *libraryHandleHint = nullptr);
 
 private:
   friend Decl;

include/swift/AST/CompilerPlugin.h

@@ -17,7 +17,6 @@
 #include "swift/AST/Type.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/StringExtras.h"
-#include "llvm/Support/DynamicLibrary.h"
 
 #ifndef SWIFT_AST_COMPILER_PLUGIN_H
 #define SWIFT_AST_COMPILER_PLUGIN_H
@@ -49,12 +48,16 @@ class CompilerPlugin {
     Kind = 2,
     // static func _rewrite(...) -> (UnsafePointer<UInt8>?, count: Int)
     Rewrite = 3,
+    // static func _genericSignature(...) -> (UnsafePointer<UInt8>?, count: Int)
+    GenericSignature = 4,
+    // static func _typeSignature(...) -> (UnsafePointer<UInt8>, count: Int)
+    TypeSignature = 5,
   };
 
   /// The plugin type metadata.
   const void *metadata;
   /// The parent dynamic library containing the plugin.
-  llvm::sys::DynamicLibrary parentLibrary;
+  void *parentLibrary;
   /// The witness table proving that the plugin conforms to `_CompilerPlugin`.
   const void *witnessTable;
   /// The plugin name, aka. result of the `_name()` method.
@@ -68,8 +71,7 @@ class CompilerPlugin {
   }
 
 protected:
-  CompilerPlugin(const void *metadata, llvm::sys::DynamicLibrary parentLibrary,
-                 ASTContext &ctx);
+  CompilerPlugin(const void *metadata, void *parentLibrary, ASTContext &ctx);
 
 private:
   /// Invoke the `_name` method. The caller assumes ownership of the result
@@ -80,12 +82,24 @@ class CompilerPlugin {
   Kind invokeKind() const;
 
 public:
-  /// Invoke the `_rewrite` method. Invoke the `_name` method. The caller
-  /// assumes ownership of the result string buffer.
+  ~CompilerPlugin();
+  CompilerPlugin(const CompilerPlugin &) = delete;
+  CompilerPlugin(CompilerPlugin &&) = default;
+
+  /// Invoke the `_rewrite` method. The caller assumes ownership of the result
+  /// string buffer.
   Optional<NullTerminatedStringRef> invokeRewrite(
       StringRef targetModuleName, StringRef filePath, StringRef sourceFileText,
       CharSourceRange range,  ASTContext &ctx) const;
 
+  /// Invoke the `_genericSignature` method. The caller assumes ownership of the
+  /// result string buffer.
+  Optional<StringRef> invokeGenericSignature() const;
+
+  /// Invoke the `_typeSignature` method. The caller assumes ownership of the
+  /// result string buffer.
+  StringRef invokeTypeSignature() const;
+
   StringRef getName() const {
     return name;
   }

include/swift/AST/DiagnosticsSema.def

@@ -6706,8 +6706,8 @@ ERROR(experimental_no_metadata_feature_can_only_be_used_when_enabled,
 ERROR(expected_macro_expansion_expr,PointsToFirstBadToken,
       "expected macro expansion to produce an expression", ())
 ERROR(macro_undefined,PointsToFirstBadToken,
-      "macro '%0' is undefined; use `-load-plugin-library` to specify dynamic "
-      "libraries that contain this macro", (StringRef))
+      "macro %0 is undefined; use `-load-plugin-library` to specify dynamic "
+      "libraries that contain this macro", (Identifier))
 
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

lib/AST/ASTContext.cpp

@@ -73,6 +73,10 @@
 #include <queue>
 #include <memory>
 
+#if !defined(_WIN32)
+#include <dlfcn.h>
+#endif
+
 #include "RequirementMachine/RewriteContext.h"
 
 using namespace swift;
@@ -6055,16 +6059,15 @@ CompilerPlugin *ASTContext::getLoadedPlugin(StringRef name) {
   return &lookup->second;
 }
 
-void *ASTContext::getAddressOfSymbol(StringRef name,
-                                     llvm::sys::DynamicLibrary *libraryHint) {
+void *ASTContext::getAddressOfSymbol(const char *name,
+                                     void *libraryHandleHint) {
   auto lookup = LoadedSymbols.try_emplace(name, nullptr);
   void *&address = lookup.first->getValue();
+#if !defined(_WIN32)
   if (lookup.second) {
-    if (libraryHint && libraryHint->isValid())
-      address = libraryHint->getAddressOfSymbol(name.data());
-    else
-      address = llvm::sys::DynamicLibrary::
-          SearchForAddressOfSymbol(name.data());
+    auto *handle = libraryHandleHint ? libraryHandleHint : RTLD_DEFAULT;
+    address = dlsym(handle, name);
   }
+#endif
   return address;
 }

lib/AST/CompilerPlugin.cpp

@@ -24,9 +24,12 @@
 #include "swift/Parse/Lexer.h"
 #include "swift/Subsystems.h"
 #include "llvm/Config/config.h"
-#include "llvm/Support/DynamicLibrary.h"
 #include <cstdlib>
 
+#if !defined(_WIN32)
+#include <dlfcn.h>
+#endif
+
 using namespace swift;
 
 #define ALLMACROS_PROPERTY_NAME "allMacros"
@@ -56,8 +59,7 @@ void swift_ASTGen_getMacroTypes(const void *getter,
 #endif
 
 static const void *
-getMacroRegistrationPropertyGetter(llvm::sys::DynamicLibrary library,
-                                   StringRef moduleName,
+getMacroRegistrationPropertyGetter(void *library, StringRef moduleName,
                                    ASTContext &ctx) {
   assert(!moduleName.empty());
   // TODO: Consider using runtime lookup to get all types that conform to the
@@ -147,23 +149,28 @@ getMacroRegistrationPropertyGetter(llvm::sys::DynamicLibrary library,
     assert(mangleResult.isSuccess());
     name = mangleResult.result();
   }
-  return ctx.getAddressOfSymbol(name, &library);
+  return ctx.getAddressOfSymbol(name.c_str(), library);
 }
 
 void ASTContext::loadCompilerPlugins() {
-  for (StringRef path : SearchPathOpts.getCompilerPluginLibraryPaths()) {
-    std::string errorMsg;
-    auto lib = llvm::sys::DynamicLibrary::getPermanentLibrary(
-        path.data(), &errorMsg);
-    if (!lib.isValid()) {
+  for (auto &path : SearchPathOpts.getCompilerPluginLibraryPaths()) {
+    void *lib = nullptr;
+#if !defined(_WIN32)
+    lib = dlopen(path.c_str(), RTLD_LAZY|RTLD_LOCAL);
+#endif
+    if (!lib) {
+      const char *errorMsg = "Unsupported platform";
+#if !defined(_WIN32)
+      errorMsg = dlerror();
+#endif
       Diags.diagnose(SourceLoc(), diag::compiler_plugin_not_loaded, path,
                      errorMsg);
       continue;
     }
     auto moduleName = llvm::sys::path::filename(path);
-    #if !defined(_WIN32)
+#if !defined(_WIN32)
     moduleName.consume_front("lib");
-    #endif
+#endif
     moduleName.consume_back(LTDL_SHLIB_EXT);
     auto *getter = getMacroRegistrationPropertyGetter(lib, moduleName, *this);
     if (!getter) {
@@ -175,7 +182,7 @@ void ASTContext::loadCompilerPlugins() {
     // Note: We don't currently have a way to poke at the contents of a Swift
     // array `[Any.Type]` from C++. But this should not be an issue for release
     // toolchains where user-defined macros will be used.
-    #if SWIFT_SWIFT_PARSER
+#if SWIFT_SWIFT_PARSER
     const void *const *metatypesAddress;
     ptrdiff_t metatypeCount;
     swift_ASTGen_getMacroTypes(getter, &metatypesAddress, &metatypeCount);
@@ -186,7 +193,7 @@ void ASTContext::loadCompilerPlugins() {
       LoadedPlugins.try_emplace(name, std::move(plugin));
     }
     free(const_cast<void *>((const void *)metatypes.data()));
-    #endif
+#endif // SWIFT_SWIFT_PARSER
   }
 }
 
@@ -197,8 +204,7 @@ using WitnessTableLookupFn = const void *(const void *type,
 extern "C" WitnessTableLookupFn swift_conformsToProtocol;
 #endif
 
-CompilerPlugin::CompilerPlugin(const void *metadata,
-                               llvm::sys::DynamicLibrary parentLibrary,
+CompilerPlugin::CompilerPlugin(const void *metadata, void *parentLibrary,
                                ASTContext &ctx)
    : metadata(metadata), parentLibrary(parentLibrary)
 {
@@ -208,6 +214,8 @@ CompilerPlugin::CompilerPlugin(const void *metadata,
 #endif
   void *protocolDescriptor =
       ctx.getAddressOfSymbol(COMPILER_PLUGIN_PROTOCOL_DESCRIPTOR);
+  assert(swift_conformsToProtocol);
+  assert(protocolDescriptor);
   witnessTable = swift_conformsToProtocol(metadata, protocolDescriptor);
   assert(witnessTable && "Type does not conform to _CompilerPlugin");
   auto returnedName = invokeName();
@@ -216,6 +224,12 @@ CompilerPlugin::CompilerPlugin(const void *metadata,
   kind = invokeKind();
 }
 
+CompilerPlugin::~CompilerPlugin() {
+#if !defined(_WIN32)
+  dlclose(parentLibrary);
+#endif
+}
+
 namespace {
 struct CharBuffer {
   const char *data;
@@ -280,3 +294,28 @@ CompilerPlugin::invokeRewrite(StringRef targetModuleName,
   llvm_unreachable("Incompatible host compiler");
 #endif
 }
+
+Optional<StringRef>
+CompilerPlugin::invokeGenericSignature() const {
+#if __clang__
+  using Method = SWIFT_CC CharBuffer(
+      SWIFT_CONTEXT const void *, const void *, const void *);
+  auto method = getWitnessMethodUnsafe<Method>(
+      WitnessTableEntry::GenericSignature);
+  return method(metadata, metadata, witnessTable).str();
+#else
+  llvm_unreachable("Incompatible host compiler");
+#endif
+}
+
+StringRef CompilerPlugin::invokeTypeSignature() const {
+#if __clang__
+  using Method = SWIFT_CC CharBuffer(
+      SWIFT_CONTEXT const void *, const void *, const void *);
+  auto method = getWitnessMethodUnsafe<Method>(
+      WitnessTableEntry::TypeSignature);
+  return method(metadata, metadata, witnessTable).str();
+#else
+  llvm_unreachable("Incompatible host compiler");
+#endif
+}

lib/CompilerPluginSupport/CompilerPluginSupport.swift

@@ -57,4 +57,16 @@ public protocol _CompilerPlugin {
     localSourceText: UnsafePointer<UInt8>,
     localSourceTextCount: Int
   ) -> (UnsafePointer<UInt8>?, count: Int)
+
+  /// Returns the generic signature of the plugin.
+  ///
+  /// - Returns: A newly allocated buffer containing the generic signature.
+  ///   The caller is responsible for managing the memory.
+  static func _genericSignature() -> (UnsafePointer<UInt8>?, count: Int)
+
+  /// Returns the type signature of the plugin.
+  ///
+  /// - Returns: A newly allocated buffer containing the type signature. The
+  ///   caller is responsible for managing the memory.
+  static func _typeSignature() -> (UnsafePointer<UInt8>, count: Int)
 }

lib/Sema/CSGen.cpp

@@ -3635,10 +3635,10 @@ namespace {
         auto macroIdent = expr->getMacroName().getBaseIdentifier();
         auto refType = CS.getTypeOfMacroReference(macroIdent.str(), expr);
         if (!refType) {
-          // FIXME: This is currently hard-coded to (Int, String) just for
-          // testing Stringify, before we can parse a type signature from the
-          // macro plugin.
-          return TupleType::get({ctx.getIntType(), ctx.getStringType()}, ctx);
+          ctx.Diags.diagnose(expr->getMacroNameLoc(), diag::macro_undefined,
+                             macroIdent)
+              .highlight(expr->getMacroNameLoc().getSourceRange());
+          return Type();
         }
         if (expr->getArgs()) {
           CS.associateArgumentList(CS.getConstraintLocator(expr), expr->getArgs());