move into shrink and remove test cases

README.md

@@ -44,10 +44,32 @@ Usage: `ic-wasm <input.wasm> info`
 
 ### Shrink
 
-Remove unused functions and debug info
+Remove unused functions and debug info.
 
 Usage: `ic-wasm <input.wasm> -o <output.wasm> shrink`
 
+Optionally invoke wasm optimizations from [`wasm-opt`](https://github.com/WebAssembly/binaryen).
+
+The optimizer exposes different optimization levels to choose from.
+
+Performance levels (optimizes for runtime):
+- O4
+- O3 (default setting: best for minimizing cycle usage)
+- O2
+- O1
+- O0 (no optimizations)
+
+Code size levels (optimizes for binary size):
+- Oz (best for minimizing code size)
+- Os
+
+The recommended setting (O3) reduces cycle usage for Motoko programs by ~10% and  Rust programs by ~4%. The code size for both languages is reduced by ~16%.
+
+Note: The `icp` metadata sections are preserved through the optimizations.
+
+
+Usage: `ic-wasm <input.wasm> -o <output.wasm> shrink --optimize <level>`
+
 ### Resource
 
 Limit resource usage, mainly used by Motoko Playground
@@ -81,30 +103,6 @@ Current limitations:
 * We cannot measure query calls.
 * No concurrent calls
 
-### Optimize
-
-Invoke wasm optimizations from [`wasm-opt`](https://github.com/WebAssembly/binaryen).
-
-The optimizer exposes different optimization levels to choose from.
-
-Performance levels (optimizes for runtime):
-- O4
-- O3 (default setting: best for minimizing cycle usage)
-- O2
-- O1
-- O0 (no optimizations)
-
-Code size levels (optimizes for binary size):
-- Oz (best for minimizing code size)
-- Os
-
-The recommended setting (O3) reduces cycle usage for Motoko programs by ~10% and  Rust programs by ~4%. The code size for both languages is reduced by ~16%.
-
-Note: The `icp` metadata sections are preserved through the optimizations.
-
-
-Usage: `ic-wasm <input.wasm> -o <output.wasm> optimize --level <optional level>`
-
 ## Library
 
 To use `ic-wasm` as a library, add this to your `Cargo.toml`:

src/bin/main.rs

@@ -48,30 +48,34 @@ enum SubCommand {
     },
     /// List information about the Wasm canister
     Info,
-    /// Remove unused functions and debug info
-    Shrink,
+    /// Remove unused functions and debug info. Optionally, optimize the Wasm code
+    Shrink {
+        #[clap(short, long, value_parser = ["O0", "O1", "O2", "O3", "O4", "Os", "Oz"])]
+        optimize: Option<String>,
+    },
     /// Instrument canister method to emit execution trace to stable memory (experimental)
     Instrument {
         #[clap(short, long)]
         trace_only: Option<Vec<String>>,
     },
-    /// Invoke wasm optimizations from wasm-opt
-    Optimize {
-        #[clap(short, long, value_parser = ["O0", "O1", "O2", "O3", "O4", "Os", "Oz"], default_value = "O3")]
-        level: String,
-    },
 }
 
 fn main() -> anyhow::Result<()> {
     let opts: Opts = Opts::parse();
-    let keep_name_section = matches!(opts.subcommand, SubCommand::Shrink);
+    let keep_name_section = matches!(opts.subcommand, SubCommand::Shrink { .. });
     let mut m = ic_wasm::utils::parse_wasm_file(opts.input, keep_name_section)?;
     match &opts.subcommand {
         SubCommand::Info => {
             let mut stdout = std::io::stdout();
             ic_wasm::info::info(&m, &mut stdout)?;
         }
-        SubCommand::Shrink => ic_wasm::shrink::shrink(&mut m),
+        SubCommand::Shrink { optimize } => {
+            use ic_wasm::shrink;
+            match optimize {
+                Some(level) => shrink::optimize(&mut m, keep_name_section, level),
+                None => shrink::shrink(&mut m),
+            }
+        }
         SubCommand::Instrument { trace_only } => match trace_only {
             None => ic_wasm::instrumentation::instrument(&mut m, &[]),
             Some(vec) => ic_wasm::instrumentation::instrument(&mut m, vec),
@@ -126,10 +130,6 @@ fn main() -> anyhow::Result<()> {
                 return Ok(());
             }
         }
-        SubCommand::Optimize { level } => {
-            use ic_wasm::optimize::optimize;
-            optimize(&mut m, keep_name_section, level)
-        }
     };
     if let Some(output) = opts.output {
         m.emit_wasm_file(output)?;

src/lib.rs

@@ -2,7 +2,6 @@ pub mod info;
 pub mod instrumentation;
 pub mod limit_resource;
 pub mod metadata;
-pub mod optimize;
 pub mod shrink;
 pub mod utils;
 

src/shrink.rs

@@ -1,6 +1,8 @@
-use walrus::*;
-
+use crate::metadata::*;
 use crate::utils::*;
+use tempfile::NamedTempFile;
+use walrus::*;
+use wasm_opt::OptimizationOptions;
 
 pub fn shrink(m: &mut Module) {
     if is_motoko_canister(m) {
@@ -25,3 +27,77 @@ pub fn shrink(m: &mut Module) {
     }
     passes::gc::run(m);
 }
+
+pub fn optimize(m: &mut Module, keep_name_section: bool, level: &str) {
+    // 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, keep_name_section, level);
+            let blob = encode_module_as_data_section(module);
+            m.data.get_mut(id).value = blob;
+        }
+    }
+
+    // pull out a copy of the custom sections to preserve
+    let m_copy = parse_wasm(&m.emit_wasm(), keep_name_section).unwrap();
+    let mut metadata_sections = Vec::new();
+    list_metadata(&m_copy).iter().for_each(|full_name| {
+        match full_name.strip_prefix("icp:public ") {
+            Some(name) => {
+                metadata_sections.push(("public", name, get_metadata(&m_copy, name).unwrap()))
+            }
+            None => match full_name.strip_prefix("icp:private ") {
+                Some(name) => {
+                    metadata_sections.push(("private", name, get_metadata(&m_copy, name).unwrap()))
+                }
+                None => unreachable!(),
+            },
+        };
+    });
+
+    // write to temp file
+    let temp_file = NamedTempFile::new().unwrap_or_else(|e| {
+        eprintln!("unable to create temp file: {}", e);
+        std::process::exit(1);
+    });
+    m.emit_wasm_file(temp_file.path()).unwrap_or_else(|e| {
+        eprintln!("unable to write to temp file: {}", e);
+        std::process::exit(1);
+    });
+
+    // 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(),
+        _ => unreachable!(),
+    }
+    .run(temp_file.path(), temp_file.path())
+    .unwrap_or_else(|e| {
+        eprintln!("unable to optimize wasm: {}", e);
+        std::process::exit(1);
+    });
+
+    // read optimized wasm back in from temp file
+    *m = parse_wasm_file(temp_file.path().to_path_buf(), keep_name_section).unwrap_or_else(|e| {
+        eprintln!("unable to read optimized wasm: {}", e);
+        std::process::exit(1);
+    });
+
+    // re-insert the custom sections
+    metadata_sections
+        .iter()
+        .for_each(|(visibility, name, data)| {
+            let visibility = match *visibility {
+                "public" => Kind::Public,
+                "private" => Kind::Private,
+                _ => unreachable!(),
+            };
+            add_metadata(m, visibility, name, data.to_vec());
+        });
+}

tests/deployable.ic-repl.sh

@@ -118,12 +118,4 @@ classes(file("ok/classes-nop-redirect.wasm"));
 motoko(file("ok/motoko-optimize.wasm"));
 rust(file("ok/rust-optimize.wasm"));
 wat(file("ok/wat-optimize.wasm"));
-classes(file("ok/classes-optimize.wasm"));
-motoko(file("ok/motoko-optimize-level-4.wasm"));
-rust(file("ok/rust-optimize-level-4.wasm"));
-wat(file("ok/wat-optimize-level-4.wasm"));
-classes(file("ok/classes-optimize-level-4.wasm"));
-motoko(file("ok/motoko-optimize-level-z.wasm"));
-rust(file("ok/rust-optimize-level-z.wasm"));
-wat(file("ok/wat-optimize-level-z.wasm"));
-classes(file("ok/classes-optimize-level-z.wasm"));
+classes(file("ok/classes-optimize.wasm"));