[Macros] Add support for wasm macros

include/swift/AST/DiagnosticsFrontend.def

@@ -127,6 +127,9 @@ ERROR(error_no_source_location_scope_map,none,
 ERROR(error_load_plugin_executable,none,
       "invalid value '%0' in '-load-plugin-executable'; "
       "make sure to use format '<plugin path>#<module names>'", (StringRef))
+ERROR(error_load_plugin,none,
+      "invalid value '%0' in '-load-plugin'; "
+      "make sure to use format '<library path>#<server path>#<module names>'", (StringRef))
 
 NOTE(note_valid_swift_versions, none,
      "valid arguments to '-swift-version' are %0", (StringRef))

include/swift/AST/SearchPathOptions.h

@@ -214,17 +214,23 @@ class PluginSearchOption {
     std::string SearchPath;
     std::string ServerPath;
   };
+  struct LoadPlugin {
+    std::string LibraryPath;
+    std::string ServerPath;
+    std::vector<std::string> ModuleNames;
+  };
 
   enum class Kind : uint8_t {
     LoadPluginLibrary,
     LoadPluginExecutable,
     PluginPath,
     ExternalPluginPath,
+    LoadPlugin,
   };
 
 private:
   using Members = ExternalUnionMembers<LoadPluginLibrary, LoadPluginExecutable,
-                                       PluginPath, ExternalPluginPath>;
+                                       PluginPath, ExternalPluginPath, LoadPlugin>;
   static Members::Index getIndexForKind(Kind kind) {
     switch (kind) {
     case Kind::LoadPluginLibrary:
@@ -235,6 +241,8 @@ class PluginSearchOption {
       return Members::indexOf<PluginPath>();
     case Kind::ExternalPluginPath:
       return Members::indexOf<ExternalPluginPath>();
+    case Kind::LoadPlugin:
+      return Members::indexOf<LoadPlugin>();
     }
   };
   using Storage = ExternalUnion<Kind, Members, getIndexForKind>;
@@ -258,6 +266,9 @@ class PluginSearchOption {
       : kind(Kind::ExternalPluginPath) {
     storage.emplace<ExternalPluginPath>(kind, v);
   }
+  PluginSearchOption(const LoadPlugin &v) : kind(Kind::LoadPlugin) {
+    storage.emplace<LoadPlugin>(kind, v);
+  }
   PluginSearchOption(const PluginSearchOption &o) : kind(o.kind) {
     storage.copyConstruct(o.kind, o.storage);
   }

include/swift/Option/Options.td

@@ -2071,6 +2071,13 @@ def in_process_plugin_server_path : Separate<["-"], "in-process-plugin-server-pa
   Flags<[FrontendOption, ArgumentIsPath]>,
   HelpText<"Path to dynamic library plugin server">;
 
+def load_plugin:
+  Separate<["-"], "load-plugin">, Group<plugin_search_Group>,
+  Flags<[FrontendOption, DoesNotAffectIncrementalBuild]>,
+  HelpText<"Path to a plugin library, a server to load it in, and a comma-separated "
+           "list of module names where the macro types are declared">,
+  MetaVarName<"<library-path>#<server-path>#<module-names>">;
+
 def disable_sandbox:
   Flag<["-"], "disable-sandbox">,
   Flags<[FrontendOption, DoesNotAffectIncrementalBuild]>,

lib/AST/PluginLoader.cpp

@@ -123,6 +123,14 @@ PluginLoader::getPluginMap() {
       continue;
     }
 
+    case PluginSearchOption::Kind::LoadPlugin: {
+      auto &val = entry.get<PluginSearchOption::LoadPlugin>();
+      for (auto &moduleName : val.ModuleNames) {
+        try_emplace(moduleName, val.LibraryPath, val.ServerPath);
+      }
+      continue;
+    }
+
     // '-plugin-path <library search path>'.
     case PluginSearchOption::Kind::PluginPath: {
       auto &val = entry.get<PluginSearchOption::PluginPath>();

lib/Basic/Sandbox.cpp

@@ -24,8 +24,8 @@ static StringRef sandboxProfile(llvm::BumpPtrAllocator &Alloc) {
   // Allow reading file metadata of any files.
   contents += "(allow file-read-metadata)\n";
 
-  // Allow reading dylibs.
-  contents += "(allow file-read* (regex #\"\\.dylib$\"))\n";
+  // Allow reading dylibs and WebAssembly macros.
+  contents += "(allow file-read* (regex #\"\\.(dylib|wasm)$\"))\n";
 
   // This is required to launch any processes (execve(2)).
   contents += "(allow process-exec*)\n";

lib/Frontend/CompilerInvocation.cpp

@@ -2154,6 +2154,30 @@ static bool ParseSearchPathArgs(SearchPathOptions &Opts, ArgList &Args,
       }
       break;
     }
+    case OPT_load_plugin: {
+      // '<path to library>#<path to server>#<module names>' where the module names are
+      // comma separated.
+      StringRef pathAndServer;
+      StringRef modulesStr;
+      std::tie(pathAndServer, modulesStr) = StringRef(A->getValue()).rsplit('#');
+      StringRef path;
+      StringRef server;
+      std::tie(path, server) = pathAndServer.rsplit('#');
+      std::vector<std::string> moduleNames;
+      for (auto name : llvm::split(modulesStr, ',')) {
+        moduleNames.emplace_back(name);
+      }
+      if (path.empty() || server.empty() || moduleNames.empty()) {
+        Diags.diagnose(SourceLoc(), diag::error_load_plugin,
+                       A->getValue());
+      } else {
+        Opts.PluginSearchOpts.emplace_back(
+            PluginSearchOption::LoadPlugin{resolveSearchPath(path),
+                                           resolveSearchPath(server),
+                                           std::move(moduleNames)});
+      }
+      break;
+    }
     case OPT_plugin_path: {
       Opts.PluginSearchOpts.emplace_back(
           PluginSearchOption::PluginPath{resolveSearchPath(A->getValue())});

lib/Frontend/ModuleInterfaceLoader.cpp

@@ -1728,6 +1728,15 @@ void InterfaceSubContextDelegateImpl::inheritOptionsForBuildingInterface(
       }
       break;
     }
+    case PluginSearchOption::Kind::LoadPlugin: {
+      auto &val = entry.get<PluginSearchOption::LoadPlugin>();
+      for (auto &moduleName : val.ModuleNames) {
+        GenericArgs.push_back("-load-plugin");
+        GenericArgs.push_back(
+            ArgSaver.save(val.LibraryPath + "#" + val.ServerPath + "#" + moduleName));
+      }
+      break;
+    }
     case PluginSearchOption::Kind::PluginPath: {
       auto &val = entry.get<PluginSearchOption::PluginPath>();
       GenericArgs.push_back("-plugin-path");

lib/Serialization/ModuleFileSharedCore.cpp

@@ -147,6 +147,9 @@ static bool readOptionsBlock(llvm::BitstreamCursor &cursor,
       case PluginSearchOptionKind::LoadPluginExecutable:
         optKind = PluginSearchOption::Kind::LoadPluginExecutable;
         break;
+      case PluginSearchOptionKind::LoadPlugin:
+        optKind = PluginSearchOption::Kind::LoadPlugin;
+        break;
       }
       extendedInfo.addPluginSearchOption({optKind, blobData});
       break;

lib/Serialization/ModuleFormat.h

@@ -58,7 +58,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 896; // @yield_once_2 produces address
+const uint16_t SWIFTMODULE_VERSION_MINOR = 897; // Add PluginSearchOptionKind::LoadPlugin
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -694,6 +694,7 @@ enum class PluginSearchOptionKind : uint8_t {
   ExternalPluginPath,
   LoadPluginLibrary,
   LoadPluginExecutable,
+  LoadPlugin,
 };
 using PluginSearchOptionKindField = BCFixed<3>;
 

lib/Serialization/Serialization.cpp

@@ -1237,6 +1237,17 @@ void Serializer::writeHeader() {
                 uint8_t(PluginSearchOptionKind::LoadPluginExecutable), optStr);
             continue;
           }
+          case PluginSearchOption::Kind::LoadPlugin: {
+            auto &opt = elem.get<PluginSearchOption::LoadPlugin>();
+            std::string optStr = opt.LibraryPath + "#" + opt.ServerPath + "#";
+            llvm::interleave(
+                opt.ModuleNames, [&](auto &name) { optStr += name; },
+                [&]() { optStr += ","; });
+            PluginSearchOpt.emit(
+                ScratchRecord,
+                uint8_t(PluginSearchOptionKind::LoadPlugin), optStr);
+            continue;
+          }
           }
         }
       }

test/Macros/Inputs/wasi_shim.h

@@ -0,0 +1,55 @@
+// definitions from
+// https://github.com/WebAssembly/wasi-libc/blob/320bbbcced68ce8e564b0dc4c8f80a5a5ad21a9c/libc-bottom-half/headers/public/wasi/api.h
+
+typedef unsigned char uint8_t;
+typedef unsigned short uint16_t;
+typedef unsigned int uint32_t;
+typedef unsigned long long uint64_t;
+
+typedef uint8_t wasi_signal_t;
+typedef uint16_t wasi_errno_t;
+typedef uint32_t wasi_exitcode_t;
+typedef int wasi_fd_t;
+typedef __SIZE_TYPE__ wasi_size_t;
+
+typedef struct wasi_ciovec_t {
+  const uint8_t *buf;
+  wasi_size_t buf_len;
+} wasi_ciovec_t;
+
+wasi_errno_t wasi_fd_write(
+  wasi_fd_t fd,
+  const wasi_ciovec_t *iovs,
+  wasi_size_t iovs_len,
+  wasi_size_t *nwritten
+) __attribute__((
+  __import_module__("wasi_snapshot_preview1"),
+  __import_name__("fd_write"),
+  __warn_unused_result__
+));
+
+_Noreturn void wasi_proc_exit(
+  wasi_exitcode_t code
+) __attribute__((
+  __import_module__("wasi_snapshot_preview1"),
+  __import_name__("proc_exit")
+));
+
+// libc shims
+
+static inline wasi_size_t swift_strlen(const char *buf) {
+  wasi_size_t len = 0;
+  while (buf[len]) len++;
+  return len;
+}
+
+static inline wasi_errno_t swift_write(int fd, const void *buf, wasi_size_t len) {
+  struct wasi_ciovec_t vec = { .buf = (const uint8_t *)buf, .buf_len = len };
+  wasi_size_t nwritten = 0;
+  return wasi_fd_write(fd, &vec, 1, &nwritten);
+}
+
+_Noreturn static inline void swift_abort(const char *message) {
+  swift_write(2, message, swift_strlen(message));
+  wasi_proc_exit(1);
+}

test/Macros/lit.local.cfg

@@ -31,3 +31,11 @@ else:
     )
     config.substitutions.append(('%c-flags', config.c_flags))
     config.substitutions.append(('%exe-linker-flags', config.exe_linker_flags))
