diff --git a/src/bin/wasm-snip.rs b/src/bin/wasm-snip.rs index 12d38d7..b33aee7 100644 --- a/src/bin/wasm-snip.rs +++ b/src/bin/wasm-snip.rs @@ -1,3 +1,4 @@ +use clap::ArgMatches; use failure::ResultExt; use std::fs; use std::io::{self, Write}; @@ -14,32 +15,26 @@ fn main() { } } +fn get_values(matches: &ArgMatches, name: &str) -> Vec { + matches + .values_of(name) + .map(|fs| fs.map(|f| f.to_string()).collect()) + .unwrap_or(vec![]) +} + fn try_main() -> Result<(), failure::Error> { let matches = parse_args(); let mut opts = wasm_snip::Options::default(); - opts.functions = matches - .values_of("function") - .map(|fs| fs.map(|f| f.to_string()).collect()) - .unwrap_or(vec![]); - - opts.patterns = matches - .values_of("pattern") - .map(|ps| ps.map(|p| p.to_string()).collect()) - .unwrap_or(vec![]); + opts.functions = get_values(&matches, "function"); + opts.patterns = get_values(&matches, "pattern"); + opts.kept_exports = get_values(&matches, "kept_export"); + opts.kept_export_patterns = get_values(&matches, "kept_export_pattern"); - if matches.is_present("snip_rust_fmt_code") { - opts.snip_rust_fmt_code = true; - } - - if matches.is_present("snip_rust_panicking_code") { - opts.snip_rust_panicking_code = true; - } - - if matches.is_present("skip_producers_section") { - opts.skip_producers_section = true; - } + opts.snip_rust_fmt_code = matches.is_present("snip_rust_fmt_code"); + opts.snip_rust_panicking_code = matches.is_present("snip_rust_panicking_code"); + opts.skip_producers_section = matches.is_present("skip_producers_section"); let config = walrus_config_from_options(&opts); let path = matches.value_of("input").unwrap(); @@ -114,6 +109,24 @@ Very helpful when shrinking the size of WebAssembly binaries! .takes_value(true) .help("Snip any function that matches the given regular expression."), ) + .arg( + clap::Arg::with_name("kept_export") + .required(false) + .multiple(true) + .short("k") + .long("kept-export") + .takes_value(true) + .help("Snip exports not included. Matches exactly."), + ) + .arg( + clap::Arg::with_name("kept_export_pattern") + .required(false) + .multiple(true) + .short("x") + .long("kept-export-pattern") + .takes_value(true) + .help("Snip exports not included. Matches exports with regular expression."), + ) .arg( clap::Arg::with_name("snip_rust_fmt_code") .required(false) diff --git a/src/lib.rs b/src/lib.rs index c475caa..12a8e88 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -140,6 +140,12 @@ pub struct Options { /// file. pub patterns: Vec, + /// The exports kept while others are snipped from the `.wasm` + pub kept_exports: Vec, + + /// The regex patterns of exports kept while others are snipped from the `.wasm`. + pub kept_export_patterns: Vec, + /// Should Rust `std::fmt` and `core::fmt` functions be snipped? pub snip_rust_fmt_code: bool, @@ -159,10 +165,15 @@ pub fn snip(module: &mut walrus::Module, options: Options) -> Result<(), failure .add_processed_by("wasm-snip", env!("CARGO_PKG_VERSION")); } + let re_kept_set = regex::RegexSet::new(&options.kept_export_patterns)?; + let export_names: HashSet = options.kept_exports.iter().cloned().collect(); + let exports_to_delete = find_exports_to_delete(&module, &export_names, &re_kept_set); + let names: HashSet = options.functions.iter().cloned().collect(); let re_set = build_regex_set(options).context("failed to compile regex")?; let to_snip = find_functions_to_snip(&module, &names, &re_set); + delete_exports(module, &exports_to_delete); replace_calls_with_unreachable(module, &to_snip); unexport_snipped_functions(module, &to_snip); unimport_snipped_functions(module, &to_snip); @@ -227,6 +238,37 @@ fn find_functions_to_snip( .collect() } +fn is_function(export: &walrus::Export) -> Option<&walrus::Export> { + match export.item { + walrus::ExportItem::Function(_) => Some(export), + _ => None, + } +} +fn find_exports_to_delete( + module: &walrus::Module, + names: &HashSet, + re_set: ®ex::RegexSet, +) -> HashSet { + module + .exports + .iter() + .filter_map(is_function) + .filter_map(|e| { + if names.contains(&e.name) || re_set.is_match(&e.name) { + None + } else { + Some(e.id()) + } + }) + .collect() +} + +fn delete_exports(module: &mut walrus::Module, to_snip: &HashSet) { + for id in to_snip.iter().cloned() { + module.exports.delete(id) + } +} + fn delete_functions_to_snip(module: &mut walrus::Module, to_snip: &HashSet) { for f in to_snip.iter().cloned() { module.funcs.delete(f); diff --git a/tests/hello.rs b/tests/hello.rs index e3195b9..a7b6e53 100755 --- a/tests/hello.rs +++ b/tests/hello.rs @@ -10,6 +10,17 @@ pub fn fluxions(x: usize) -> usize { } } +#[no_mangle] +pub fn keep_me() -> usize { + fluxions(1) +} + +#[no_mangle] +pub fn keep_me_too() -> usize { + let x = 1; + return x; +} + #[no_mangle] pub fn quicksilver(x: usize) { if x > 100 { diff --git a/tests/hello.wasm b/tests/hello.wasm old mode 100644 new mode 100755 index 7eea55f..8ffbcdc Binary files a/tests/hello.wasm and b/tests/hello.wasm differ diff --git a/tests/kept_me.wasm b/tests/kept_me.wasm new file mode 100644 index 0000000..125a79c Binary files /dev/null and b/tests/kept_me.wasm differ diff --git a/tests/kept_me_too.wasm b/tests/kept_me_too.wasm new file mode 100644 index 0000000..b746b6e Binary files /dev/null and b/tests/kept_me_too.wasm differ diff --git a/tests/no_alloc.wasm b/tests/no_alloc.wasm index 861ccfb..9cb1603 100644 Binary files a/tests/no_alloc.wasm and b/tests/no_alloc.wasm differ diff --git a/tests/no_fmt.wasm b/tests/no_fmt.wasm index a439fe8..9cb1603 100644 Binary files a/tests/no_fmt.wasm and b/tests/no_fmt.wasm differ diff --git a/tests/no_panicking.wasm b/tests/no_panicking.wasm index 2bffda5..9cb1603 100644 Binary files a/tests/no_panicking.wasm and b/tests/no_panicking.wasm differ diff --git a/tests/snip_me.wasm b/tests/snip_me.wasm index 622f57e..9cb1603 100644 Binary files a/tests/snip_me.wasm and b/tests/snip_me.wasm differ diff --git a/tests/tests.rs b/tests/tests.rs index b7382ff..00113fb 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -83,3 +83,13 @@ fn snip_rust_panicking_code() { "no_panicking.wasm", ); } + +#[test] +fn keep_exports() { + assert_snip(wasm_snip().arg("-k").arg("keep_me"), "kept_me.wasm"); +} + +#[test] +fn keep_export_patterns() { + assert_snip(wasm_snip().arg("-x").arg("keep_me"), "kept_me_too.wasm"); +}