Merge pull request #55 from paritytech/update-stack

Add stack limit alteration in wasm-build

build/src/main.rs

@@ -14,7 +14,7 @@ use std::path::PathBuf;
 use clap::{App, Arg};
 use parity_wasm::elements;
 
-use wasm_utils::{CREATE_SYMBOL, CALL_SYMBOL, ununderscore_funcs, externalize_mem};
+use wasm_utils::{CREATE_SYMBOL, CALL_SYMBOL, ununderscore_funcs, externalize_mem, shrink_unknown_stack};
 
 #[derive(Debug)]
 pub enum Error {
@@ -98,6 +98,10 @@ fn main() {
 			.help("Save intermediate raw bytecode to path")
 			.takes_value(true)
 			.long("save-raw"))
+		.arg(Arg::with_name("shrink_stack")
+			.help("Shrinks the new stack size for wasm32-unknown-unknown")
+			.takes_value(true)
+			.long("shrink-stack"))
 		.get_matches();
 
     let target_dir = matches.value_of("target").expect("is required; qed");
@@ -129,7 +133,15 @@ fn main() {
 	}
 
 	if let source::SourceTarget::Unknown = source_input.target() {
-		module = externalize_mem(module);
+		// 49152 is 48kb!
+		let stack_size: u32 = matches.value_of("shrink_stack").unwrap_or_else(|| "49152").parse().expect("New stack size is not valid u32");
+		assert!(stack_size <= 1024*1024);
+		let (new_module, new_stack_top) = shrink_unknown_stack(module, 1024 * 1024 - stack_size);
+		module = new_module;
+		let mut stack_top_page = new_stack_top / 65536;
+		if new_stack_top % 65536 > 0 { stack_top_page += 1 };
+
+		module = externalize_mem(module, Some(stack_top_page));
 	}
 
 	if let Some(runtime_type) = matches.value_of("runtime_type") {

src/ext.rs

@@ -1,5 +1,6 @@
 use parity_wasm::{elements, builder};
 use optimizer::{import_section, export_section};
+use byteorder::{LittleEndian, ByteOrder};
 
 type Insertion = (usize, u32, u32, String);
 
@@ -31,13 +32,17 @@ pub fn memory_section<'a>(module: &'a mut elements::Module) -> Option<&'a mut el
 	None
 }
 
-pub fn externalize_mem(mut module: elements::Module) -> elements::Module {
-	let entry = memory_section(&mut module)
+pub fn externalize_mem(mut module: elements::Module, adjust_pages: Option<u32>) -> elements::Module {
+	let mut entry = memory_section(&mut module)
 		.expect("Memory section to exist")
 		.entries_mut()
 		.pop()
 		.expect("Own memory entry to exist in memory section");
 
+	if let Some(adjust_pages) = adjust_pages {
+		entry = elements::MemoryType::new(adjust_pages, None);
+	}
+
 	import_section(&mut module).expect("Import section to exist").entries_mut().push(
 		elements::ImportEntry::new(
 			"env".to_owned(),
@@ -75,11 +80,35 @@ pub fn underscore_funcs(module: elements::Module) -> elements::Module {
 	foreach_public_func_name(module, |n| n.insert(0, '_'))
 }
 
-
 pub fn ununderscore_funcs(module: elements::Module) -> elements::Module {
 	foreach_public_func_name(module, |n| { n.remove(0); })
 }
 
+pub fn shrink_unknown_stack(
+	mut module: elements::Module,
+	// for example, `shrink_amount = (1MB - 64KB)` will limit stack to 64KB
+	shrink_amount: u32,
+) -> (elements::Module, u32) {
+	let mut new_stack_top = 0;
+	for section in module.sections_mut() {
+		match section {
+			&mut elements::Section::Data(ref mut data_section) => {
+				for ref mut data_segment in data_section.entries_mut() {
+					if data_segment.offset().code() == &[elements::Opcode::I32Const(4), elements::Opcode::End] {
+						assert_eq!(data_segment.value().len(), 4);
+						let current_val = LittleEndian::read_u32(data_segment.value());
+						let new_val = current_val - shrink_amount;
+						LittleEndian::write_u32(data_segment.value_mut(), new_val);
+						new_stack_top = new_val;
+					}
+				}
+			},
+			_ => continue
+		}
+	}
+	(module, new_stack_top)
+}
+
 pub fn externalize(
 	module: elements::Module,
 	replaced_funcs: Vec<&str>,

src/lib.rs

@@ -21,7 +21,7 @@ mod runtime_type;
 pub use optimizer::{optimize, Error as OptimizerError};
 pub use gas::inject_gas_counter;
 pub use logger::init_log;
-pub use ext::{externalize, externalize_mem, underscore_funcs, ununderscore_funcs};
+pub use ext::{externalize, externalize_mem, underscore_funcs, ununderscore_funcs, shrink_unknown_stack};
 pub use pack::pack_instance;
 pub use nondeterminism_check::is_deterministic;
 pub use runtime_type::inject_runtime_type;