diff --git a/Cargo.lock b/Cargo.lock index 00ae807..ff3a6fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -468,7 +468,7 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" [[package]] name = "ic-wasm" -version = "0.4.0" +version = "0.5.0" dependencies = [ "anyhow", "assert_cmd", diff --git a/Cargo.toml b/Cargo.toml index a9fec45..71a10e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ic-wasm" -version = "0.4.0" +version = "0.5.0" authors = ["DFINITY Stiftung"] edition = "2021" description = "A library for performing Wasm transformations specific to canisters running on the Internet Computer" diff --git a/src/bin/main.rs b/src/bin/main.rs index 14ae389..3ad64cc 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -48,10 +48,20 @@ enum SubCommand { }, /// List information about the Wasm canister Info, - /// Remove unused functions and debug info. Optionally, optimize the Wasm code + /// Remove unused functions and debug info Shrink { - #[clap(long, value_parser = ["O0", "O1", "O2", "O3", "O4", "Os", "Oz"])] - optimize: Option, + #[clap(short, long)] + keep_name_section: bool, + }, + /// Optimize the Wasm module using wasm-opt + #[cfg(feature = "wasm-opt")] + Optimize { + #[clap()] + level: ic_wasm::optimize::OptLevel, + #[clap(long("inline-functions-with-loops"))] + inline_functions_with_loops: bool, + #[clap(long("always-inline-max-function-size"))] + always_inline_max_function_size: Option, #[clap(short, long)] keep_name_section: bool, }, @@ -65,7 +75,9 @@ enum SubCommand { fn main() -> anyhow::Result<()> { let opts: Opts = Opts::parse(); let keep_name_section = match opts.subcommand { - SubCommand::Shrink { + SubCommand::Shrink { keep_name_section } => keep_name_section, + #[cfg(feature = "wasm-opt")] + SubCommand::Optimize { keep_name_section, .. } => keep_name_section, _ => false, @@ -76,18 +88,20 @@ fn main() -> anyhow::Result<()> { let mut stdout = std::io::stdout(); ic_wasm::info::info(&m, &mut stdout)?; } - SubCommand::Shrink { optimize, .. } => { - use ic_wasm::shrink; - match optimize { - Some(level) => { - #[cfg(not(feature = "wasm-opt"))] - panic!("Please build with wasm-opt feature"); - #[cfg(feature = "wasm-opt")] - shrink::shrink_with_wasm_opt(&mut m, level, keep_name_section)? - } - None => shrink::shrink(&mut m), - } - } + SubCommand::Shrink { .. } => ic_wasm::shrink::shrink(&mut m), + #[cfg(feature = "wasm-opt")] + SubCommand::Optimize { + level, + inline_functions_with_loops, + always_inline_max_function_size, + .. + } => ic_wasm::optimize::optimize( + &mut m, + level, + *inline_functions_with_loops, + always_inline_max_function_size, + keep_name_section, + )?, SubCommand::Instrument { trace_only } => match trace_only { None => ic_wasm::instrumentation::instrument(&mut m, &[]), Some(vec) => ic_wasm::instrumentation::instrument(&mut m, vec), diff --git a/src/lib.rs b/src/lib.rs index 6c067b1..9df083c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,8 @@ pub mod info; pub mod instrumentation; pub mod limit_resource; pub mod metadata; +#[cfg(feature = "wasm-opt")] +pub mod optimize; pub mod shrink; pub mod utils; diff --git a/src/optimize.rs b/src/optimize.rs new file mode 100644 index 0000000..605c869 --- /dev/null +++ b/src/optimize.rs @@ -0,0 +1,100 @@ +use crate::metadata::*; +use crate::utils::*; +use clap::ValueEnum; +use walrus::*; + +#[derive(Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum OptLevel { + #[clap(name = "O0")] + O0, + #[clap(name = "O1")] + O1, + #[clap(name = "O2")] + O2, + #[clap(name = "O3")] + O3, + #[clap(name = "O4")] + O4, + #[clap(name = "Os")] + Os, + #[clap(name = "Oz")] + Oz, +} + +pub fn optimize( + m: &mut Module, + level: &OptLevel, + inline_functions_with_loops: bool, + always_inline_max_function_size: &Option, + keep_name_section: bool, +) -> anyhow::Result<()> { + use tempfile::NamedTempFile; + use wasm_opt::OptimizationOptions; + // recursively optimize embedded modules in Motoko actor classes + if is_motoko_canister(m) { + let data = get_motoko_wasm_data_sections(m); + for (id, mut module) in data.into_iter() { + optimize( + &mut module, + level, + inline_functions_with_loops, + always_inline_max_function_size, + keep_name_section, + )?; + let blob = encode_module_as_data_section(module); + m.data.get_mut(id).value = blob; + } + } + + // write module to temp file + let temp_file = NamedTempFile::new()?; + m.emit_wasm_file(temp_file.path())?; + + // pull out a copy of the custom sections to preserve + let metadata_sections: Vec<(Kind, &str, Vec)> = m + .customs + .iter() + .filter(|(_, section)| section.name().starts_with("icp:")) + .map(|(_, section)| { + let data = section.data(&IdsToIndices::default()).to_vec(); + let full_name = section.name(); + match full_name.strip_prefix("public ") { + Some(name) => (Kind::Public, name, data), + None => match full_name.strip_prefix("private ") { + Some(name) => (Kind::Private, name, data), + None => unreachable!(), + }, + } + }) + .collect(); + + // read in from temp file and optimize + let mut optimizations = match level { + OptLevel::O0 => OptimizationOptions::new_opt_level_0(), + OptLevel::O1 => OptimizationOptions::new_opt_level_1(), + OptLevel::O2 => OptimizationOptions::new_opt_level_2(), + OptLevel::O3 => OptimizationOptions::new_opt_level_3(), + OptLevel::O4 => OptimizationOptions::new_opt_level_4(), + OptLevel::Os => OptimizationOptions::new_optimize_for_size(), + OptLevel::Oz => OptimizationOptions::new_optimize_for_size_aggressively(), + }; + optimizations.debug_info(keep_name_section); + optimizations.allow_functions_with_loops(inline_functions_with_loops); + if let Some(max_size) = always_inline_max_function_size { + optimizations.always_inline_max_size(*max_size); + } + optimizations.run(temp_file.path(), temp_file.path())?; + + // read optimized wasm back in from temp file + let mut m_opt = parse_wasm_file(temp_file.path().to_path_buf(), keep_name_section)?; + + // re-insert the custom sections + metadata_sections + .into_iter() + .for_each(|(visibility, name, data)| { + add_metadata(&mut m_opt, visibility, name, data); + }); + + *m = m_opt; + Ok(()) +} diff --git a/src/shrink.rs b/src/shrink.rs index c85f7de..0e2e6b6 100644 --- a/src/shrink.rs +++ b/src/shrink.rs @@ -1,4 +1,3 @@ -use crate::metadata::*; use crate::utils::*; use walrus::*; @@ -25,71 +24,3 @@ pub fn shrink(m: &mut Module) { } passes::gc::run(m); } - -#[cfg(feature = "wasm-opt")] -pub fn shrink_with_wasm_opt( - m: &mut Module, - level: &str, - keep_name_section: bool, -) -> anyhow::Result<()> { - use tempfile::NamedTempFile; - use wasm_opt::OptimizationOptions; - // recursively optimize embedded modules in Motoko actor classes - if is_motoko_canister(m) { - let data = get_motoko_wasm_data_sections(m); - for (id, mut module) in data.into_iter() { - shrink_with_wasm_opt(&mut module, level, keep_name_section)?; - let blob = encode_module_as_data_section(module); - m.data.get_mut(id).value = blob; - } - } - - // write module to temp file - let temp_file = NamedTempFile::new()?; - m.emit_wasm_file(temp_file.path())?; - - // pull out a copy of the custom sections to preserve - let metadata_sections: Vec<(Kind, &str, Vec)> = m - .customs - .iter() - .filter(|(_, section)| section.name().starts_with("icp:")) - .map(|(_, section)| { - let data = section.data(&IdsToIndices::default()).to_vec(); - let full_name = section.name(); - match full_name.strip_prefix("public ") { - Some(name) => (Kind::Public, name, data), - None => match full_name.strip_prefix("private ") { - Some(name) => (Kind::Private, name, data), - None => unreachable!(), - }, - } - }) - .collect(); - - // read in from temp file and optimize - match level { - "O0" => OptimizationOptions::new_opt_level_0(), - "O1" => OptimizationOptions::new_opt_level_1(), - "O2" => OptimizationOptions::new_opt_level_2(), - "O3" => OptimizationOptions::new_opt_level_3(), - "O4" => OptimizationOptions::new_opt_level_4(), - "Os" => OptimizationOptions::new_optimize_for_size(), - "Oz" => OptimizationOptions::new_optimize_for_size_aggressively(), - _ => anyhow::bail!("invalid optimization level"), - } - .debug_info(keep_name_section) - .run(temp_file.path(), temp_file.path())?; - - // read optimized wasm back in from temp file - let mut m_opt = parse_wasm_file(temp_file.path().to_path_buf(), keep_name_section)?; - - // re-insert the custom sections - metadata_sections - .into_iter() - .for_each(|(visibility, name, data)| { - add_metadata(&mut m_opt, visibility, name, data); - }); - - *m = m_opt; - Ok(()) -} diff --git a/tests/tests.rs b/tests/tests.rs index dc68565..e846f5a 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -62,12 +62,6 @@ fn instrumentation() { #[test] fn shrink() { - let expected_metadata = r#"icp:public candid:service -icp:private candid:args -icp:private motoko:stable-types -icp:private motoko:compiler -"#; - wasm_input("motoko.wasm", true) .arg("shrink") .assert() @@ -88,10 +82,18 @@ icp:private motoko:compiler .assert() .success(); assert_wasm("classes-shrink.wasm"); +} + +#[test] +fn optimize() { + let expected_metadata = r#"icp:public candid:service +icp:private candid:args +icp:private motoko:stable-types +icp:private motoko:compiler +"#; wasm_input("classes.wasm", true) - .arg("shrink") - .arg("--optimize") + .arg("optimize") .arg("O3") .assert() .success(); @@ -102,8 +104,7 @@ icp:private motoko:compiler .stdout(expected_metadata) .success(); wasm_input("classes.wasm", true) - .arg("shrink") - .arg("--optimize") + .arg("optimize") .arg("O3") .arg("--keep-name-section") .assert()