+
+config.substitutions.insert(
+    0,
+    (
+        '%swift-build-wasm-c-plugin',
+        f'{config.clang} -target wasm32 -nostdlib -iquote %S'
+    )
+)

test/Macros/macro_plugin_wasm.swift

@@ -0,0 +1,60 @@
+// REQUIRES: swift_swift_parser
+// REQUIRES: CODEGENERATOR=WebAssembly
+
+// RUN: %empty-directory(%t)
+
+// RUN: split-file %s %t
+
+//#-- Prepare the Wasm macro plugin.
+// RUN: %swift-build-wasm-c-plugin %t/MacroDefinition.c -o %t/Plugin.wasm
+
+// RUN: env SWIFT_DUMP_PLUGIN_MESSAGING=1 %target-swift-frontend \
+// RUN:   -typecheck \
+// RUN:   -swift-version 5 \
+// RUN:   -load-plugin %t/Plugin.wasm#%swift-plugin-server#MacroDefinition \
+// RUN:   -Rmacro-loading \
+// RUN:   -module-name MyApp \
+// RUN:   %t/test.swift \
+// RUN:   > %t/macro-loading.txt 2>&1
+
+// RUN: %FileCheck %s < %t/macro-loading.txt
+
+// CHECK: ->(plugin:[[#PID:]]) {"getCapability":{"capability":{"protocolVersion":[[#PROTOCOL_VERSION:]]}}}
+// CHECK: <-(plugin:[[#PID]]) {"getCapabilityResult":{"capability":{"features":["load-plugin-library"],"protocolVersion":7}}}
+// CHECK: ->(plugin:[[#PID]]) {"loadPluginLibrary":{"libraryPath":"{{.*[\\/]}}Plugin.wasm","moduleName":"MacroDefinition"}}
+// CHECK: <-(plugin:[[#PID]]) {"loadPluginLibraryResult":{"diagnostics":[],"loaded":true}}
+// CHECK: ->(plugin:[[#PID]]) {"expandFreestandingMacro":{"discriminator":"$s{{.+}}","lexicalContext":[{{.*}}],"macro":{"moduleName":"MacroDefinition","name":"constInt","typeName":"ConstMacro"},"macroRole":"expression","syntax":{"kind":"expression","location":{"column":16,"fileID":"MyApp/test.swift","fileName":"{{.+}}test.swift","line":4,"offset":143},"source":"#constInt"}}}
+// CHECK: <-(plugin:[[#PID]]) {"expandMacroResult":{"diagnostics":[],"expandedSource":"1"}}
+
+//--- test.swift
+@freestanding(expression) macro constInt() -> Int = #externalMacro(module: "MacroDefinition", type: "ConstMacro")
+
+func foo() {
+  let _: Int = #constInt
+}
+
+//--- MacroDefinition.c
+#include "Inputs/wasi_shim.h"
+
+static void write_json(const char *json) {
+  wasi_size_t len = swift_strlen(json);
+  uint64_t len64 = (uint64_t)len;
+  swift_write(1, &len64, sizeof(len64));
+  swift_write(1, json, len);
+}
+
+int was_start_called = 0;
+int pump_calls = 0;
+
+__attribute__((export_name("_start")))
+void _start(void) {
+  if (was_start_called) swift_abort("_start called twice!");
+  was_start_called = 1;
+}
+
+__attribute__((export_name("swift_wasm_macro_v1_pump")))
+void pump(void) {
+  if (!was_start_called) swift_abort("_start not called!");
+  if (pump_calls++ != 0) swift_abort("expected pump to be called once");
+  write_json("{\"expandMacroResult\": {\"expandedSource\": \"1\", \"diagnostics\": []}}");
+}

