diff --git a/ext/Cargo.toml b/ext/Cargo.toml index cfaa1c5..f69ebed 100644 --- a/ext/Cargo.toml +++ b/ext/Cargo.toml @@ -4,4 +4,5 @@ version = "0.1.0" authors = ["NikVolf "] [dependencies] -parity-wasm = { git="https://github.com/nikvolf/parity-wasm" } \ No newline at end of file +parity-wasm = { git="https://github.com/nikvolf/parity-wasm" } +wasm_utils = { path = "../" } \ No newline at end of file diff --git a/ext/src/main.rs b/ext/src/main.rs index 541b118..e4fb0ed 100644 --- a/ext/src/main.rs +++ b/ext/src/main.rs @@ -1,28 +1,7 @@ extern crate parity_wasm; +extern crate wasm_utils; use std::env; -use parity_wasm::{builder, elements}; - -type Insertion = (usize, u32, u32, String); - -pub fn update_call_index(opcodes: &mut elements::Opcodes, original_imports: usize, inserts: &[Insertion]) { - use parity_wasm::elements::Opcode::*; - for opcode in opcodes.elements_mut().iter_mut() { - match opcode { - &mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => { - update_call_index(block, original_imports, inserts) - }, - &mut Call(ref mut call_index) => { - if let Some(pos) = inserts.iter().position(|x| x.1 == *call_index) { - *call_index = (original_imports + pos) as u32; - } else if *call_index as usize > original_imports { - *call_index += inserts.len() as u32; - } - }, - _ => { } - } - } -} fn main() { @@ -32,84 +11,10 @@ fn main() { return; } - // Loading module - let module = parity_wasm::deserialize_file(&args[1]).unwrap(); + let module = wasm_utils::externalize( + parity_wasm::deserialize_file(&args[1]).expect("Module to deserialize ok"), + vec!["_free", "_malloc"], + ); - let replaced_funcs = vec!["_free", "_malloc"]; - - // Save import functions number for later - let import_funcs_total = module - .import_section().expect("Import section to exist") - .entries() - .iter() - .filter(|e| if let &elements::External::Function(_) = e.external() { true } else { false }) - .count(); - - // First, we find functions indices that are to be rewired to externals - // Triple is (function_index (callable), type_index, function_name) - let mut replaces: Vec = replaced_funcs - .into_iter() - .filter_map(|f| { - let export = module - .export_section().expect("Export section to exist") - .entries().iter().enumerate() - .find(|&(_, entry)| entry.field() == f) - .expect("All functions of interest to exist"); - - if let &elements::Internal::Function(func_idx) = export.1.internal() { - let type_ref = module - .functions_section().expect("Functions section to exist") - .entries()[func_idx as usize - import_funcs_total] - .type_ref(); - - Some((export.0, func_idx, type_ref, export.1.field().to_owned())) - } else { - None - } - }) - .collect(); - - replaces.sort_by_key(|e| e.0); - - // Second, we duplicate them as import definitions - let mut mbuilder = builder::from_module(module); - for &(_, _, type_ref, ref field) in replaces.iter() { - mbuilder.push_import( - builder::import() - .module("env") - .field(field) - .external().func(type_ref) - .build() - ); - } - - // Back to mutable access - let mut module = mbuilder.build(); - - // Third, rewire all calls to imported functions and update all other calls indices - for section in module.sections_mut() { - match section { - &mut elements::Section::Code(ref mut code_section) => { - for ref mut func_body in code_section.bodies_mut() { - update_call_index(func_body.code_mut(), import_funcs_total, &replaces); - } - }, - &mut elements::Section::Export(ref mut export_section) => { - for ref mut export in export_section.entries_mut() { - match export.internal_mut() { - &mut elements::Internal::Function(ref mut func_index) => { - if *func_index >= import_funcs_total as u32 { *func_index += replaces.len() as u32; } - }, - _ => {} - } - } - }, - _ => { } - } - } - - // Forth step could be to eliminate now dead code in actual functions - // but it is managed by `wasm-opt` - - parity_wasm::serialize_to_file(&args[2], module).unwrap(); + parity_wasm::serialize_to_file(&args[2], module).expect("Module to serialize ok"); } diff --git a/src/ext.rs b/src/ext.rs new file mode 100644 index 0000000..1cfd528 --- /dev/null +++ b/src/ext.rs @@ -0,0 +1,101 @@ +use parity_wasm::{elements, builder}; + +type Insertion = (usize, u32, u32, String); + +pub fn update_call_index(opcodes: &mut elements::Opcodes, original_imports: usize, inserts: &[Insertion]) { + use parity_wasm::elements::Opcode::*; + for opcode in opcodes.elements_mut().iter_mut() { + match opcode { + &mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => { + update_call_index(block, original_imports, inserts) + }, + &mut Call(ref mut call_index) => { + if let Some(pos) = inserts.iter().position(|x| x.1 == *call_index) { + *call_index = (original_imports + pos) as u32; + } else if *call_index as usize > original_imports { + *call_index += inserts.len() as u32; + } + }, + _ => { } + } + } +} + +pub fn externalize( + module: elements::Module, + replaced_funcs: Vec<&str>, +) -> elements::Module { + // Save import functions number for later + let import_funcs_total = module + .import_section().expect("Import section to exist") + .entries() + .iter() + .filter(|e| if let &elements::External::Function(_) = e.external() { true } else { false }) + .count(); + + // First, we find functions indices that are to be rewired to externals + // Triple is (function_index (callable), type_index, function_name) + let mut replaces: Vec = replaced_funcs + .into_iter() + .filter_map(|f| { + let export = module + .export_section().expect("Export section to exist") + .entries().iter().enumerate() + .find(|&(_, entry)| entry.field() == f) + .expect("All functions of interest to exist"); + + if let &elements::Internal::Function(func_idx) = export.1.internal() { + let type_ref = module + .functions_section().expect("Functions section to exist") + .entries()[func_idx as usize - import_funcs_total] + .type_ref(); + + Some((export.0, func_idx, type_ref, export.1.field().to_owned())) + } else { + None + } + }) + .collect(); + + replaces.sort_by_key(|e| e.0); + + // Second, we duplicate them as import definitions + let mut mbuilder = builder::from_module(module); + for &(_, _, type_ref, ref field) in replaces.iter() { + mbuilder.push_import( + builder::import() + .module("env") + .field(field) + .external().func(type_ref) + .build() + ); + } + + // Back to mutable access + let mut module = mbuilder.build(); + + // Third, rewire all calls to imported functions and update all other calls indices + for section in module.sections_mut() { + match section { + &mut elements::Section::Code(ref mut code_section) => { + for ref mut func_body in code_section.bodies_mut() { + update_call_index(func_body.code_mut(), import_funcs_total, &replaces); + } + }, + &mut elements::Section::Export(ref mut export_section) => { + for ref mut export in export_section.entries_mut() { + match export.internal_mut() { + &mut elements::Internal::Function(ref mut func_index) => { + if *func_index >= import_funcs_total as u32 { *func_index += replaces.len() as u32; } + }, + _ => {} + } + } + }, + _ => { } + } + } + + module + +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index da5e3dc..dbc3c5c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,7 +7,9 @@ mod optimizer; mod gas; mod symbols; mod logger; +mod ext; pub use optimizer::optimize; pub use gas::inject_gas_counter; -pub use logger::init_log; \ No newline at end of file +pub use logger::init_log; +pub use ext::externalize; \ No newline at end of file