ic_wasm/
optimize.rs

1use crate::metadata::*;
2use crate::utils::*;
3use clap::ValueEnum;
4use walrus::*;
5
6#[derive(Clone, Copy, PartialEq, Eq, ValueEnum)]
7pub enum OptLevel {
8    #[clap(name = "O0")]
9    O0,
10    #[clap(name = "O1")]
11    O1,
12    #[clap(name = "O2")]
13    O2,
14    #[clap(name = "O3")]
15    O3,
16    #[clap(name = "O4")]
17    O4,
18    #[clap(name = "Os")]
19    Os,
20    #[clap(name = "Oz")]
21    Oz,
22}
23
24pub fn optimize(
25    m: &mut Module,
26    level: &OptLevel,
27    inline_functions_with_loops: bool,
28    always_inline_max_function_size: &Option<u32>,
29    keep_name_section: bool,
30) -> anyhow::Result<()> {
31    use tempfile::NamedTempFile;
32    use wasm_opt::OptimizationOptions;
33    // recursively optimize embedded modules in Motoko actor classes
34    if is_motoko_canister(m) {
35        let data = get_motoko_wasm_data_sections(m);
36        for (id, mut module) in data.into_iter() {
37            let old_size = module.emit_wasm().len();
38            optimize(
39                &mut module,
40                level,
41                inline_functions_with_loops,
42                always_inline_max_function_size,
43                keep_name_section,
44            )?;
45            let new_size = module.emit_wasm().len();
46            // Guard against embedded actor class overriding the parent module
47            if new_size <= old_size {
48                let blob = encode_module_as_data_section(module);
49                m.data.get_mut(id).value = blob;
50            } else {
51                eprintln!("Warning: embedded actor class module was not optimized because the optimized module is larger than the original module");
52            }
53        }
54    }
55
56    // write module to temp file
57    let temp_file = NamedTempFile::new()?;
58    m.emit_wasm_file(temp_file.path())?;
59
60    // pull out a copy of the custom sections to preserve
61    let metadata_sections: Vec<(Kind, &str, Vec<u8>)> = m
62        .customs
63        .iter()
64        .filter(|(_, section)| section.name().starts_with("icp:"))
65        .map(|(_, section)| {
66            let data = section.data(&IdsToIndices::default()).to_vec();
67            let full_name = section.name();
68            match full_name.strip_prefix("public ") {
69                Some(name) => (Kind::Public, name, data),
70                None => match full_name.strip_prefix("private ") {
71                    Some(name) => (Kind::Private, name, data),
72                    None => unreachable!(),
73                },
74            }
75        })
76        .collect();
77
78    // read in from temp file and optimize
79    let mut optimizations = match level {
80        OptLevel::O0 => OptimizationOptions::new_opt_level_0(),
81        OptLevel::O1 => OptimizationOptions::new_opt_level_1(),
82        OptLevel::O2 => OptimizationOptions::new_opt_level_2(),
83        OptLevel::O3 => OptimizationOptions::new_opt_level_3(),
84        OptLevel::O4 => OptimizationOptions::new_opt_level_4(),
85        OptLevel::Os => OptimizationOptions::new_optimize_for_size(),
86        OptLevel::Oz => OptimizationOptions::new_optimize_for_size_aggressively(),
87    };
88    optimizations.debug_info(keep_name_section);
89    optimizations.allow_functions_with_loops(inline_functions_with_loops);
90    if let Some(max_size) = always_inline_max_function_size {
91        optimizations.always_inline_max_size(*max_size);
92    }
93
94    for feature in IC_ENABLED_WASM_FEATURES {
95        optimizations.enable_feature(feature);
96    }
97    optimizations.run(temp_file.path(), temp_file.path())?;
98
99    // read optimized wasm back in from temp file
100    let mut m_opt = parse_wasm_file(temp_file.path().to_path_buf(), keep_name_section)?;
101
102    // re-insert the custom sections
103    metadata_sections
104        .into_iter()
105        .for_each(|(visibility, name, data)| {
106            add_metadata(&mut m_opt, visibility, name, data);
107        });
108
109    *m = m_opt;
110    Ok(())
111}