Wasm-opt integration (#28)

Add option under `shrink` to optimize using `wasm-opt`

Author: kentosugama

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "ic-wasm"
-version = "0.3.5"
+version = "0.3.6"
 authors = ["DFINITY Stiftung"]
 edition = "2021"
 description = "A library for performing Wasm transformations specific to canisters running on the Internet Computer"
@@ -22,6 +22,8 @@ walrus = "0.19.0"
 candid = "0.9.0-beta.2"
 rustc-demangle = "0.1"
 thiserror = "1.0.35"
+wasm-opt = "0.112.0"
+tempfile = "3.5.0"
 
 anyhow = { version = "1.0.34", optional = true }
 clap = { version = "4.1", features = ["derive", "cargo"], optional = true }

Cargo.toml

@@ -17,7 +17,7 @@ Manage metadata in the Wasm module.
 Usage: `ic-wasm <input.wasm> [-o <output.wasm>] metadata [name] [-d <text content> | -f <file content>] [-v <public|private>]`
 
 * List current metadata sections
-``` 
+```
 $ ic-wasm input.wasm metadata
 ```
 
@@ -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
@@ -56,7 +78,7 @@ Usage: `ic-wasm <input.wasm> -o <output.wasm> resource --remove_cycles_transfer
 
 ### Instrument (experimental)
 
-Instrument canister method to emit execution trace to stable memory. 
+Instrument canister method to emit execution trace to stable memory.
 
 Usage: `ic-wasm <input.wasm> -o <output.wasm> instrument --trace-only func1 --trace-only func2`
 
@@ -81,7 +103,6 @@ Current limitations:
 * We cannot measure query calls.
 * No concurrent calls
 
-
 ## Library
 
 To use `ic-wasm` as a library, add this to your `Cargo.toml`:

README.md

@@ -48,8 +48,11 @@ 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)]
@@ -59,14 +62,20 @@ enum SubCommand {
 
 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::shrink_with_wasm_opt(&mut m, 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),

src/bin/main.rs

@@ -1,5 +1,6 @@
 use walrus::{IdsToIndices, Module, RawCustomSection};
 
+#[derive(Clone, Copy)]
 pub enum Kind {
     Public,
     Private,

src/metadata.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,63 @@ pub fn shrink(m: &mut Module) {
     }
     passes::gc::run(m);
 }
+
+pub fn shrink_with_wasm_opt(m: &mut Module, level: &str) -> anyhow::Result<()> {
+    // 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)?;
+            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<u8>)> = 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(),
+        _ => unreachable!(),
+    }
+    .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(), false)?;
+
+    // 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(())
+}

src/shrink.rs

@@ -21,6 +21,11 @@ function motoko(wasm) {
   call S.inc();
   call S.get();
   assert _ == (43 : nat);
+
+  call S.inc();
+  call S.inc();
+  call S.get();
+  assert _ == (45 : nat);
   S
 };
 function rust(wasm) {
@@ -29,6 +34,11 @@ function rust(wasm) {
   call S.inc();
   call S.read();
   assert _ == (43 : nat);
+
+  call S.inc();
+  call S.inc();
+  call S.read();
+  assert _ == (45 : nat);
   S
 };
 function wat(wasm) {
@@ -37,6 +47,11 @@ function wat(wasm) {
   call S.inc();
   call S.get();
   assert _ == (43 : int64);
+
+  call S.inc();
+  call S.inc();
+  call S.get();
+  assert _ == (45 : int64);
   S
 };
 function classes(wasm) {
@@ -46,6 +61,12 @@ function classes(wasm) {
   call S.put(42, "text");
   call S.get(42);
   assert _ == opt "text";
+
+  call S.put(40, "text0");
+  call S.put(41, "text1");
+  call S.put(42, "text2");
+  call S.get(42);
+  assert _ == opt "text2";
   S
 };
 function classes_limit(wasm) {
@@ -67,26 +88,31 @@ function classes_redirect(wasm) {
 
 let S = motoko(file("ok/motoko-instrument.wasm"));
 call S.__get_cycles();
-assert _ == (7199 : int64);
+assert _ == (9003 : int64);
 let S = motoko(file("ok/motoko-gc-instrument.wasm"));
 call S.__get_cycles();
-assert _ == (177 : int64);
+assert _ == (295 : int64);
 motoko(file("ok/motoko-shrink.wasm"));
 motoko(file("ok/motoko-limit.wasm"));
 
 let S = rust(file("ok/rust-instrument.wasm"));
 call S.__get_cycles();
-assert _ == (66016 : int64);
+assert _ == (136378 : int64);
 rust(file("ok/rust-shrink.wasm"));
 rust(file("ok/rust-limit.wasm"));
 
 let S = wat(file("ok/wat-instrument.wasm"));
 call S.__get_cycles();
-assert _ == (121 : int64);
+assert _ == (189 : int64);
 wat(file("ok/wat-shrink.wasm"));
 wat(file("ok/wat-limit.wasm"));
 
 classes(file("ok/classes-shrink.wasm"));
 classes_limit(file("ok/classes-limit.wasm"));
 classes_redirect(file("ok/classes-redirect.wasm"));
 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"));