Aggressive inlining options for wasm-opt

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"

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<String>,
+        #[clap(short, long)]
+        keep_name_section: bool,
+    },
+    /// Optimize the Wasm module using wasm-opt
+    #[cfg(feature = "wasm-opt")]
+    Optimize {
+        #[clap()]
+        level: ic_wasm::shrink::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<u32>,
         #[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::shrink::shrink_with_wasm_opt(
+            &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),

src/shrink.rs

@@ -1,5 +1,6 @@
 use crate::metadata::*;
 use crate::utils::*;
+use clap::ValueEnum;
 use walrus::*;
 
 pub fn shrink(m: &mut Module) {
@@ -26,10 +27,30 @@ pub fn shrink(m: &mut Module) {
     passes::gc::run(m);
 }
 
+#[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,
+}
+
 #[cfg(feature = "wasm-opt")]
 pub fn shrink_with_wasm_opt(
     m: &mut Module,
-    level: &str,
+    level: &OptLevel,
+    inline_functions_with_loops: bool,
+    always_inline_max_function_size: &Option<u32>,
     keep_name_section: bool,
 ) -> anyhow::Result<()> {
     use tempfile::NamedTempFile;
@@ -38,7 +59,13 @@ pub fn shrink_with_wasm_opt(
     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)?;
+            shrink_with_wasm_opt(
+                &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;
         }
@@ -67,18 +94,21 @@ pub fn shrink_with_wasm_opt(
         .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"),
+    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);
     }
-    .debug_info(keep_name_section)
-    .run(temp_file.path(), temp_file.path())?;
+    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)?;

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()