Implement wasm proc-macros using wasmi

README.md

@@ -1,3 +1,11 @@
+```
+WASI_SDK_PATH=../wasi-sdk-33.0-x86_64-linux ./x.py build library --target wasm32-wasip1,x86_64-unknown-linux-gnu
+echo 'extern crate proc_macro; #[proc_macro] pub fn foo(a: proc_macro::TokenStream) -> proc_macro::TokenStream { println!("foo"); (a.to_string() + ";println!(\"Hello from wasm proc macro!\");").parse().unwrap() }' | rustc +stage1 - --crate-type proc-macro --target wasm32-wasip1 -Zwasm-proc-macros
+echo 'fn main() { foo::foo!(println!("Hello World")); }' | rustc +stage1 --extern foo=rust_out.wasm - --edition 2024 && ./rust_out
+WASI_SDK_PATH=../wasi-sdk-33.0-x86_64-linux ./x.py test tests/ui/proc-macro --set rust.wasm-proc-macros=true
+```
+
+
 <div align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/rust-lang/www.rust-lang.org/master/static/images/rust-social-wide-dark.svg">

compiler/rustc_builtin_macros/src/lib.rs

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use rustc_data_structures::sync::IntoDynSyncSend;
 use rustc_expand::base::{MacroExpanderFn, ResolverExpand, SyntaxExtensionKind};
 use rustc_expand::proc_macro::BangProcMacro;
 use rustc_span::sym;
@@ -145,8 +146,12 @@ pub fn register_builtin_macros(resolver: &mut dyn ResolverExpand) {
         From: from::expand_deriving_from,
     }
 
-    let client = rustc_proc_macro::bridge::client::Client::expand1(rustc_proc_macro::quote);
-    register(sym::quote, SyntaxExtensionKind::Bang(Arc::new(BangProcMacro { client })));
+    let client = rustc_proc_macro::bridge::client::Client::expand1(rustc_proc_macro::quote)
+        .into_dyn_client();
+    register(
+        sym::quote,
+        SyntaxExtensionKind::Bang(Arc::new(BangProcMacro { client: IntoDynSyncSend(client) })),
+    );
     let requires = SyntaxExtensionKind::Attr(Arc::new(contracts::ExpandRequires));
     register(sym::contracts_requires, requires);
     let ensures = SyntaxExtensionKind::Attr(Arc::new(contracts::ExpandEnsures));

compiler/rustc_codegen_ssa/src/back/linker.rs

