feat: add option to strip exports except argument

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<String> {
+    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)

src/lib.rs

@@ -140,6 +140,12 @@ pub struct Options {
     /// file.
     pub patterns: Vec<String>,
 
+    /// The exports kept while others are snipped from the `.wasm`
+    pub kept_exports: Vec<String>,
+
+    /// The regex patterns of exports kept while others are snipped from the `.wasm`.
+    pub kept_export_patterns: Vec<String>,
+
     /// 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<String> = options.kept_exports.iter().cloned().collect();
+    let exports_to_delete = find_exports_to_delete(&module, &export_names, &re_kept_set);
+
     let names: HashSet<String> = 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<String>,
+    re_set: &regex::RegexSet,
+) -> HashSet<walrus::ExportId> {
+    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<walrus::ExportId>) {
+    for id in to_snip.iter().cloned() {
+        module.exports.delete(id)
+    }
+}
+
 fn delete_functions_to_snip(module: &mut walrus::Module, to_snip: &HashSet<walrus::FunctionId>) {
     for f in to_snip.iter().cloned() {
         module.funcs.delete(f);

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 {

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");
+}