feat: router-based wasm bundle splitting (#3683)
* feat: wasm-split base implementation * kinda sorta works with the cli * okay great it actually works * fix: cache busting for chunks * cleaned up version * it works with the router * wip: migrate to a faster linear solution using id remapping * add a good harness * it works completely on the docsite * the super tight build works too * make parallelizable * works on the docsite but not the harness * light cleanups, fix for both harness and docsite * make codegen smaller by extract some shared functions to prevent monos * chunking works, but tbd on tuning heuristics * extract used module * wip: cleaning up module * wip: much better understanding of how imports work * fix the call graph * working again... * wip: thinking about ifuncs * ifunc approach works * yes the ifunc approach works!! * it all works completely even with chunks * small cleanups * comments * actually locate all of the symbols * clean up final implementation * clean up tomls, revert some mono changs * clean up, add some more comments / docs in places * hoist more packages * clean up impl a bit more * clean cli * move harness to playwright * update playwright * extract codgen * add lazy to components * Add docs to split loader * clippy * typos * clippy, don't generate splits if not on wasm * whoops, const * Typos and docs * enable optimizations in ci * add optimizations to ci * add wasm-split router feature * parallel playwright? * fix no case check * disable pre-compress... think that might be the issue * better logging, no default pre-compress, caching of shared symbols * Fix: compressing assets * fix non-interactive logging for serve * cache playwright even on failure, use the cli with optimizations * fix prebuild * fix playwright webserver
This commit is contained in:
Jonathan Kelley authored and GitHub committed
1 parent
14eb7597df
commit
bdeedc13eb
59 files changed
+4249
-299
No files matched your search
@@ -0,0 +1,10 @@
|
||||
# This folder contains the wasm-split sub-workspace
|
||||
|
||||
wasm-split is a tool that allows you to split a wasm binary into multiple chunks that are lazily loaded on demand.
|
||||
|
||||
This workspace is comprised of:
|
||||
- the harness on which we test against
|
||||
- the wasm-split user-facing crate
|
||||
- the wasm-split-macro crate which is used to generate the wasm-split loader
|
||||
- the wasm-split-cli crate which is imported by the dioxus-cli and used when building the wasm modules
|
||||
- the wasm-used crate which provides a custom walrus `Used` struct that makes it easier to debug why a emit_wasm might be failing
|
||||
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
name = "wasm-split-cli"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
anyhow = { workspace = true }
|
||||
itertools = { workspace = true }
|
||||
walrus = { workspace = true, features = ["parallel"]}
|
||||
wasmparser = { workspace = true }
|
||||
id-arena = { workspace = true }
|
||||
rayon = { workspace = true }
|
||||
tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] }
|
||||
tracing = { workspace = true }
|
||||
clap = { workspace = true, features = ["derive"] }
|
||||
wasm-used = { path = "../wasm-used" }
|
||||
@@ -0,0 +1,4 @@
|
||||
# this folder contains the wasm output for the test harnessa, so ignore it
|
||||
|
||||
*.wasm
|
||||
bindgen/
|
||||
@@ -0,0 +1,63 @@
|
||||
// when running the harness we need to make sure to uncommon this out...
|
||||
|
||||
export function makeLoad(url, deps, fusedImports, initIt) {
|
||||
let alreadyLoaded = false;
|
||||
return async (callbackIndex, callbackData) => {
|
||||
await Promise.all(deps.map((dep) => dep()));
|
||||
if (alreadyLoaded) return;
|
||||
try {
|
||||
const response = await fetch(url);
|
||||
const initSync = initIt || window.__wasm_split_main_initSync;
|
||||
const mainExports = initSync(undefined, undefined);
|
||||
|
||||
let imports = {
|
||||
env: {
|
||||
memory: mainExports.memory,
|
||||
},
|
||||
__wasm_split: {
|
||||
__indirect_function_table: mainExports.__indirect_function_table,
|
||||
__stack_pointer: mainExports.__stack_pointer,
|
||||
__tls_base: mainExports.__tls_base,
|
||||
memory: mainExports.memory,
|
||||
},
|
||||
};
|
||||
|
||||
for (let mainExport in mainExports) {
|
||||
imports["__wasm_split"][mainExport] = mainExports[mainExport];
|
||||
}
|
||||
|
||||
for (let name in fusedImports) {
|
||||
imports["__wasm_split"][name] = fusedImports[name];
|
||||
}
|
||||
|
||||
let new_exports = await WebAssembly.instantiateStreaming(
|
||||
response,
|
||||
imports
|
||||
);
|
||||
|
||||
alreadyLoaded = true;
|
||||
|
||||
for (let name in new_exports.instance.exports) {
|
||||
fusedImports[name] = new_exports.instance.exports[name];
|
||||
}
|
||||
|
||||
if (callbackIndex !== undefined) {
|
||||
mainExports.__indirect_function_table.get(callbackIndex)(
|
||||
callbackData,
|
||||
true
|
||||
);
|
||||
}
|
||||
} catch (e) {
|
||||
console.error(
|
||||
"Failed to load wasm-split module",
|
||||
e,
|
||||
url,
|
||||
deps,
|
||||
fusedImports
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
let fusedImports = {};
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,178 @@
|
||||
use clap::Parser;
|
||||
use std::path::PathBuf;
|
||||
use wasm_split_cli::SplitModule;
|
||||
|
||||
fn main() {
|
||||
tracing_subscriber::fmt()
|
||||
.without_time()
|
||||
.compact()
|
||||
.with_env_filter("debug,walrus=info")
|
||||
.init();
|
||||
|
||||
match Commands::parse() {
|
||||
Commands::Split(split_args) => split(split_args),
|
||||
Commands::Validate(validate_args) => validate(validate_args),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Parser)]
|
||||
enum Commands {
|
||||
/// Split a wasm module into multiple chunks
|
||||
#[clap(name = "split")]
|
||||
Split(SplitArgs),
|
||||
|
||||
/// Validate the main module of a wasm module
|
||||
#[clap(name = "validate")]
|
||||
Validate(ValidateArgs),
|
||||
}
|
||||
|
||||
#[derive(Parser)]
|
||||
struct SplitArgs {
|
||||
/// The wasm module emitted by rustc
|
||||
original: PathBuf,
|
||||
|
||||
/// The wasm module emitted by wasm-bindgen
|
||||
bindgened: PathBuf,
|
||||
|
||||
/// The output *directory* to write the split wasm files to
|
||||
out_dir: PathBuf,
|
||||
}
|
||||
|
||||
fn split(args: SplitArgs) {
|
||||
let original = std::fs::read(&args.original).expect("failed to read input file");
|
||||
let bindgened = std::fs::read(&args.bindgened).expect("failed to read input file");
|
||||
|
||||
_ = std::fs::remove_dir_all(&args.out_dir);
|
||||
std::fs::create_dir_all(&args.out_dir).expect("failed to create output dir");
|
||||
|
||||
tracing::info!("Building split module");
|
||||
|
||||
let module = wasm_split_cli::Splitter::new(&original, &bindgened).unwrap();
|
||||
|
||||
let mut chunks = module.emit().unwrap();
|
||||
|
||||
// Write out the main module
|
||||
tracing::info!(
|
||||
"Writing main module to {}",
|
||||
args.out_dir.join("main.wasm").display()
|
||||
);
|
||||
std::fs::write(args.out_dir.join("main.wasm"), &chunks.main.bytes).unwrap();
|
||||
|
||||
// Write the js module
|
||||
std::fs::write(
|
||||
args.out_dir.join("__wasm_split.js"),
|
||||
emit_js(&chunks.chunks, &chunks.modules),
|
||||
)
|
||||
.expect("failed to write js module");
|
||||
|
||||
for (idx, chunk) in chunks.chunks.iter().enumerate() {
|
||||
tracing::info!(
|
||||
"Writing chunk {} to {}",
|
||||
idx,
|
||||
args.out_dir
|
||||
.join(format!("chunk_{}_{}.wasm", idx, chunk.module_name))
|
||||
.display()
|
||||
);
|
||||
std::fs::write(
|
||||
args.out_dir
|
||||
.join(format!("chunk_{}_{}.wasm", idx, chunk.module_name)),
|
||||
&chunk.bytes,
|
||||
)
|
||||
.expect("failed to write chunk");
|
||||
}
|
||||
|
||||
for (idx, module) in chunks.modules.iter_mut().enumerate() {
|
||||
tracing::info!(
|
||||
"Writing module {} to {}",
|
||||
idx,
|
||||
args.out_dir
|
||||
.join(format!(
|
||||
"module_{}_{}.wasm",
|
||||
idx,
|
||||
module.component_name.as_ref().unwrap()
|
||||
))
|
||||
.display()
|
||||
);
|
||||
std::fs::write(
|
||||
args.out_dir.join(format!(
|
||||
"module_{}_{}.wasm",
|
||||
idx,
|
||||
module.component_name.as_ref().unwrap()
|
||||
)),
|
||||
&module.bytes,
|
||||
)
|
||||
.expect("failed to write chunk");
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_js(chunks: &[SplitModule], modules: &[SplitModule]) -> String {
|
||||
use std::fmt::Write;
|
||||
let mut glue = format!(
|
||||
r#"import {{ initSync }} from "./main.js";
|
||||
{}"#,
|
||||
include_str!("./__wasm_split.js")
|
||||
);
|
||||
|
||||
for (idx, chunk) in chunks.iter().enumerate() {
|
||||
tracing::debug!("emitting chunk: {:?}", chunk.module_name);
|
||||
writeln!(
|
||||
glue,
|
||||
"export const __wasm_split_load_chunk_{idx} = makeLoad(\"/harness/split/chunk_{idx}_{module}.wasm\", [], fusedImports, initSync);",
|
||||
module = chunk.module_name
|
||||
).expect("failed to write to string");
|
||||
}
|
||||
|
||||
// Now write the modules
|
||||
for (idx, module) in modules.iter().enumerate() {
|
||||
let deps = module
|
||||
.relies_on_chunks
|
||||
.iter()
|
||||
.map(|idx| format!("__wasm_split_load_chunk_{idx}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let hash_id = module.hash_id.as_ref().unwrap();
|
||||
|
||||
writeln!(
|
||||
glue,
|
||||
"export const __wasm_split_load_{module}_{hash_id}_{cname} = makeLoad(\"/harness/split/module_{idx}_{cname}.wasm\", [{deps}], fusedImports, initSync);",
|
||||
module = module.module_name,
|
||||
idx = idx,
|
||||
cname = module.component_name.as_ref().unwrap(),
|
||||
deps = deps
|
||||
)
|
||||
.expect("failed to write to string");
|
||||
}
|
||||
|
||||
glue
|
||||
}
|
||||
|
||||
#[derive(Parser)]
|
||||
struct ValidateArgs {
|
||||
/// The input wasm file to validate
|
||||
main: PathBuf,
|
||||
|
||||
chunks: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
fn validate(args: ValidateArgs) {
|
||||
let bytes = std::fs::read(&args.main).expect("failed to read input file");
|
||||
let main_module = walrus::Module::from_buffer(&bytes).unwrap();
|
||||
|
||||
for chunk in args.chunks {
|
||||
let bytes = std::fs::read(chunk).expect("failed to read input file");
|
||||
let chunk_module = walrus::Module::from_buffer(&bytes).unwrap();
|
||||
|
||||
assert!(chunk_module.tables.iter().count() == 1);
|
||||
|
||||
for import in chunk_module.imports.iter() {
|
||||
let matching = main_module.exports.iter().find(|e| e.name == import.name);
|
||||
|
||||
let Some(matching) = matching else {
|
||||
tracing::error!("Could not find matching export for import {import:#?}");
|
||||
continue;
|
||||
};
|
||||
|
||||
tracing::debug!("import: {:?}", matching.name);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
[package]
|
||||
name = "wasm-split-macro"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
syn = { workspace = true, features = ["full"] }
|
||||
base16 = { workspace = true }
|
||||
digest = { workspace = true }
|
||||
quote = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
proc-macro2 = { workspace = true }
|
||||
|
||||
[lib]
|
||||
proc-macro = true
|
||||
@@ -0,0 +1,226 @@
|
||||
use proc_macro::TokenStream;
|
||||
|
||||
use digest::Digest;
|
||||
use quote::{format_ident, quote};
|
||||
use syn::{parse_macro_input, parse_quote, FnArg, Ident, ItemFn, ReturnType, 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);
|
||||
|
||||
if item_fn.sig.asyncness.is_none() {
|
||||
panic!("wasm_split functions must be async. Use a LazyLoader with synchronous functions instead.");
|
||||
}
|
||||
|
||||
let LoaderNames {
|
||||
split_loader_ident,
|
||||
impl_import_ident,
|
||||
impl_export_ident,
|
||||
load_module_ident,
|
||||
..
|
||||
} = LoaderNames::new(item_fn.sig.ident.clone(), module_ident.to_string());
|
||||
|
||||
let mut desugard_async_sig = item_fn.sig.clone();
|
||||
desugard_async_sig.asyncness = None;
|
||||
desugard_async_sig.output = match &desugard_async_sig.output {
|
||||
ReturnType::Default => {
|
||||
parse_quote! { -> std::pin::Pin<Box<dyn std::future::Future<Output = ()>>> }
|
||||
}
|
||||
ReturnType::Type(_, ty) => {
|
||||
parse_quote! { -> std::pin::Pin<Box<dyn std::future::Future<Output = #ty>>> }
|
||||
}
|
||||
};
|
||||
|
||||
let import_sig = Signature {
|
||||
ident: impl_import_ident.clone(),
|
||||
..desugard_async_sig.clone()
|
||||
};
|
||||
|
||||
let export_sig = Signature {
|
||||
ident: impl_export_ident.clone(),
|
||||
..desugard_async_sig.clone()
|
||||
};
|
||||
|
||||
let default_item = item_fn.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::Receiver(_) => args.push(format_ident!("self")),
|
||||
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,
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let attrs = &item_fn.attrs;
|
||||
let stmts = &item_fn.block.stmts;
|
||||
|
||||
quote! {
|
||||
#[cfg(target_arch = "wasm32")]
|
||||
#wrapper_sig {
|
||||
#(#attrs)*
|
||||
#[allow(improper_ctypes_definitions)]
|
||||
#[no_mangle]
|
||||
pub extern "C" #export_sig {
|
||||
Box::pin(async move { #(#stmts)* })
|
||||
}
|
||||
|
||||
#[link(wasm_import_module = "./__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;
|
||||
}
|
||||
|
||||
thread_local! {
|
||||
static #split_loader_ident: wasm_split::LazySplitLoader = unsafe {
|
||||
wasm_split::LazySplitLoader::new(#load_module_ident)
|
||||
};
|
||||
}
|
||||
|
||||
// Initiate the download by calling the load_module_ident function which will kick-off the loader
|
||||
if !wasm_split::LazySplitLoader::ensure_loaded(&#split_loader_ident).await {
|
||||
panic!("Failed to load wasm-split module");
|
||||
}
|
||||
|
||||
unsafe { #impl_import_ident( #(#args),* ) }.await
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
#default_item
|
||||
}
|
||||
.into()
|
||||
}
|
||||
|
||||
/// Create a lazy loader for a given function. Meant to be used in statics. Designed for libraries to
|
||||
/// integrate with.
|
||||
///
|
||||
/// ```rust, no_run
|
||||
/// fn SomeFunction(args: Args) -> Ret {}
|
||||
///
|
||||
/// static LOADER: wasm_split::LazyLoader<Args, Ret> = lazy_loader!(SomeFunction);
|
||||
///
|
||||
/// LOADER.load().await.call(args)
|
||||
/// ```
|
||||
#[proc_macro]
|
||||
pub fn lazy_loader(input: TokenStream) -> TokenStream {
|
||||
// We can only accept idents/paths that will be the source function
|
||||
let sig = parse_macro_input!(input as Signature);
|
||||
let params = sig.inputs.clone();
|
||||
let outputs = sig.output.clone();
|
||||
let Some(FnArg::Typed(arg)) = params.first().cloned() else {
|
||||
panic!(
|
||||
"Lazy Loader must define a single input argument to satisfy the LazyLoader signature"
|
||||
)
|
||||
};
|
||||
let arg_ty = arg.ty.clone();
|
||||
let LoaderNames {
|
||||
name,
|
||||
split_loader_ident,
|
||||
impl_import_ident,
|
||||
impl_export_ident,
|
||||
load_module_ident,
|
||||
..
|
||||
} = LoaderNames::new(
|
||||
sig.ident.clone(),
|
||||
sig.abi
|
||||
.as_ref()
|
||||
.and_then(|abi| abi.name.as_ref().map(|f| f.value()))
|
||||
.expect("abi to be module name")
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
quote! {
|
||||
{
|
||||
#[cfg(target_arch = "wasm32")]
|
||||
{
|
||||
#[link(wasm_import_module = "./__wasm_split.js")]
|
||||
extern "C" {
|
||||
// The function we'll use to initiate the download of the module
|
||||
#[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]
|
||||
fn #impl_import_ident(arg: #arg_ty) #outputs;
|
||||
}
|
||||
|
||||
|
||||
#[allow(improper_ctypes_definitions)]
|
||||
#[no_mangle]
|
||||
pub extern "C" fn #impl_export_ident(arg: #arg_ty) #outputs {
|
||||
#name(arg)
|
||||
}
|
||||
|
||||
thread_local! {
|
||||
static #split_loader_ident: wasm_split::LazySplitLoader = unsafe {
|
||||
wasm_split::LazySplitLoader::new(#load_module_ident)
|
||||
};
|
||||
};
|
||||
|
||||
unsafe {
|
||||
wasm_split::LazyLoader::new(#impl_import_ident, &#split_loader_ident)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
{
|
||||
wasm_split::LazyLoader::preloaded(#name)
|
||||
}
|
||||
}
|
||||
}
|
||||
.into()
|
||||
}
|
||||
|
||||
struct LoaderNames {
|
||||
name: Ident,
|
||||
split_loader_ident: Ident,
|
||||
impl_import_ident: Ident,
|
||||
impl_export_ident: Ident,
|
||||
load_module_ident: Ident,
|
||||
}
|
||||
|
||||
impl LoaderNames {
|
||||
fn new(name: Ident, module: String) -> Self {
|
||||
let unique_identifier = base16::encode_lower(
|
||||
&sha2::Sha256::digest(format!("{name} {span:?}", name = name, span = name.span()))
|
||||
[..16],
|
||||
);
|
||||
|
||||
Self {
|
||||
split_loader_ident: format_ident!("__wasm_split_loader_{module}"),
|
||||
impl_export_ident: format_ident!(
|
||||
"__wasm_split_00___{module}___00_export_{unique_identifier}_{name}"
|
||||
),
|
||||
impl_import_ident: format_ident!(
|
||||
"__wasm_split_00___{module}___00_import_{unique_identifier}_{name}"
|
||||
),
|
||||
load_module_ident: format_ident!(
|
||||
"__wasm_split_load_{module}_{unique_identifier}_{name}"
|
||||
),
|
||||
name,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
[package]
|
||||
name = "wasm-split"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
async-once-cell = { workspace = true, features = ["std"] }
|
||||
wasm-split-macro = { workspace = true }
|
||||
@@ -0,0 +1,214 @@
|
||||
use std::{
|
||||
cell::Cell,
|
||||
ffi::c_void,
|
||||
future::Future,
|
||||
pin::Pin,
|
||||
rc::Rc,
|
||||
task::{Context, Poll, Waker},
|
||||
thread::LocalKey,
|
||||
};
|
||||
|
||||
pub use wasm_split_macro::{lazy_loader, wasm_split};
|
||||
|
||||
pub type Result<T> = std::result::Result<T, SplitLoaderError>;
|
||||
|
||||
#[non_exhaustive]
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum SplitLoaderError {
|
||||
FailedToLoad,
|
||||
}
|
||||
impl std::fmt::Display for SplitLoaderError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
SplitLoaderError::FailedToLoad => write!(f, "Failed to load wasm-split module"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A lazy loader that can be used to load a function from a split out `.wasm` file.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// To use the split loader, you must first create the loader using the `lazy_loader` macro. This macro
|
||||
/// requires the complete signature of the function you want to load. The extern abi string denotes
|
||||
/// which module the function should be loaded from. If you don't know which module to use, use `auto`
|
||||
/// and wasm-split will automatically combine all the modules into one.
|
||||
///
|
||||
/// ```rust, no_run
|
||||
/// static LOADER: wasm_split::LazyLoader<Args, Ret> = wasm_split::lazy_loader!(extern "auto" fn SomeFunction(args: Args) -> Ret);
|
||||
///
|
||||
/// fn SomeFunction(args: Args) -> Ret {
|
||||
/// // Implementation
|
||||
/// }
|
||||
/// ```
|
||||
///
|
||||
/// ## The `#[component(lazy)]` macro
|
||||
///
|
||||
/// If you're using wasm-split with Dioxus, the `#[component(lazy)]` macro is provided that wraps
|
||||
/// the lazy loader with suspense. This means that the component will suspense until its body has
|
||||
/// been loaded.
|
||||
///
|
||||
/// ```rust, ignore
|
||||
/// fn app() -> Element {
|
||||
/// rsx! {
|
||||
/// Suspense {
|
||||
/// fallback: rsx! { "Loading..." },
|
||||
/// LazyComponent { abc: 0 }
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
///
|
||||
/// #[component(lazy)]
|
||||
/// fn LazyComponent(abc: i32) -> Element {
|
||||
/// rsx! {
|
||||
/// div {
|
||||
/// "This is a lazy component! {abc}"
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub struct LazyLoader<Args, Ret> {
|
||||
imported: unsafe extern "C" fn(arg: Args) -> Ret,
|
||||
key: &'static LocalKey<LazySplitLoader>,
|
||||
}
|
||||
|
||||
impl<Args, Ret> LazyLoader<Args, Ret> {
|
||||
/// Create a new lazy loader from a lazy imported function and a LazySplitLoader
|
||||
///
|
||||
/// # Safety
|
||||
/// This is unsafe because we're taking an arbitrary function pointer and using it as the loader.
|
||||
/// This function is likely not instantiated when passed here, so it should never be called directly.
|
||||
#[doc(hidden)]
|
||||
pub const unsafe fn new(
|
||||
imported: unsafe extern "C" fn(arg: Args) -> Ret,
|
||||
key: &'static LocalKey<LazySplitLoader>,
|
||||
) -> Self {
|
||||
Self { imported, key }
|
||||
}
|
||||
|
||||
/// Create a new lazy loader that is already resolved.
|
||||
pub const fn preloaded(f: fn(Args) -> Ret) -> Self {
|
||||
let imported =
|
||||
unsafe { std::mem::transmute::<fn(Args) -> Ret, unsafe extern "C" fn(Args) -> Ret>(f) };
|
||||
|
||||
thread_local! {
|
||||
static LAZY: LazySplitLoader = LazySplitLoader::preloaded();
|
||||
};
|
||||
|
||||
Self {
|
||||
imported,
|
||||
key: &LAZY,
|
||||
}
|
||||
}
|
||||
|
||||
/// Load the lazy loader, returning an boolean indicating whether it loaded successfully
|
||||
pub async fn load(&'static self) -> bool {
|
||||
*self.key.with(|inner| inner.lazy.clone()).as_ref().await
|
||||
}
|
||||
|
||||
/// Call the lazy loader with the given arguments
|
||||
pub fn call(&'static self, args: Args) -> Result<Ret> {
|
||||
let Some(true) = self.key.with(|inner| inner.lazy.try_get().copied()) else {
|
||||
return Err(SplitLoaderError::FailedToLoad);
|
||||
};
|
||||
|
||||
Ok(unsafe { (self.imported)(args) })
|
||||
}
|
||||
}
|
||||
|
||||
type Lazy = async_once_cell::Lazy<bool, SplitLoaderFuture>;
|
||||
type LoadCallbackFn = unsafe extern "C" fn(*const c_void, bool) -> ();
|
||||
type LoadFn = unsafe extern "C" fn(LoadCallbackFn, *const c_void) -> ();
|
||||
|
||||
pub struct LazySplitLoader {
|
||||
lazy: Pin<Rc<Lazy>>,
|
||||
}
|
||||
|
||||
impl LazySplitLoader {
|
||||
/// Create a new lazy split loader from a load function that is generated by the wasm-split macro
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// This is unsafe because we're taking an arbitrary function pointer and using it as the loader.
|
||||
/// It is likely not instantiated when passed here, so it should never be called directly.
|
||||
#[doc(hidden)]
|
||||
pub unsafe fn new(load: LoadFn) -> Self {
|
||||
Self {
|
||||
lazy: Rc::pin(Lazy::new({
|
||||
SplitLoaderFuture {
|
||||
loader: Rc::new(SplitLoader {
|
||||
state: Cell::new(SplitLoaderState::Deferred(load)),
|
||||
waker: Cell::new(None),
|
||||
}),
|
||||
}
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
fn preloaded() -> Self {
|
||||
Self {
|
||||
lazy: Rc::pin(Lazy::new({
|
||||
SplitLoaderFuture {
|
||||
loader: Rc::new(SplitLoader {
|
||||
state: Cell::new(SplitLoaderState::Completed(true)),
|
||||
waker: Cell::new(None),
|
||||
}),
|
||||
}
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// Wait for the lazy loader to load
|
||||
pub async fn ensure_loaded(loader: &'static std::thread::LocalKey<LazySplitLoader>) -> bool {
|
||||
*loader.with(|inner| inner.lazy.clone()).as_ref().await
|
||||
}
|
||||
}
|
||||
|
||||
struct SplitLoader {
|
||||
state: Cell<SplitLoaderState>,
|
||||
waker: Cell<Option<Waker>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum SplitLoaderState {
|
||||
Deferred(LoadFn),
|
||||
Pending,
|
||||
Completed(bool),
|
||||
}
|
||||
|
||||
struct SplitLoaderFuture {
|
||||
loader: Rc<SplitLoader>,
|
||||
}
|
||||
|
||||
impl Future for SplitLoaderFuture {
|
||||
type Output = bool;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<bool> {
|
||||
unsafe extern "C" fn load_callback(loader: *const c_void, success: bool) {
|
||||
let loader = unsafe { Rc::from_raw(loader as *const SplitLoader) };
|
||||
loader.state.set(SplitLoaderState::Completed(success));
|
||||
if let Some(waker) = loader.waker.take() {
|
||||
waker.wake()
|
||||
}
|
||||
}
|
||||
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
[package]
|
||||
name = "wasm-used"
|
||||
edition = "2021"
|
||||
version.workspace = true
|
||||
|
||||
[dependencies]
|
||||
walrus = { workspace = true, features = ["parallel"] }
|
||||
id-arena = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
@@ -0,0 +1,346 @@
|
||||
use std::collections::HashSet;
|
||||
|
||||
use id_arena::Id;
|
||||
use walrus::{ir::*, ExportId};
|
||||
use walrus::{ConstExpr, Data, DataId, DataKind, Element, ExportItem, Function};
|
||||
use walrus::{ElementId, ElementItems, ElementKind, Module, RefType, Type, TypeId};
|
||||
use walrus::{FunctionId, FunctionKind, Global, GlobalId};
|
||||
use walrus::{GlobalKind, Memory, MemoryId, Table, TableId};
|
||||
|
||||
type IdHashSet<T> = HashSet<Id<T>>;
|
||||
|
||||
/// Set of all root used items in a wasm module.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct Roots {
|
||||
tables: Vec<TableId>,
|
||||
funcs: Vec<(FunctionId, Location)>,
|
||||
globals: Vec<GlobalId>,
|
||||
memories: Vec<MemoryId>,
|
||||
data: Vec<DataId>,
|
||||
elements: Vec<ElementId>,
|
||||
used: Used,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug)]
|
||||
pub enum Location {
|
||||
Start,
|
||||
Export { export: ExportId },
|
||||
Table { table: TableId },
|
||||
Memory { memory: MemoryId },
|
||||
Global { global: GlobalId },
|
||||
Data,
|
||||
Element { element: ElementId },
|
||||
Code { func: FunctionId },
|
||||
}
|
||||
|
||||
impl Roots {
|
||||
/// Creates a new set of empty roots.
|
||||
pub fn new() -> Roots {
|
||||
Roots::default()
|
||||
}
|
||||
|
||||
/// Adds a new function to the set of roots
|
||||
pub fn push_func(&mut self, func: FunctionId, from: Location) -> &mut Roots {
|
||||
if self.used.funcs.insert(func) {
|
||||
// log::trace!("function is used: {:?}", func);
|
||||
self.funcs.push((func, from));
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
/// Adds a new table to the set of roots
|
||||
pub fn push_table(&mut self, table: TableId) -> &mut Roots {
|
||||
if self.used.tables.insert(table) {
|
||||
// log::trace!("table is used: {:?}", table);
|
||||
self.tables.push(table);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
/// Adds a new memory to the set of roots
|
||||
pub fn push_memory(&mut self, memory: MemoryId) -> &mut Roots {
|
||||
if self.used.memories.insert(memory) {
|
||||
// log::trace!("memory is used: {:?}", memory);
|
||||
self.memories.push(memory);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
/// Adds a new global to the set of roots
|
||||
pub fn push_global(&mut self, global: GlobalId) -> &mut Roots {
|
||||
if self.used.globals.insert(global) {
|
||||
// log::trace!("global is used: {:?}", global);
|
||||
self.globals.push(global);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
fn push_data(&mut self, data: DataId) -> &mut Roots {
|
||||
if self.used.data.insert(data) {
|
||||
// log::trace!("data is used: {:?}", data);
|
||||
self.data.push(data);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
fn push_element(&mut self, element: ElementId) -> &mut Roots {
|
||||
if self.used.elements.insert(element) {
|
||||
// log::trace!("element is used: {:?}", element);
|
||||
self.elements.push(element);
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Finds the things within a module that are used.
|
||||
///
|
||||
/// This is useful for implementing something like a linker's `--gc-sections` so
|
||||
/// that our emitted `.wasm` binaries are small and don't contain things that
|
||||
/// are not used.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct Used {
|
||||
/// The module's used tables.
|
||||
pub tables: IdHashSet<Table>,
|
||||
/// The module's used types.
|
||||
pub types: IdHashSet<Type>,
|
||||
/// The module's used functions.
|
||||
pub funcs: IdHashSet<Function>,
|
||||
/// The module's used globals.
|
||||
pub globals: IdHashSet<Global>,
|
||||
/// The module's used memories.
|
||||
pub memories: IdHashSet<Memory>,
|
||||
/// The module's used passive element segments.
|
||||
pub elements: IdHashSet<Element>,
|
||||
/// The module's used passive data segments.
|
||||
pub data: IdHashSet<Data>,
|
||||
}
|
||||
|
||||
impl Used {
|
||||
/// Construct a new `Used` set for the given module.
|
||||
pub fn new(module: &Module, deleted: &HashSet<FunctionId>) -> Used {
|
||||
// log::debug!("starting to calculate used set");
|
||||
let mut stack = Roots::default();
|
||||
|
||||
// All exports are roots
|
||||
for export in module.exports.iter() {
|
||||
match export.item {
|
||||
ExportItem::Function(f) => stack.push_func(
|
||||
f,
|
||||
Location::Export {
|
||||
export: export.id(),
|
||||
},
|
||||
),
|
||||
ExportItem::Table(t) => stack.push_table(t),
|
||||
ExportItem::Memory(m) => stack.push_memory(m),
|
||||
ExportItem::Global(g) => stack.push_global(g),
|
||||
};
|
||||
}
|
||||
|
||||
// The start function is an implicit root as well
|
||||
if let Some(f) = module.start {
|
||||
stack.push_func(f, Location::Start);
|
||||
}
|
||||
|
||||
// Initialization of memories or tables is a side-effectful operation
|
||||
// because they can be out-of-bounds, so keep all active segments.
|
||||
for data in module.data.iter() {
|
||||
if let DataKind::Active { .. } = &data.kind {
|
||||
stack.push_data(data.id());
|
||||
}
|
||||
}
|
||||
for elem in module.elements.iter() {
|
||||
match elem.kind {
|
||||
// Active segments are rooted because they initialize imported
|
||||
// tables.
|
||||
ElementKind::Active { table, .. } => {
|
||||
if module.tables.get(table).import.is_some() {
|
||||
stack.push_element(elem.id());
|
||||
}
|
||||
}
|
||||
// Declared segments can probably get gc'd but for now we're
|
||||
// conservative and we root them
|
||||
ElementKind::Declared => {
|
||||
stack.push_element(elem.id());
|
||||
}
|
||||
ElementKind::Passive => {}
|
||||
}
|
||||
}
|
||||
|
||||
// // And finally ask custom sections for their roots
|
||||
// for (_id, section) in module.customs.iter() {
|
||||
// section.add_gc_roots(&mut stack);
|
||||
// }
|
||||
// tracing::info!("Used roots: {:#?}", stack);
|
||||
|
||||
// Iteratively visit all items until our stack is empty
|
||||
while !stack.funcs.is_empty()
|
||||
|| !stack.tables.is_empty()
|
||||
|| !stack.memories.is_empty()
|
||||
|| !stack.globals.is_empty()
|
||||
|| !stack.data.is_empty()
|
||||
|| !stack.elements.is_empty()
|
||||
{
|
||||
while let Some((f, _loc)) = stack.funcs.pop() {
|
||||
if deleted.contains(&f) {
|
||||
let func = module.funcs.get(f);
|
||||
let name = func
|
||||
.name
|
||||
.as_ref()
|
||||
.cloned()
|
||||
.unwrap_or_else(|| format!("unknown - {}", f.index()));
|
||||
// panic!(
|
||||
// "Found a function that should be deleted but is still used: {:?} - {:?}",
|
||||
// name, f
|
||||
// );
|
||||
tracing::error!(
|
||||
"Found a function that should be deleted but is still used: {:?} - {:?} - {:?}",
|
||||
name,
|
||||
f,
|
||||
_loc
|
||||
);
|
||||
if let Location::Code { func } = _loc {
|
||||
let func_name = module.funcs.get(func).name.as_ref().unwrap();
|
||||
tracing::error!("Function {:?} is used by {:?}", f, func_name);
|
||||
}
|
||||
|
||||
// continue;
|
||||
}
|
||||
|
||||
let func = module.funcs.get(f);
|
||||
stack.used.types.insert(func.ty());
|
||||
|
||||
match &func.kind {
|
||||
FunctionKind::Local(func) => {
|
||||
let mut visitor = UsedVisitor {
|
||||
cur_func: f,
|
||||
stack: &mut stack,
|
||||
};
|
||||
dfs_in_order(&mut visitor, func, func.entry_block());
|
||||
}
|
||||
FunctionKind::Import(_) => {}
|
||||
FunctionKind::Uninitialized(_) => unreachable!(),
|
||||
}
|
||||
}
|
||||
|
||||
while let Some(t) = stack.tables.pop() {
|
||||
for elem in module.tables.get(t).elem_segments.iter() {
|
||||
stack.push_element(*elem);
|
||||
}
|
||||
}
|
||||
|
||||
while let Some(t) = stack.globals.pop() {
|
||||
match &module.globals.get(t).kind {
|
||||
GlobalKind::Import(_) => {}
|
||||
GlobalKind::Local(ConstExpr::Global(global)) => {
|
||||
stack.push_global(*global);
|
||||
}
|
||||
GlobalKind::Local(ConstExpr::RefFunc(func)) => {
|
||||
stack.push_func(*func, Location::Global { global: t });
|
||||
}
|
||||
GlobalKind::Local(ConstExpr::Value(_))
|
||||
| GlobalKind::Local(ConstExpr::RefNull(_)) => {}
|
||||
}
|
||||
}
|
||||
|
||||
while let Some(t) = stack.memories.pop() {
|
||||
for data in &module.memories.get(t).data_segments {
|
||||
stack.push_data(*data);
|
||||
}
|
||||
}
|
||||
|
||||
while let Some(d) = stack.data.pop() {
|
||||
let d = module.data.get(d);
|
||||
if let DataKind::Active { memory, offset } = &d.kind {
|
||||
stack.push_memory(*memory);
|
||||
if let ConstExpr::Global(g) = offset {
|
||||
stack.push_global(*g);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
while let Some(e) = stack.elements.pop() {
|
||||
let e = module.elements.get(e);
|
||||
if let ElementItems::Functions(function_ids) = &e.items {
|
||||
function_ids.iter().for_each(|f| {
|
||||
stack.push_func(*f, Location::Element { element: e.id() });
|
||||
});
|
||||
}
|
||||
if let ElementItems::Expressions(RefType::Funcref, items) = &e.items {
|
||||
for item in items {
|
||||
match item {
|
||||
ConstExpr::Global(g) => {
|
||||
stack.push_global(*g);
|
||||
}
|
||||
ConstExpr::RefFunc(f) => {
|
||||
stack.push_func(*f, Location::Element { element: e.id() });
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
if let ElementKind::Active { offset, table } = &e.kind {
|
||||
if let ConstExpr::Global(g) = offset {
|
||||
stack.push_global(*g);
|
||||
}
|
||||
stack.push_table(*table);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Wabt seems to have weird behavior where a `data` segment, if present
|
||||
// even if passive, requires a `memory` declaration. Our GC pass is
|
||||
// pretty aggressive and if you have a passive data segment and only
|
||||
// `data.drop` instructions you technically don't need the `memory`.
|
||||
// Let's keep `wabt` passing though and just say that if there are data
|
||||
// segments kept, but no memories, then we try to add the first memory,
|
||||
// if any, to the used set.
|
||||
if !stack.used.data.is_empty() && stack.used.memories.is_empty() {
|
||||
if let Some(mem) = module.memories.iter().next() {
|
||||
stack.used.memories.insert(mem.id());
|
||||
}
|
||||
}
|
||||
|
||||
stack.used
|
||||
}
|
||||
}
|
||||
|
||||
struct UsedVisitor<'a> {
|
||||
cur_func: FunctionId,
|
||||
stack: &'a mut Roots,
|
||||
}
|
||||
|
||||
impl Visitor<'_> for UsedVisitor<'_> {
|
||||
fn visit_function_id(&mut self, &func: &FunctionId) {
|
||||
self.stack.push_func(
|
||||
func,
|
||||
Location::Code {
|
||||
func: self.cur_func,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn visit_memory_id(&mut self, &m: &MemoryId) {
|
||||
self.stack.push_memory(m);
|
||||
}
|
||||
|
||||
fn visit_global_id(&mut self, &g: &GlobalId) {
|
||||
self.stack.push_global(g);
|
||||
}
|
||||
|
||||
fn visit_table_id(&mut self, &t: &TableId) {
|
||||
self.stack.push_table(t);
|
||||
}
|
||||
|
||||
fn visit_type_id(&mut self, &t: &TypeId) {
|
||||
self.stack.used.types.insert(t);
|
||||
}
|
||||
|
||||
fn visit_data_id(&mut self, &d: &DataId) {
|
||||
self.stack.push_data(d);
|
||||
}
|
||||
|
||||
fn visit_element_id(&mut self, &e: &ElementId) {
|
||||
self.stack.push_element(e);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user