[Macros] Add support for wasm macros

include/swift/AST/DiagnosticsFrontend.def

@@ -132,6 +132,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

@@ -213,17 +213,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:
@@ -234,6 +240,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>;
@@ -257,6 +265,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

@@ -2002,6 +2002,13 @@ def load_plugin_executable:
            "of module names where the macro types are declared">,
   MetaVarName<"<path>#<module-names>">;
 
+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/ModuleDependencies.cpp

@@ -584,6 +584,14 @@ void SwiftDependencyTracker::addCommonSearchPathDeps(
       recordFiles(val.SearchPath);
       break;
     }
+
+    // '-load-plugin <library path>:<server path>#<module name>,...'.
+    case PluginSearchOption::Kind::LoadPlugin: {
+      auto &val = entry.get<PluginSearchOption::LoadPlugin>();
+      FS->status(val.LibraryPath);
+      // TODO: what else do we need?
+      break;
+    }
     }
   }
 }

lib/AST/PluginLoader.cpp

@@ -122,6 +122,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

@@ -2027,6 +2027,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/Serialization/ModuleFileSharedCore.cpp

@@ -144,6 +144,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

@@ -59,7 +59,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// 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 =
-    875; // Add package field to SerializedKind_t
+    876; // Add PluginSearchOptionKind::LoadPlugin
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -688,6 +688,7 @@ enum class PluginSearchOptionKind : uint8_t {
   ExternalPluginPath,
   LoadPluginLibrary,
   LoadPluginExecutable,
+  LoadPlugin,
 };
 using PluginSearchOptionKindField = BCFixed<3>;
 

lib/Serialization/Serialization.cpp

@@ -1213,6 +1213,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;
+          }
           }
         }
       }

tools/lldb-moduleimport-test/lldb-moduleimport-test.cpp

@@ -105,6 +105,9 @@ static bool validateModule(
       case swift::PluginSearchOption::Kind::LoadPluginExecutable:
         optStr = "-load-plugin-executable";
         break;
+      case swift::PluginSearchOption::Kind::LoadPlugin:
+        optStr = "-load-plugin";
+        break;
       }
       llvm::outs() << "    " << optStr << " " << opt.second << "\n";
     }

tools/swift-ide-test/swift-ide-test.cpp

@@ -347,6 +347,11 @@ LoadPluginExecutable("load-plugin-executable",
                llvm::cl::desc("load plugin executable"),
                llvm::cl::cat(Category));
 
+static llvm::cl::list<std::string>
+LoadPlugin("load-plugin",
+           llvm::cl::desc("load plugin"),
+           llvm::cl::cat(Category));
+
 
 static llvm::cl::opt<bool>
 EnableSourceImport("enable-source-import", llvm::cl::Hidden,
@@ -4563,6 +4568,24 @@ int main(int argc, char *argv[]) {
                                                    std::move(moduleNames)});
     }
   }