@@ -1738,7 +1738,38 @@ fn exported_symbols_for_proc_macro_crate(tcx: TyCtxt<'_>) -> Vec<SymbolExport> {
     let stable_crate_id = tcx.stable_crate_id(LOCAL_CRATE);
     let proc_macro_decls_name = rustc_session::generate_proc_macro_decls_symbol(stable_crate_id);
 
-    vec![symbol_export_from_raw_name(tcx, proc_macro_decls_name, SymbolExportKind::Data)]
+    if tcx.sess.target.is_like_wasm {
+        vec![
+            symbol_export_from_raw_name(tcx, proc_macro_decls_name, SymbolExportKind::Data),
+            symbol_export_from_raw_name(
+                tcx,
+                "__rustc_proc_macro_alloc_buffer".to_owned(),
+                SymbolExportKind::Text,
+            ),
+            symbol_export_from_raw_name(
+                tcx,
+                "__rustc_proc_macro_buffer_replace".to_owned(),
+                SymbolExportKind::Text,
+            ),
+            symbol_export_from_raw_name(
+                tcx,
+                "__rustc_proc_macro_buffer_ptr".to_owned(),
+                SymbolExportKind::Text,
+            ),
+            symbol_export_from_raw_name(
+                tcx,
+                "__rustc_proc_macro_buffer_len".to_owned(),
+                SymbolExportKind::Text,
+            ),
+            symbol_export_from_raw_name(
+                tcx,
+                "__rustc_proc_macro_call_client".to_owned(),
+                SymbolExportKind::Text,
+            ),
+        ]
+    } else {
+        vec![symbol_export_from_raw_name(tcx, proc_macro_decls_name, SymbolExportKind::Data)]
+    }
 }
 
 pub(crate) fn linked_symbols(

compiler/rustc_expand/src/proc_macro.rs

@@ -1,5 +1,6 @@
 use rustc_ast as ast;
 use rustc_ast::tokenstream::TokenStream;
+use rustc_data_structures::marker::IntoDynSyncSend;
 use rustc_data_structures::profiling::TimingGuard;
 use rustc_errors::ErrorGuaranteed;
 use rustc_middle::ty::{self, TyCtxt};
@@ -31,7 +32,7 @@ fn record_expand_proc_macro<'a>(
 }
 
 pub struct BangProcMacro {
-    pub client: pm::bridge::client::Client,
+    pub client: IntoDynSyncSend<pm::bridge::server::DynClient>,
 }
 
 impl base::BangProcMacro for BangProcMacro {
@@ -58,7 +59,7 @@ impl base::BangProcMacro for BangProcMacro {
 }
 
 pub struct AttrProcMacro {
-    pub client: pm::bridge::client::Client,
+    pub client: IntoDynSyncSend<pm::bridge::server::DynClient>,
 }
 
 impl base::AttrProcMacro for AttrProcMacro {
@@ -88,7 +89,7 @@ impl base::AttrProcMacro for AttrProcMacro {
 }
 
 pub struct DeriveProcMacro {
-    pub client: DeriveClient,
+    pub client: IntoDynSyncSend<DeriveClient>,
 }
 
 impl MultiItemModifier for DeriveProcMacro {
@@ -117,12 +118,12 @@ impl MultiItemModifier for DeriveProcMacro {
                 let input = &*tcx.arena.alloc(input);
                 let key: (LocalExpnId, &TokenStream) = (invoc_id, input);
 
-                QueryDeriveExpandCtx::enter(ecx, self.client, move || {
+                QueryDeriveExpandCtx::enter(ecx, self.client.0.clone(), move || {
                     tcx.derive_macro_expansion(key).cloned()
                 })
             })
         } else {
-            expand_derive_macro(invoc_id, input, ecx, self.client)
+            expand_derive_macro(invoc_id, input, ecx, &self.client)
         };
 
         let Ok(output) = res else {
@@ -178,13 +179,13 @@ pub(super) fn provide_derive_macro_expansion<'tcx>(
     })
 }
 
-type DeriveClient = pm::bridge::client::Client;
+type DeriveClient = pm::bridge::server::DynClient;
 
 fn expand_derive_macro(
     invoc_id: LocalExpnId,
     input: TokenStream,
     ecx: &mut ExtCtxt<'_>,
-    client: DeriveClient,
+    client: &DeriveClient,
 ) -> Result<TokenStream, ()> {
     let _timer =
         ecx.sess.prof.generic_activity_with_arg_recorder("expand_proc_macro", |recorder| {
@@ -239,7 +240,7 @@ impl QueryDeriveExpandCtx {
     /// Must be called while the `enter` function is active.
     fn with<F, R>(f: F) -> R
     where
-        F: for<'a, 'b> FnOnce(&'b mut ExtCtxt<'a>, DeriveClient) -> R,
+        F: for<'a, 'b> FnOnce(&'b mut ExtCtxt<'a>, &DeriveClient) -> R,
     {
         DERIVE_EXPAND_CTX.with(|ctx| {
             let ectx = {
@@ -252,7 +253,7 @@ impl QueryDeriveExpandCtx {
                 unsafe { casted.as_mut().unwrap() }
             };
 
-            f(ectx, ctx.client)
+            f(ectx, &ctx.client)
         })
     }
 }

compiler/rustc_metadata/Cargo.toml

@@ -32,6 +32,8 @@ rustc_span = { path = "../rustc_span" }
 rustc_target = { path = "../rustc_target" }
 tempfile = "3.7.1"
 tracing = "0.1"
+wasmi = "1.0.9"
+wasmi_wasi = "1.0.9"
 # tidy-alphabetical-end
 
 [target.'cfg(target_os = "aix")'.dependencies]

compiler/rustc_metadata/src/creader.rs

@@ -11,7 +11,7 @@ use rustc_crate_store::{CrateDepKind, CrateSource, ExternCrate, ExternCrateSourc
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::owned_slice::OwnedSlice;
 use rustc_data_structures::svh::Svh;
-use rustc_data_structures::sync::{self, FreezeReadGuard, FreezeWriteGuard};
+use rustc_data_structures::sync::{self, FreezeReadGuard, FreezeWriteGuard, IntoDynSyncSend};
 use rustc_data_structures::unord::UnordMap;
 use rustc_expand::base::SyntaxExtension;
 use rustc_hir as hir;
@@ -23,7 +23,7 @@ use rustc_lint_defs::builtin::UNUSED_CRATE_DEPENDENCIES;
 use rustc_middle::bug;
 use rustc_middle::ty::data_structures::IndexSet;
 use rustc_middle::ty::{TyCtxt, TyCtxtFeed};
-use rustc_proc_macro::bridge::client::Client as ProcMacroClient;
+use rustc_proc_macro::bridge::server::DynClient;
 use rustc_session::Session;
 use rustc_session::config::mitigation_coverage::DeniedPartialMitigationLevel;
 use rustc_session::config::{
@@ -941,7 +941,7 @@ impl CStore {
         &self,
         path: &Path,
         stable_crate_id: StableCrateId,
-    ) -> Result<&'static [ProcMacroClient], CrateError> {
+    ) -> Result<Vec<IntoDynSyncSend<DynClient>>, CrateError> {
         Ok(crate::host_dylib::dlsym_proc_macros(path, stable_crate_id)?)
     }
 

compiler/rustc_metadata/src/host_dylib.rs

@@ -1,9 +1,12 @@
 use std::error::Error;
 use std::path::Path;
+use std::sync::Arc;
 use std::time::Duration;
 
+use rustc_data_structures::sync::IntoDynSyncSend;
 use rustc_fs_util::try_canonicalize;
 use rustc_proc_macro::bridge::client::Client as ProcMacroClient;
+use rustc_proc_macro::bridge::server::DynClient;
 use rustc_session::StableCrateId;
 use tracing::debug;
 
@@ -126,22 +129,182 @@ pub unsafe fn load_symbol_from_dylib<T: Copy>(
 pub(crate) fn dlsym_proc_macros(
     path: &Path,
     stable_crate_id: StableCrateId,
-) -> Result<&'static [ProcMacroClient], DylibError> {
+) -> Result<Vec<IntoDynSyncSend<DynClient>>, DylibError> {
     let sym_name = rustc_session::generate_proc_macro_decls_symbol(stable_crate_id);
     debug!("trying to dlsym proc_macros {} for symbol `{}`", path.display(), sym_name);
 
     unsafe {
         // FIXME(bjorn3) this depends on the unstable slice memory layout
-        let result = crate::load_symbol_from_dylib::<*const &[ProcMacroClient]>(path, &sym_name);
+        let result = load_symbol_from_dylib::<*const &[ProcMacroClient]>(path, &sym_name);
         match result {
             Ok(result) => {
                 debug!("loaded dlsym proc_macros {} for symbol `{}`", path.display(), sym_name);
-                Ok(*result)
+                Ok((*result)
+                    .iter()
+                    .map(|proc_macro| IntoDynSyncSend(proc_macro.into_dyn_client()))
+                    .collect())
             }
-            Err(err) => {
-                debug!("failed to dlsym proc_macros {} for symbol `{}`", path.display(), sym_name);
-                Err(err.into())
+            Err(_err) => {
+                let engine = wasmi::Engine::default();
+                let module = wasmi::Module::new(&engine, std::fs::read(path).unwrap()).unwrap();
+
+                let mut store =
+                    wasmi::Store::new(&engine, wasmi_wasi::WasiCtxBuilder::new().build());
+                let mut linker = wasmi::Linker::new(&engine);
+                linker
+                    .func_wrap("env", "__rustc_proc_macro_dispatch", |_: u32, _: u32| -> () {
+                        unreachable!()
+                    })
+                    .unwrap();
+                wasmi_wasi::add_to_linker(&mut linker, |ctx| ctx).unwrap();
+                let instance = linker.instantiate_and_start(&mut store, &module).unwrap();
+
+                let memory = instance.get_export(&store, "memory").unwrap().into_memory().unwrap();
+
+                let sym = instance
+                    .get_export(&store, &sym_name)
+                    .unwrap()
+                    .into_global()
+                    .unwrap()
+                    .get(&store)
+                    .i32()
+                    .unwrap();
+
+                let mut data = [0; 8];
+                memory.read(&store, sym as usize, &mut data).unwrap();
+                let ptr = u32::from_le_bytes(data[..4].try_into().unwrap());
+                let len = u32::from_le_bytes(data[4..8].try_into().unwrap());
+
+                Ok((0..len)
+                    .map(|i| {
+                        let mut data = [0; 4];
+                        memory.read(&store, (ptr + 4 * i) as usize, &mut data).unwrap();
+                        let func_ptr = u32::from_le_bytes(data);
+
+                        IntoDynSyncSend(wasm_macro_client(engine.clone(), module.clone(), func_ptr))
+                    })
+                    .collect::<Vec<_>>())
             }
         }
     }
 }
+
+fn wasm_macro_client(engine: wasmi::Engine, module: wasmi::Module, func_ptr: u32) -> DynClient {
+    DynClient {
+        run: Arc::new(move |config| {
+            struct Ctx<'a> {
+                wasi_ctx: wasmi_wasi::WasiCtx,
+                dispatch: rustc_proc_macro::bridge::Closure<'a>,
+                client_refs: Option<ClientRefs>,
+            }
+
+            #[derive(Clone)]
+            struct ClientRefs {
+                memory: wasmi::Memory,
+                buffer_replace: wasmi::TypedFunc<(u32, u32), ()>,
+                buffer_ptr: wasmi::TypedFunc<(u32,), (u32,)>,
+                buffer_len: wasmi::TypedFunc<(u32,), (u32,)>,
+            }
+
+            impl ClientRefs {
+                fn read(
+                    &self,
+                    store: &mut impl wasmi::AsContextMut,
+                    buffer: u32,
+                ) -> Result<Vec<u8>, wasmi::Error> {
+                    let (buffer_ptr,) = self.buffer_ptr.call(&mut *store, (buffer,))?;
+                    let (buffer_len,) = self.buffer_len.call(&mut *store, (buffer,))?;
+                    let mut data = vec![0; buffer_len as usize];
+                    self.memory.read(store, buffer_ptr as usize, &mut data)?;
+                    Ok(data)
+                }
+            }
+
+            let mut store = wasmi::Store::new(
+                &engine,
+                Ctx {
+                    wasi_ctx: wasmi_wasi::WasiCtxBuilder::new()
+                        .inherit_stdout()
+                        .inherit_stderr()
+                        .build(),
+                    dispatch: config.dispatch,
+                    client_refs: None,
+                },
+            );
+            let mut linker: wasmi::Linker<Ctx<'_>> = wasmi::Linker::new(&engine);
+            linker
+                .func_new(
+                    "env",
+                    "__rustc_proc_macro_dispatch",
+                    wasmi::FuncType::new([wasmi::ValType::I32, wasmi::ValType::I32], []),
+                    |mut caller, inputs, _outputs| {
+                        let client_refs = caller.data().client_refs.clone().unwrap();
+                        let input_buffer = inputs[0].i32().unwrap().cast_unsigned();
+                        let output_buffer = inputs[1].i32().unwrap().cast_unsigned();
+
+                        let input_buffer_data = client_refs.read(&mut caller, input_buffer)?;
+
+                        let output_buffer_data =
+                            caller.data_mut().dispatch.call(input_buffer_data.into());
+
+                        client_refs.buffer_replace.call(
+                            &mut caller,
+                            (output_buffer, output_buffer_data.len().try_into().unwrap()),
+                        )?;
+                        let (output_buffer_ptr,) =
+                            client_refs.buffer_ptr.call(&mut caller, (output_buffer,))?;
+                        client_refs.memory.write(
+                            &mut caller,
+                            output_buffer_ptr as usize,
+                            &output_buffer_data,
+                        )?;
+
+                        Ok(())
+                    },
+                )
+                .unwrap();
+            wasmi_wasi::add_to_linker(&mut linker, |ctx| &mut ctx.wasi_ctx).unwrap();
+            let instance = linker.instantiate_and_start(&mut store, &module).unwrap();
+
+            fn get_func<T: wasmi::WasmParams, U: wasmi::WasmResults>(
+                instance: &wasmi::Instance,
+                store: &impl wasmi::AsContext,
+                name: &str,
+            ) -> wasmi::TypedFunc<T, U> {
+                instance.get_export(store, name).unwrap().into_func().unwrap().typed(store).unwrap()
+            }
+
+            let client_refs = ClientRefs {
+                memory: instance.get_export(&store, "memory").unwrap().into_memory().unwrap(),
+                buffer_replace: get_func(&instance, &store, "__rustc_proc_macro_buffer_replace"),
+                buffer_ptr: get_func(&instance, &store, "__rustc_proc_macro_buffer_ptr"),
+                buffer_len: get_func(&instance, &store, "__rustc_proc_macro_buffer_len"),
+            };
+            store.data_mut().client_refs = Some(client_refs.clone());
+
+            let alloc_buffer: wasmi::TypedFunc<(u32,), (u32,)> =
+                get_func(&instance, &store, "__rustc_proc_macro_alloc_buffer");
+
+            let call_client: wasmi::TypedFunc<(u32, u32), (u32,)> =
+                get_func(&instance, &store, "__rustc_proc_macro_call_client");
+
+            let (input_buffer,) =
+                alloc_buffer.call(&mut store, (config.input.len().try_into().unwrap(),)).unwrap();
+            let input_buffer_ptr =
+                client_refs.buffer_ptr.call(&mut store, (input_buffer,)).unwrap().0;
+            instance
+                .get_export(&store, "memory")
+                .unwrap()
+                .into_memory()
+                .unwrap()
+                .write(&mut store, input_buffer_ptr as usize, &config.input)
+                .unwrap();
+
+            let (output_buffer,) = call_client.call(&mut store, (input_buffer, func_ptr)).unwrap();
+
+            let output_buffer_data = client_refs.read(&mut store, output_buffer).unwrap();
+
+            output_buffer_data.into()
+        }),
+    }
+}