test/Macros/macro_plugin_wasm_badversion.swift

@@ -0,0 +1,33 @@
+// REQUIRES: swift_swift_parser
+// REQUIRES: CODEGENERATOR=WebAssembly
+
+// RUN: %empty-directory(%t)
+
+// RUN: split-file %s %t
+
+//#-- Prepare the Wasm macro plugin.
+// RUN: %swift-build-wasm-c-plugin %t/MacroDefinition.c -o %t/Plugin.wasm
+
+// RUN: %target-swift-frontend \
+// RUN:   -typecheck -verify \
+// RUN:   -swift-version 5 \
+// RUN:   -load-plugin %t/Plugin.wasm#%swift-plugin-server#MacroDefinition \
+// RUN:   -module-name MyApp \
+// RUN:   %t/test.swift
+
+//--- test.swift
+// expected-warning @+2 {{Wasm plugin has an unknown ABI (could not find 'swift_wasm_macro_v1_pump')}}
+// expected-note @+1 {{declared here}}
+@freestanding(expression) macro constInt() -> Int = #externalMacro(module: "MacroDefinition", type: "ConstMacro")
+
+func foo() {
+  // expected-error @+1 {{Wasm plugin has an unknown ABI (could not find 'swift_wasm_macro_v1_pump')}}
+  let _: Int = #constInt
+}
+
+//--- MacroDefinition.c
+__attribute__((export_name("_start")))
+void _start(void) {}
+
+__attribute__((export_name("swift_wasm_macro_v100_pump")))
+void pump(void) {}