+  if (!options::LoadPlugin.empty()) {
+    for (auto arg: options::LoadPlugin) {
+      StringRef pathAndServer;
+      StringRef modulesStr;
+      std::tie(pathAndServer, modulesStr) = StringRef(arg).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);
+      }
+      InitInvok.getSearchPathOptions().PluginSearchOpts.emplace_back(
+        PluginSearchOption::LoadPlugin{std::string(path),
+                                       std::string(server),
+                                       std::move(moduleNames)});
+    }
+  }
   for (auto path : options::PluginPath) {
     InitInvok.getSearchPathOptions().PluginSearchOpts.emplace_back(
         PluginSearchOption::PluginPath{path});

tools/swift-plugin-server/.gitignore

@@ -0,0 +1 @@
+Package.resolved

tools/swift-plugin-server/CMakeLists.txt

@@ -1,10 +1,21 @@
 if (SWIFT_BUILD_SWIFT_SYNTAX)
+  # override the remote SwiftSystem and ArgParser dependencies in WasmKit
+  FetchContent_Declare(SwiftSystem SOURCE_DIR "${PROJECT_SOURCE_DIR}/../swift-system")
+  FetchContent_Declare(ArgumentParser SOURCE_DIR "${PROJECT_SOURCE_DIR}/../swift-argument-parser")
+
+  FetchContent_Declare(WasmKit SOURCE_DIR "${PROJECT_SOURCE_DIR}/../wasmkit")
+  FetchContent_MakeAvailable(WasmKit)
+
   add_pure_swift_host_tool(swift-plugin-server
     Sources/swift-plugin-server/swift-plugin-server.swift
+    Sources/swift-plugin-server/WasmEngine.swift
+    Sources/swift-plugin-server/WasmKitEngine.swift
+    Sources/swift-plugin-server/WasmMessageHandler.swift
     SWIFT_COMPONENT
       compiler
     SWIFT_DEPENDENCIES
       SwiftCompilerPluginMessageHandling
       SwiftLibraryPluginProvider
+      WasmKitWASI
   )
 endif()

tools/swift-plugin-server/Package.swift

@@ -9,12 +9,16 @@ let package = Package(
   ],
   dependencies: [
     .package(path: "../../../swift-syntax"),
+    .package(path: "../../../wasmkit"),
   ],
   targets: [
     .executableTarget(
       name: "swift-plugin-server",
       dependencies: [
         .product(name: "SwiftCompilerPluginMessageHandling", package: "swift-syntax"),
+        .product(name: "SwiftLibraryPluginProvider", package: "swift-syntax"),
+        .product(name: "WASI", package: "WasmKit"),
+        .product(name: "WasmKitWASI", package: "WasmKit"),
       ]
     ),
   ],

tools/swift-plugin-server/Sources/swift-plugin-server/WasmEngine.swift

@@ -0,0 +1,91 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import WASI
+import WasmTypes
+import SystemPackage
+
+typealias WasmFunction = ([UInt32]) throws -> [UInt32]
+
+protocol WasmEngine {
+  init(path: FilePath, imports: WASIBridgeToHost) throws
+
+  func function(named name: String) throws -> WasmFunction?
+}
+
+typealias DefaultWasmPlugin = WasmEnginePlugin<DefaultWasmEngine>
+
+// a WasmPlugin implementation that delegates to a WasmEngine
+struct WasmEnginePlugin<Engine: WasmEngine>: WasmPlugin {
+  private let hostToPlugin: FileDescriptor
+  private let pluginToHost: FileDescriptor
+  private let pumpFunction: WasmFunction
+  let engine: Engine
+
+  init(path: FilePath) throws {
+    let hostToPluginPipes = try FileDescriptor.pipe()
+    let pluginToHostPipes = try FileDescriptor.pipe()
+    self.hostToPlugin = hostToPluginPipes.writeEnd
+    self.pluginToHost = pluginToHostPipes.readEnd
+
+    let bridge = try WASIBridgeToHost(
+      stdin: hostToPluginPipes.readEnd,
+      stdout: pluginToHostPipes.writeEnd,
+      stderr: .standardError
+    )
+    engine = try Engine(path: path, imports: bridge)
+
+    let exportName = "swift_wasm_macro_v1_pump"
+    guard let pump = try engine.function(named: exportName) else {
+      throw WasmEngineError(message: "Wasm plugin has an unknown ABI (could not find '\(exportName)')")
+    }
+    self.pumpFunction = pump
+
+    guard let start = try engine.function(named: "_start") else {
+      throw WasmEngineError(message: "Wasm plugin does not have a '_start' entrypoint")
+    }
+    _ = try start([])
+  }
+
+  func handleMessage(_ json: [UInt8]) throws -> [UInt8] {
+    try withUnsafeBytes(of: UInt64(json.count).littleEndian) {
+      _ = try hostToPlugin.writeAll($0)
+    }
+    try hostToPlugin.writeAll(json)
+
+    _ = try pumpFunction([])
+
+    let lengthRaw = try withUnsafeTemporaryAllocation(of: UInt8.self, capacity: 8) { buffer in
+      let lengthCount = try pluginToHost.read(into: UnsafeMutableRawBufferPointer(buffer))
+      guard lengthCount == 8 else {
+        throw WasmEngineError(message: "Wasm plugin sent invalid response")
+      }
+      return buffer.withMemoryRebound(to: UInt64.self, \.baseAddress!.pointee)
+    }
+    let length = Int(UInt64(littleEndian: lengthRaw))
+    return try [UInt8](unsafeUninitializedCapacity: length) { buffer, size in
+      let received = try pluginToHost.read(into: UnsafeMutableRawBufferPointer(buffer))
+      guard received == length else {
+        throw WasmEngineError(message: "Wasm plugin sent truncated response")
+      }
+      size = received
+    }
+  }
+}
+
+struct WasmEngineError: Error, CustomStringConvertible {
+  let description: String
+
+  init(message: String) {
+    self.description = message
+  }
+}