Expose optimize feature

docs/dfx-json-schema.json

@@ -415,6 +415,14 @@
             "boolean",
             "null"
           ]
+        },
+        "optimize": {
+          "title": "Optimize Canister WASM",
+          "description": "Whether run `ic-wasm shrink --optimize <level>` after building the Canister. Disabled by default.",
+          "type": [
+            "string",
+            "null"
+          ]
         }
       }
     },
@@ -858,4 +866,4 @@
       ]
     }
   }
-}
+}

src/dfx-core/src/config/model/dfinity.rs

@@ -195,6 +195,11 @@ pub struct ConfigCanistersCanister {
     /// Disabled by default for custom canisters.
     pub shrink: Option<bool>,
 
+    /// # Optimize Canister WASM
+    /// Whether run `ic-wasm shrink --optimize <level>` after building the Canister.
+    /// Disabled by default.
+    pub optimize: Option<String>,
+
     /// # Metadata
     /// Defines metadata sections to set in the canister .wasm
     #[serde(default)]

src/dfx/src/lib/builders/custom.rs

@@ -146,11 +146,20 @@ impl CanisterBuilder for CustomBuilder {
             }
         }
 
+        let optimize = info.get_optimize();
         let shrink = info.get_shrink().unwrap_or(false);
         // Custom canister may have WASM gzipped
-        if shrink && is_wasm_format(&wasm)? {
-            info!(self.logger, "Shrink WASM module size.");
-            super::shrink_wasm(&wasm)?;
+        if is_wasm_format(&wasm)? {
+            if let Some(level) = optimize {
+                info!(
+                    self.logger,
+                    "Optimize and shrink WASM module at level {}", level
+                );
+                super::optimize_wasm(&wasm, level)?;
+            } else if shrink {
+                info!(self.logger, "Shrink WASM module size.");
+                super::shrink_wasm(&wasm)?;
+            }
         }
 
         Ok(BuildOutput {

src/dfx/src/lib/builders/mod.rs

@@ -491,6 +491,20 @@ fn shrink_wasm(wasm_path: impl AsRef<Path>) -> DfxResult {
     Ok(())
 }
 
+#[context("Failed to optimize wasm at {}.", &wasm_path.as_ref().display())]
+fn optimize_wasm(wasm_path: impl AsRef<Path>, level: &str) -> DfxResult {
+    let wasm_path = wasm_path.as_ref();
+    let wasm = std::fs::read(wasm_path).context("Could not read the WASM module.")?;
+    let mut module =
+        ic_wasm::utils::parse_wasm(&wasm, true).context("Could not parse the WASM module.")?;
+    ic_wasm::shrink::shrink_with_wasm_opt(&mut module, level)
+        .context("Could not optimize the WASM module.")?;
+    module
+        .emit_wasm_file(wasm_path)
+        .with_context(|| format!("Could not write optimized WASM to {:?}.", wasm_path))?;
+    Ok(())
+}
+
 pub struct BuilderPool {
     builders: BTreeMap<&'static str, Arc<dyn CanisterBuilder>>,
 }

src/dfx/src/lib/builders/motoko.rs

@@ -166,8 +166,15 @@ impl CanisterBuilder for MotokoBuilder {
         };
         motoko_compile(&self.logger, cache.as_ref(), &params)?;
 
+        let optimize = canister_info.get_optimize();
         let shrink = canister_info.get_shrink().unwrap_or(true);
-        if shrink {
+        if let Some(level) = optimize {
+            info!(
+                self.logger,
+                "Optimize and shrink WASM module at level {}", level
+            );
+            super::optimize_wasm(motoko_info.get_output_wasm_path(), level)?;
+        } else if shrink {
             info!(self.logger, "Shrink WASM module size.");
             super::shrink_wasm(motoko_info.get_output_wasm_path())?;
         }

src/dfx/src/lib/builders/rust.rs

@@ -97,8 +97,15 @@ impl CanisterBuilder for RustBuilder {
         );
         let output = cargo.output().context("Failed to run 'cargo build'. You might need to run `cargo update` (or a similar command like `cargo vendor`) if you have updated `Cargo.toml`, because `dfx build` uses the --locked flag with Cargo.")?;
 
+        let optimize = canister_info.get_optimize();
         let shrink = canister_info.get_shrink().unwrap_or(true);
-        if shrink {
+        if let Some(level) = optimize {
+            info!(
+                self.logger,
+                "Optimize and shrink WASM module at level {}", level
+            );
+            super::optimize_wasm(rust_info.get_output_wasm_path(), level)?;
+        } else if shrink {
             info!(self.logger, "Shrink WASM module size.");
             super::shrink_wasm(rust_info.get_output_wasm_path())?;
         }

src/dfx/src/lib/canister_info.rs

@@ -52,6 +52,7 @@ pub struct CanisterInfo {
     post_install: Vec<String>,
     main: Option<PathBuf>,
     shrink: Option<bool>,
+    optimize: Option<String>,
     metadata: CanisterMetadataConfig,
     pull_ready: bool,
     pull_dependencies: Vec<(String, CanisterId)>,
@@ -154,6 +155,7 @@ impl CanisterInfo {
             post_install,
             main: canister_config.main.clone(),
             shrink: canister_config.shrink,
+            optimize: canister_config.optimize.clone(),
             metadata,
             pull_ready: canister_config.pull_ready,
             pull_dependencies,
@@ -229,6 +231,10 @@ impl CanisterInfo {
         self.shrink
     }
 
+    pub fn get_optimize(&self) -> &Option<String> {
+        &self.optimize
+    }
+
     pub fn get_build_wasm_path(&self) -> PathBuf {
         self.output_root.join(&self.name).with_extension("wasm")
     }