test/Macros/macro_plugin_wasm_guest_error.swift

@@ -0,0 +1,40 @@
+// REQUIRES: swift_swift_parser
+// REQUIRES: CODEGENERATOR=WebAssembly
+
+// RUN: %empty-directory(%t)
+
+// RUN: split-file %s %t
+
+//#-- Prepare the Wasm macro plugin.
+// RUN: %swift-build-wasm-c-plugin %t/MacroDefinition.c -o %t/Plugin.wasm
+
+// RUN: %target-swift-frontend \
+// RUN:   -typecheck -verify \
+// RUN:   -swift-version 5 \
+// RUN:   -load-plugin %t/Plugin.wasm#%swift-plugin-server#MacroDefinition \
+// RUN:   -module-name MyApp \
+// RUN:   %t/test.swift \
+// RUN:   2>%t/macro-loading.txt
+
+// RUN: %FileCheck %s < %t/macro-loading.txt
+
+// CHECK: guest error!
+
+//--- test.swift
+@freestanding(expression) macro constInt() -> Int = #externalMacro(module: "MacroDefinition", type: "ConstMacro")
+
+func foo() {
+  // expected-error @+1 {{failed to communicate with external macro}}
+  let _: Int = #constInt
+}
+
+//--- MacroDefinition.c
+#include "Inputs/wasi_shim.h"
+
+__attribute__((export_name("_start")))
+void _start(void) {}
+
+__attribute__((export_name("swift_wasm_macro_v1_pump")))
+void pump(void) {
+  swift_abort("guest error!\n");
+}