[Macros] Support user-defined macros as compiler plugins

Allow user-defined macros to be loaded from dynamic libraries and evaluated.

- Introduce a _CompilerPluginSupport module installed into the toolchain. Its `_CompilerPlugin` protocol acts as a stable interface between the compiler and user-defined macros.
- Introduce a `-compiler-plugin-library-path <path>` attribute which allows users to specify dynamic libraries to be loaded into the compiler.

A macro library must declare a public top-level computed property `public var allMacros: [Any.Type]` and be compiled to a dynamic library. The compiler will call the getter of this property to obtain and regsiter all macros.

Known issues:
- We current do not have a way to strip out unnecessary symbols from the plugin dylib, aka. produce a plugin library that does not contain SwiftSyntax symbols that will collide with the compiler itself.
- `MacroExpansionExpr`'s type is hard-coded as `(Int, String)`. It should instead be specified by the macro via protocol requirements such as `signature` and `genericSignature`. We need more protocol requirements in `_CompilerPlugin` to handle this.
- `dlopen` is not secure and is only for prototyping use here.

Author: rxwei

include/swift/AST/ASTContext.h

@@ -44,6 +44,7 @@
 #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>
@@ -132,6 +133,7 @@ namespace swift {
   class IndexSubset;
   struct SILAutoDiffDerivativeFunctionKey;
   struct InterfaceSubContextDelegate;
+  class CompilerPlugin;
 
   enum class KnownProtocolKind : uint8_t;
 
@@ -347,6 +349,12 @@ class ASTContext final {
       llvm::SmallPtrSet<DerivativeAttr *, 1>>
       DerivativeAttrs;
 
+  /// Cache of compiler plugins keyed by their name.
+  llvm::StringMap<CompilerPlugin> LoadedPlugins;
+
+  /// Cache of loaded symbols.
+  llvm::StringMap<void *> LoadedSymbols;
+
 private:
   /// The current generation number, which reflects the number of
   /// times that external modules have been loaded.
@@ -1435,6 +1443,14 @@ class ASTContext final {
   /// The declared interface type of Builtin.TheTupleType.
   BuiltinTupleType *getBuiltinTupleType();
 
+  /// Finds the loaded compiler plugin given its name.
+  CompilerPlugin *getLoadedPlugin(StringRef name);
+
+  /// Finds the address of the given symbol. If `libraryHint` is non-null,
+  /// search within the library.
+  void *getAddressOfSymbol(StringRef name,
+                           llvm::sys::DynamicLibrary *libraryHint = nullptr);
+
 private:
   friend Decl;
 
@@ -1447,6 +1463,8 @@ class ASTContext final {
   Optional<StringRef> getBriefComment(const Decl *D);
   void setBriefComment(const Decl *D, StringRef Comment);
 
+  void loadCompilerPlugins();
+
   friend TypeBase;
   friend ArchetypeType;
   friend OpaqueTypeDecl;

include/swift/AST/CompilerPlugin.h

@@ -0,0 +1,99 @@
+//===--- CompilerPlugin.h - Compiler Plugin Support -------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 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
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines supporting data structures for compiler plugins.
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/AST/Type.h"
+#include "swift/Basic/LLVM.h"
+#include "llvm/Support/DynamicLibrary.h"
+
+#ifndef SWIFT_AST_COMPILER_PLUGIN_H
+#define SWIFT_AST_COMPILER_PLUGIN_H
+
+namespace swift {
+
+class ASTContext;
+
+// Must be modified together with `_CompilerPluginKind` in
+// stdlib/toolchain/CompilerPluginSupport.swift.
+enum class CompilerPluginKind: uint32_t {
+  ExpressionMacro,
+};
+
+/// A compiler plugin that corresponds to a dynamically loaded Swift type that
+/// conforms to the `_CompilerPluginSupport._CompilerPlugin` protocol.
+class CompilerPlugin {
+  friend class ASTContext;
+
+  // Must be modified together with `_CompilerPlugin` in
+  // stdlib/toolchain/CompilerPluginSupport.swift.
+  enum class WitnessTableEntry: unsigned {
+    ConformanceDescriptor = 0,
+    // static func _name() -> (UnsafePointer<UInt8>, count: Int8)
+    Name = 1,
+    // static func _kind() -> _CompilerPluginKind
+    Kind = 2,
+    // static func _rewrite(...) -> (UnsafePointer<UInt8>?, count: Int)
+    Rewrite = 3,
+  };
+
+  /// The plugin type metadata.
+  const void *metadata;
+  /// The parent dynamic library containing the plugin.
+  llvm::sys::DynamicLibrary parentLibrary;
+  /// The witness table proving that the plugin conforms to `_CompilerPlugin`.
+  const void *witnessTable;
+  /// The plugin name, aka. result of the `_name()` method.
+  StringRef name;
+  /// The plugin's kind, aka. result of the `_kind()` method.
+  CompilerPluginKind kind;
+
+  template<typename Func>
+  const Func *getWitnessMethodUnsafe(WitnessTableEntry entry) const {
+    return reinterpret_cast<const Func *const *>(witnessTable)[(unsigned)entry];
+  }
+
+protected:
+  CompilerPlugin(const void *metadata, llvm::sys::DynamicLibrary parentLibrary,
+                 ASTContext &ctx);
+
+private:
+  /// Invoke the `_name` method. The caller assumes ownership of the result
+  /// string buffer.
+  StringRef invokeName() const;
+
+  /// Invoke the `_kind` method.
+  CompilerPluginKind invokeKind() const;
+
+public:
+  /// Invoke the `_rewrite` method. Invoke the `_name` method. The caller
+  /// assumes ownership of the result string buffer.
+  Optional<StringRef> invokeRewrite(StringRef targetModuleName,
+                                    StringRef filePath,
+                                    StringRef sourceFileText,
+                                    CharSourceRange range,
+                                    ASTContext &ctx) const;
+
+  StringRef getName() const {
+    return name;
+  }
+
+  CompilerPluginKind getKind() const {
+    return kind;
+  }
+};
+
+} // end namespace swift
+
+#endif // SWIFT_AST_COMPILER_PLUGIN_H

include/swift/AST/DiagnosticsFrontend.def

@@ -472,6 +472,12 @@ REMARK(warning_in_access_notes_file,none,
        "ignored invalid content in access notes file: %0",
        (StringRef))
 
+WARNING(compiler_plugin_not_loaded,none,
+        "compiler plugin not loaded: %0; loader error: %1", (StringRef, StringRef))
+WARNING(compiler_plugin_missing_macro_declaration,none,
+        "compiler plugin module '%0' (in %1) is missing a top-level computed "
+        "property 'public var %2: [Any.Type]' to declare all macros; "
+        "undeclared macros will be ignored", (StringRef, StringRef, StringRef))
 
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/AST/DiagnosticsSema.def

@@ -6696,8 +6696,14 @@ ERROR(experimental_no_metadata_feature_can_only_be_used_when_enabled,
 //------------------------------------------------------------------------------
 // MARK: Macros
 //------------------------------------------------------------------------------
+
 ERROR(expected_macro_expansion_expr,PointsToFirstBadToken,
       "expected macro expansion to produce an expression", ())
+ERROR(macro_expansion_missing_feature_flag,PointsToFirstBadToken,
+      "macro expansion requires experimental feature 'Macros'", ())
+ERROR(macro_undefined,PointsToFirstBadToken,
+      "macro '%0' is undefined; use `-load-plugin` to specify dynamic "
+      "libraries that contain this macro", (StringRef))
 
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/AST/KnownIdentifiers.def

@@ -333,6 +333,9 @@ IDENTIFIER(unsafe)
 // The singleton instance of TupleTypeDecl in the Builtin module
 IDENTIFIER(TheTupleType)
 
+// Macros / compiler plugins
+IDENTIFIER_(CompilerPluginSupport)
+
 #undef IDENTIFIER
 #undef IDENTIFIER_
 #undef IDENTIFIER_WITH_NAME

include/swift/AST/KnownProtocols.def

@@ -114,6 +114,9 @@ PROTOCOL(AsyncIteratorProtocol)
 
 PROTOCOL(FloatingPoint)
 
+// Compiler Plugins
+PROTOCOL_(CompilerPlugin)
+
 EXPRESSIBLE_BY_LITERAL_PROTOCOL(ExpressibleByArrayLiteral, "Array", false)
 EXPRESSIBLE_BY_LITERAL_PROTOCOL(ExpressibleByBooleanLiteral, "BooleanLiteralType", true)
 EXPRESSIBLE_BY_LITERAL_PROTOCOL(ExpressibleByDictionaryLiteral, "Dictionary", false)

include/swift/AST/SearchPathOptions.h

@@ -34,6 +34,7 @@ enum class ModuleSearchPathKind {
   Framework,
   DarwinImplicitFramework,
   RuntimeLibrary,
+  CompilerPlugin,
 };
 
 /// A single module search path that can come from different sources, e.g.
@@ -230,6 +231,9 @@ class SearchPathOptions {
   /// \c ModuleSearchPath.
   std::vector<std::string> DarwinImplicitFrameworkSearchPaths;
 
+  /// Compiler plugin library search paths.
+  std::vector<std::string> CompilerPluginLibraryPaths;
+
   /// Add a single import search path. Must only be called from
   /// \c ASTContext::addSearchPath.
   void addImportSearchPath(StringRef Path, llvm::vfs::FileSystem *FS) {
@@ -239,6 +243,14 @@ class SearchPathOptions {
                            ImportSearchPaths.size() - 1);
   }
 
+  void addCompilerPluginLibraryPath(StringRef Path, llvm::vfs::FileSystem *FS) {
+    CompilerPluginLibraryPaths.push_back(Path.str());
+    Lookup.searchPathAdded(FS, CompilerPluginLibraryPaths.back(),
+                           ModuleSearchPathKind::CompilerPlugin,
+                           /*isSystem=*/false,
+                           CompilerPluginLibraryPaths.size() - 1);
+  }
+
   /// Add a single framework search path. Must only be called from
   /// \c ASTContext::addSearchPath.
   void addFrameworkSearchPath(FrameworkSearchPath NewPath,
@@ -302,6 +314,16 @@ class SearchPathOptions {
     Lookup.searchPathsDidChange();
   }
 
+  void setCompilerPluginLibraryPaths(
+      std::vector<std::string> NewCompilerPluginLibraryPaths) {
+    CompilerPluginLibraryPaths = NewCompilerPluginLibraryPaths;
+    Lookup.searchPathsDidChange();
+  }
+
+  ArrayRef<std::string> getCompilerPluginLibraryPaths() const {
+    return CompilerPluginLibraryPaths;
+  }
+
   /// Path(s) to virtual filesystem overlay YAML files.
   std::vector<std::string> VFSOverlayFiles;
 

include/swift/Frontend/Frontend.h

@@ -192,6 +192,14 @@ class CompilerInvocation {
     return SearchPathOpts.getFrameworkSearchPaths();
   }
 
+  void setCompilerPluginLibraryPaths(const std::vector<std::string> &Paths) {
+    SearchPathOpts.setCompilerPluginLibraryPaths(Paths);
+  }
+
+  ArrayRef<std::string> getCompilerPluginLibraryPaths() {
+    return SearchPathOpts.getCompilerPluginLibraryPaths();
+  }
+
   void setExtraClangArgs(const std::vector<std::string> &Args) {
     ClangImporterOpts.ExtraArgs = Args;
   }

include/swift/Option/Options.td

@@ -1695,4 +1695,11 @@ def nostartfiles:
 
 // END ONLY SUPPORTED IN NEW DRIVER
 
+def load_plugin:
+  Separate<["-"], "load-plugin">,
+  Flags<[FrontendOption, DoesNotAffectIncrementalBuild, ArgumentIsPath]>,
+  HelpText<"Path to a dynamic library containing compiler plugins such as "
+           "macros">,
+  MetaVarName<"<path>">;
+
 include "FrontendOptions.td"

lib/AST/ASTContext.cpp

@@ -19,9 +19,11 @@
 #include "ForeignRepresentationInfo.h"
 #include "SubstitutionMapStorage.h"
 #include "swift/AST/ClangModuleLoader.h"
+#include "swift/AST/CompilerPlugin.h"
 #include "swift/AST/ConcreteDeclRef.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsSema.h"
+#include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/AST/DistributedDecl.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/ExtInfo.h"
@@ -64,6 +66,7 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Config/config.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Compiler.h"
@@ -96,7 +99,8 @@ llvm::StringRef swift::getProtocolName(KnownProtocolKind kind) {
 namespace {
   enum class SearchPathKind : uint8_t {
     Import = 1 << 0,
-    Framework = 1 << 1
+    Framework = 1 << 1,
+    CompilerPlugin = 1 << 2
   };
 } // end anonymous namespace
 
@@ -680,10 +684,14 @@ ASTContext::ASTContext(
     getImpl().SearchPathsSet[path] |= SearchPathKind::Import;
   for (const auto &framepath : SearchPathOpts.getFrameworkSearchPaths())
     getImpl().SearchPathsSet[framepath.Path] |= SearchPathKind::Framework;
+  for (StringRef path : SearchPathOpts.getCompilerPluginLibraryPaths())
+    getImpl().SearchPathsSet[path] |= SearchPathKind::CompilerPlugin;
 
   // Register any request-evaluator functions available at the AST layer.
   registerAccessRequestFunctions(evaluator);
   registerNameLookupRequestFunctions(evaluator);
+
+  loadCompilerPlugins();
 }
 
 ASTContext::~ASTContext() {
@@ -1118,6 +1126,9 @@ ProtocolDecl *ASTContext::getProtocol(KnownProtocolKind kind) const {
   case KnownProtocolKind::UnsafeCxxRandomAccessIterator:
     M = getLoadedModule(Id_Cxx);
     break;
+  case KnownProtocolKind::CompilerPlugin:
+    M = getLoadedModule(Id_CompilerPluginSupport);
+    break;
   default:
     M = getStdlibModule();
     break;
@@ -6032,3 +6043,24 @@ BuiltinTupleType *ASTContext::getBuiltinTupleType() {
 
   return result;
 }
+
+CompilerPlugin *ASTContext::getLoadedPlugin(StringRef name) {
+  auto lookup = LoadedPlugins.find(name);
+  if (lookup == LoadedPlugins.end())
+    return nullptr;
+  return &lookup->second;
+}
+
+void *ASTContext::getAddressOfSymbol(StringRef name,
+                                     llvm::sys::DynamicLibrary *libraryHint) {
+  auto lookup = LoadedSymbols.try_emplace(name, nullptr);
+  void *&address = lookup.first->getValue();
+  if (lookup.second) {
+    if (libraryHint && libraryHint->isValid())
+      address = libraryHint->getAddressOfSymbol(name.data());
+    else
+      address = llvm::sys::DynamicLibrary::
+          SearchForAddressOfSymbol(name.data());
+  }
+  return address;
+}

lib/AST/ASTWalker.cpp

@@ -1247,8 +1247,9 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   Expr *visitMacroExpansionExpr(MacroExpansionExpr *E) {
-    auto *macro = doIt(E->getMacro());
-    if (!macro) return nullptr;
+    // Don't want the callee expr yet as it will stay as a UDRE for now.
+    // auto *macro = doIt(E->getMacro());
+    // if (!macro) return nullptr;
     Expr *rewritten = nullptr;
     if (E->getRewritten()) {
       rewritten = doIt(E->getRewritten());
@@ -1259,7 +1260,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       args = doIt(E->getArgs());
       if (!args) return nullptr;
     }
-    E->setMacro(macro);
+    // E->setMacro(macro);
     E->setRewritten(rewritten);
     E->setArgs(args);
     return E;