Skip to main content

wamex_cli/
lib.rs

1use std::path::PathBuf;
2
3use clap::{Args, Parser, Subcommand, ValueEnum};
4
5// todo: Refactor analysis and emit modules.
6pub mod analysis;
7pub mod emit;
8mod helpers;
9#[macro_use]
10mod index;
11mod diff;
12mod incremental;
13pub mod read;
14
15pub use analysis::split_point::{ModuleIdentifier, SplitModuleIdentifier, SplitProgramInfo};
16pub use anyhow::Result;
17pub use incremental::{
18    IncrementalSplitResult, IncrementalSplitState, ModuleDeps, ModuleUpdate, SplitResult,
19};
20pub use read::InputModule;
21pub use wamex_types::{BumpVersion, ModuleId};
22
23use crate::emit::CommonEmitInfo;
24
25#[derive(Debug, Parser)]
26#[command(name = "wasm-split")]
27pub struct Cli {
28    #[command(subcommand)]
29    pub command: Command,
30}
31#[derive(Debug, Args)]
32#[command(name = "wasm-split")]
33pub struct Split {
34    /// Input .wasm file.
35    pub input: PathBuf,
36
37    /// Output directory.
38    pub output: PathBuf,
39
40    /// Print verbose split information.
41    /// Also if metadata is enabled, it will print it in pretty JSON format.
42    #[arg(short, long)]
43    pub verbose: bool,
44
45    /// Parse instructions when creating pic modules (slower, but more robust to wasm spec changes).
46    #[arg(short, long)]
47    pub precise_modification: bool,
48
49    /// Skip writing files (for benchmarking).
50    #[arg(long)]
51    pub dry_run: bool,
52
53    /// Specify the split point extraction strategy.
54    #[arg(value_enum, default_value_t = SplitPointExtractor::Wamex)]
55    pub split_point_extractor: SplitPointExtractor,
56}
57
58/// This is temporary solution to support old __wamex__ split points
59#[derive(Debug, ValueEnum, Clone, Copy)]
60pub enum SplitPointExtractor {
61    /// Use regexp and _wasm_split_ prefix to identify split points.
62    Legacy,
63    /// Use _wamex_ prefix and .start_with instead of regexp.
64    Wamex,
65}
66
67#[derive(Debug, Args)]
68pub struct Diff {
69    pub left: PathBuf,
70    pub right: PathBuf,
71    #[arg(short, long)]
72    pub structural: bool,
73}
74
75#[derive(Debug, Args)]
76pub struct Debug {
77    pub input: PathBuf,
78}
79
80#[derive(Debug, Args)]
81pub struct Roundtrip {
82    pub input: PathBuf,
83    pub output: PathBuf,
84}
85
86#[derive(Debug, Subcommand)]
87pub enum Command {
88    /// Split wasm module into multiple parts.
89    Split(Split),
90
91    /// Incremental split wasm module into multiple parts.
92    /// loop and repeat split on keypresses.
93    IncrementalSplit(Split),
94
95    /// Compare two wasm modules.
96    Diff(Diff),
97
98    /// Roundtrip wasm module.
99    Roundtrip(Roundtrip),
100
101    Debug(Debug),
102}
103
104//The flow of the program is simple:
105// 1. Parse the input wasm file.
106// 2. Analyze the wasm module to gather information about its structural dependencies. And indetify split points.
107// 3. Emit processed modules to the output directory.
108// 3.1 Modify module:
109//    - Relocate functions.
110//    - Patch data lookups. (e.g. const.get -> global.get)
111// Also there should be a routine that can compare and reload changed chunks.
112
113pub fn main(args: Cli) -> Result<()> {
114    match args.command {
115        Command::Split(args) => split(args)?,
116        Command::IncrementalSplit(args) => incremental_split(args)?,
117        Command::Diff(args) => diff(args)?,
118        Command::Roundtrip(args) => roundtrip(args)?,
119        Command::Debug(args) => debug(args)?,
120    };
121    Ok(())
122}
123pub fn roundtrip(args: Roundtrip) -> Result<()> {
124    let input_wasm = std::fs::read(&args.input)?;
125    let module = InputModule::parse(&input_wasm)?;
126    let info = analysis::ModuleInfo::from_raw_module(module)?;
127    let dep_graph = analysis::dep_graph::get_dependencies(&info)?;
128
129    let split_program_info =
130        SplitProgramInfo::compute_split_modules(&info, &dep_graph, &[], &Default::default())?;
131
132    assert!(
133        split_program_info.output_modules.len() == 1,
134        "Roundtrip should produce single module",
135    );
136    crate::emit::emit_modules(
137        &info,
138        false,
139        &split_program_info,
140        &Default::default(),
141        false,
142        None,
143        Default::default(),
144        |_: &SplitModuleIdentifier, data: &[u8]| -> Result<()> {
145            std::fs::write(&args.output, data)?;
146            Ok(())
147        },
148    )?;
149
150    Ok(())
151}
152pub fn split(args: Split) -> Result<()> {
153    let input_wasm = std::fs::read(&args.input)?;
154    let _ = split_inner(
155        &input_wasm,
156        args.verbose,
157        args.precise_modification,
158        args.split_point_extractor,
159        |identifier: ModuleId, data: &[u8]| -> Result<()> {
160            let output_filename = format!("{}.wasm", identifier.module_name());
161            if !args.dry_run {
162                std::fs::create_dir_all(&args.output)?;
163                std::fs::write(args.output.join(output_filename), data)?;
164            } else {
165                log::info!("Skipping writing module {output_filename} (dry run)");
166            }
167            Ok(())
168        },
169    )?;
170    Ok(())
171}
172
173#[doc(hidden)]
174// Full split routine, but without file I/O reading
175// Returns a map of dependency for modules in format (ModuleId -> Vec<ModuleId>)
176pub fn split_inner(
177    input_wasm: &[u8],
178    verbose: bool,
179    precise_modification: bool,
180    split_point_extractor: SplitPointExtractor,
181    mut emit_module_fn: impl FnMut(ModuleId, &[u8]) -> Result<()>,
182) -> Result<ModuleDeps> {
183    let mut state = IncrementalSplitState::new();
184    let split_result = state.split_incremental(
185        input_wasm,
186        verbose,
187        precise_modification,
188        split_point_extractor,
189        |identifier: ModuleId, data: &[u8]| -> Result<()> { emit_module_fn(identifier, data) },
190    )?;
191    Ok(split_result.deps)
192}
193
194fn incremental_split(args: Split) -> Result<()> {
195    println!("Starting incremental split loop...");
196    let mut state = IncrementalSplitState::new();
197
198    loop {
199        let input_wasm = std::fs::read(&args.input)?;
200        let split_result = state.split_incremental(
201            &input_wasm,
202            args.verbose,
203            args.precise_modification,
204            args.split_point_extractor,
205            |identifier: ModuleId, data: &[u8]| -> Result<()> {
206                let output_filename = format!("{}.wasm", identifier.module_full_name());
207                if !args.dry_run {
208                    std::fs::create_dir_all(&args.output)?;
209                    std::fs::write(args.output.join(output_filename), data)?;
210                } else {
211                    log::info!("Skipping writing module {output_filename} (dry run)");
212                }
213                Ok(())
214            },
215        )?;
216
217        match split_result.incremental_result {
218            IncrementalSplitResult::Unchanged => {
219                println!("No changes detected, all modules are up to date.");
220            }
221            IncrementalSplitResult::UpdatedModules(modules) => {
222                println!(
223                    "Updated modules: {}",
224                    modules
225                        .iter()
226                        .map(|m| m.module_id.module_full_name())
227                        .collect::<Vec<_>>()
228                        .join(", ")
229                );
230            }
231            IncrementalSplitResult::FullResplit => {
232                println!("Full resplit performed, all modules were regenerated.");
233            }
234        }
235        println!("Current module dependencies:");
236        for (module, deps) in split_result.deps.iter() {
237            println!(
238                "  {} -> [{}]",
239                module.module_full_name(),
240                deps.iter()
241                    .map(|d| d.module_full_name())
242                    .collect::<Vec<_>>()
243                    .join(", ")
244            );
245        }
246
247        println!("Press Enter to re-split, or Ctrl+C to exit.");
248        let mut input = String::new();
249        std::io::stdin().read_line(&mut input)?;
250    }
251}
252
253pub fn diff(args: Diff) -> Result<()> {
254    let left = std::fs::read(&args.left)?;
255    let right = std::fs::read(&args.right)?;
256    let left_module = InputModule::parse(&left)?;
257    let right_module = InputModule::parse(&right)?;
258    let left_module_info = analysis::ModuleInfo::from_raw_module(left_module)?;
259    let right_module_info = analysis::ModuleInfo::from_raw_module(right_module)?;
260
261    let diff = diff::Compare::new(&left_module_info, &right_module_info, args.structural);
262    diff.print_diff()?;
263
264    Ok(())
265}
266
267pub fn debug(args: Debug) -> Result<()> {
268    let input = std::fs::read(&args.input)?;
269    let module = InputModule::parse(&input)?;
270    let info = analysis::ModuleInfo::from_raw_module(module)?;
271
272    let program_info = analysis::split_point::SplitProgramInfo::default();
273    // verbose flag will print debug info as side effect.
274    // TODO: make it more functional.
275    let _ci = CommonEmitInfo::new(&info, true, &program_info)?;
276
277    info.symbols.print_debug();
278    Ok(())
279}
280
281// 1. check imports - exports
282// 1.1. no wamex_split.rs imports should be present in main module
283// 1.2. all exports should be used in some module
284// 1.3. "lazy" imports (indirect fns) should be reserved for specific modules only. This place should be inited in elem section of modules.
285// 2. Check that data segments are same from original module (no data loss, and no extra fields).
286// 3. same for functions - no loss, no extra functions (only trampolines).
287
288// pub fn validate() {
289
290// }