feat: wasm-splitting library support for future cargo-leptos integration

leptos/Cargo.toml

@@ -56,6 +56,7 @@ serde_qs = "0.14.0"
 slotmap = "1.0"
 futures = "0.3.31"
 send_wrapper = "0.6.0"
+wasm_split.workspace = true
 
 [features]
 hydration = [

leptos/src/lib.rs

@@ -343,5 +343,6 @@ pub use serde_json;
 pub use tracing;
 #[doc(hidden)]
 pub use wasm_bindgen;
+pub use wasm_split;
 #[doc(hidden)]
 pub use web_sys;

leptos_macro/src/lazy.rs

@@ -24,9 +24,14 @@ pub fn lazy_impl(
         fun.sig.ident.span(),
     );
 
-    quote! {
-        #[cfg_attr(feature = "split", wasm_split::wasm_split(#converted_name))]
-        #fun
+    let is_wasm = cfg!(feature = "csr") || cfg!(feature = "hydrate");
+    if is_wasm {
+        quote! {
+            #[::leptos::wasm_split::wasm_split(#converted_name)]
+            #fun
+        }
+    } else {
+        quote! { #fun }
     }
     .into()
 }

wasm_split/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "wasm_split"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+async-once-cell = { version = "0.5.3", features = ["std"] }
+wasm_split_macros.workspace = true

wasm_split/src/lib.rs

@@ -0,0 +1,110 @@
+use std::{
+    cell::Cell,
+    ffi::c_void,
+    future::Future,
+    pin::Pin,
+    rc::Rc,
+    task::{Context, Poll, Waker},
+};
+
+pub type LoadCallbackFn = unsafe extern "C" fn(*const c_void, bool) -> ();
+pub type LoadFn = unsafe extern "C" fn(LoadCallbackFn, *const c_void) -> ();
+
+type Lazy = async_once_cell::Lazy<Option<()>, SplitLoaderFuture>;
+
+pub use wasm_split_macros::wasm_split;
+
+pub struct LazySplitLoader {
+    lazy: Pin<Rc<Lazy>>,
+}
+
+impl LazySplitLoader {
+    pub unsafe fn new(load: LoadFn) -> Self {
+        Self {
+            lazy: Rc::pin(Lazy::new(SplitLoaderFuture::new(SplitLoader::new(
+                load,
+            )))),
+        }
+    }
+}
+
+pub async fn ensure_loaded(
+    loader: &'static std::thread::LocalKey<LazySplitLoader>,
+) -> Option<()> {
+    *loader.with(|inner| inner.lazy.clone()).as_ref().await
+}
+
+#[derive(Clone, Copy, Debug)]
+enum SplitLoaderState {
+    Deferred(LoadFn),
+    Pending,
+    Completed(Option<()>),
+}
+
+struct SplitLoader {
+    state: Cell<SplitLoaderState>,
+    waker: Cell<Option<Waker>>,
+}
+
+impl SplitLoader {
+    fn new(load: LoadFn) -> Rc<Self> {
+        Rc::new(SplitLoader {
+            state: Cell::new(SplitLoaderState::Deferred(load)),
+            waker: Cell::new(None),
+        })
+    }
+
+    fn complete(&self, value: bool) {
+        self.state.set(SplitLoaderState::Completed(if value {
+            Some(())
+        } else {
+            None
+        }));
+        match self.waker.take() {
+            Some(waker) => {
+                waker.wake();
+            }
+            _ => {}
+        }
+    }
+}
+
+struct SplitLoaderFuture {
+    loader: Rc<SplitLoader>,
+}
+
+impl SplitLoaderFuture {
+    fn new(loader: Rc<SplitLoader>) -> Self {
+        SplitLoaderFuture { loader }
+    }
+}
+
+impl Future for SplitLoaderFuture {
+    type Output = Option<()>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<()>> {
+        match self.loader.state.get() {
+            SplitLoaderState::Deferred(load) => {
+                self.loader.state.set(SplitLoaderState::Pending);
+                self.loader.waker.set(Some(cx.waker().clone()));
+                unsafe {
+                    load(
+                        load_callback,
+                        Rc::<SplitLoader>::into_raw(self.loader.clone())
+                            as *const c_void,
+                    )
+                };
+                Poll::Pending
+            }
+            SplitLoaderState::Pending => {
+                self.loader.waker.set(Some(cx.waker().clone()));
+                Poll::Pending
+            }
+            SplitLoaderState::Completed(value) => Poll::Ready(value),
+        }
+    }
+}
+
+unsafe extern "C" fn load_callback(loader: *const c_void, success: bool) {
+    unsafe { Rc::from_raw(loader as *const SplitLoader) }.complete(success);
+}

wasm_split_macros/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "wasm_split_macros"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+base16 = "0.2.1"
+digest = "0.10.7"
+quote = "1.0.36"
+sha2 = "0.10.8"
+syn = "2.0.59"
+wasm-bindgen = "0.2.92"
+
+[lib]
+proc-macro = true

wasm_split_macros/src/lib.rs

@@ -0,0 +1,92 @@
+use digest::Digest;
+use proc_macro::TokenStream;
+use quote::{format_ident, quote};
+use syn::{parse_macro_input, Ident, ItemFn, Signature};
+
+#[proc_macro_attribute]
+pub fn wasm_split(args: TokenStream, input: TokenStream) -> TokenStream {
+    let module_ident = parse_macro_input!(args as Ident);
+    let item_fn = parse_macro_input!(input as ItemFn);
+
+    let name = &item_fn.sig.ident;
+
+    let unique_identifier = base16::encode_lower(
+        &sha2::Sha256::digest(format!("{name} {span:?}", span = name.span()))
+            [..16],
+    );
+
+    let load_module_ident = format_ident!("__wasm_split_load_{module_ident}");
+    let split_loader_ident = format_ident!("__wasm_split_loader");
+    let impl_import_ident = format_ident!(
+        "__wasm_split_00{module_ident}00_import_{unique_identifier}_{name}"
+    );
+    let impl_export_ident = format_ident!(
+        "__wasm_split_00{module_ident}00_export_{unique_identifier}_{name}"
+    );
+
+    let import_sig = Signature {
+        ident: impl_import_ident.clone(),
+        asyncness: None,
+        ..item_fn.sig.clone()
+    };
+    let export_sig = Signature {
+        ident: impl_export_ident.clone(),
+        asyncness: None,
+        ..item_fn.sig.clone()
+    };
+
+    let mut wrapper_sig = item_fn.sig;
+    wrapper_sig.asyncness = Some(Default::default());
+    let mut args = Vec::new();
+    for (i, param) in wrapper_sig.inputs.iter_mut().enumerate() {
+        match param {
+            syn::FnArg::Typed(pat_type) => {
+                let param_ident = format_ident!("__wasm_split_arg_{i}");
+                args.push(param_ident.clone());
+                pat_type.pat = Box::new(syn::Pat::Ident(syn::PatIdent {
+                    attrs: vec![],
+                    by_ref: None,
+                    mutability: None,
+                    ident: param_ident,
+                    subpat: None,
+                }));
+            }
+            syn::FnArg::Receiver(_) => {
+                args.push(format_ident!("self"));
+            }
+        }
+    }
+
+    let attrs = item_fn.attrs;
+
+    let stmts = &item_fn.block.stmts;
+
+    quote! {
+        #wrapper_sig {
+            thread_local! {
+                static #split_loader_ident: ::leptos::wasm_split::LazySplitLoader = unsafe { ::leptos::wasm_split::LazySplitLoader::new(#load_module_ident) };
+            }
+
+            #[link(wasm_import_module = "/pkg/__wasm_split.js")]
+            extern "C" {
+                #[no_mangle]
+                fn #load_module_ident (callback: unsafe extern "C" fn(*const ::std::ffi::c_void, bool), data: *const ::std::ffi::c_void) -> ();
+
+                #[allow(improper_ctypes)]
+                #[no_mangle]
+                #import_sig;
+            }
+
+            #(#attrs)*
+            #[allow(improper_ctypes_definitions)]
+            #[no_mangle]
+            pub extern "C" #export_sig {
+                #(#stmts)*
+            }
+
+            ::leptos::wasm_split::ensure_loaded(&#split_loader_ident).await.unwrap();
+            unsafe { #impl_import_ident( #(#args),* ) }
+        }
+    }
+    .into()
+}