Merge pull request #42 from paritytech/fix-constructor

Refactor constructor packing

build/Cargo.lock

@@ -1,14 +1,3 @@
-[root]
-name = "wasm-build"
-version = "0.1.0"
-dependencies = [
- "clap 2.27.1 (registry+https://github.com/rust-lang/crates.io-index)",
- "glob 0.2.11 (registry+https://github.com/rust-lang/crates.io-index)",
- "parity-wasm 0.15.3 (registry+https://github.com/rust-lang/crates.io-index)",
- "tempdir 0.3.5 (registry+https://github.com/rust-lang/crates.io-index)",
- "wasm-utils 0.1.0",
-]
-
 [[package]]
 name = "aho-corasick"
 version = "0.6.3"
@@ -134,7 +123,7 @@ dependencies = [
 
 [[package]]
 name = "parity-wasm"
-version = "0.15.3"
+version = "0.15.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 dependencies = [
  "byteorder 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -280,6 +269,17 @@ name = "void"
 version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 
+[[package]]
+name = "wasm-build"
+version = "0.1.0"
+dependencies = [
+ "clap 2.27.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "glob 0.2.11 (registry+https://github.com/rust-lang/crates.io-index)",
+ "parity-wasm 0.15.4 (registry+https://github.com/rust-lang/crates.io-index)",
+ "tempdir 0.3.5 (registry+https://github.com/rust-lang/crates.io-index)",
+ "wasm-utils 0.1.0",
+]
+
 [[package]]
 name = "wasm-utils"
 version = "0.1.0"
@@ -290,7 +290,7 @@ dependencies = [
  "glob 0.2.11 (registry+https://github.com/rust-lang/crates.io-index)",
  "lazy_static 0.2.10 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)",
- "parity-wasm 0.15.3 (registry+https://github.com/rust-lang/crates.io-index)",
+ "parity-wasm 0.15.4 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -321,7 +321,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 "checksum log 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)" = "880f77541efa6e5cc74e76910c9884d9859683118839d6a1dc3b11e63512565b"
 "checksum memchr 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)" = "148fab2e51b4f1cfc66da2a7c32981d1d3c083a803978268bb11fe4b86925e7a"
 "checksum owning_ref 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)" = "cdf84f41639e037b484f93433aa3897863b561ed65c6e59c7073d7c561710f37"
-"checksum parity-wasm 0.15.3 (registry+https://github.com/rust-lang/crates.io-index)" = "8431a184ad88cfbcd71a792aaca319cc7203a94300c26b8dce2d0df0681ea87d"
+"checksum parity-wasm 0.15.4 (registry+https://github.com/rust-lang/crates.io-index)" = "235801e9531998c4bb307f4ea6833c9f40a4cf132895219ac8c2cd25a9b310f7"
 "checksum parking_lot 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)" = "149d8f5b97f3c1133e3cfcd8886449959e856b557ff281e292b733d7c69e005e"
 "checksum parking_lot_core 0.2.6 (registry+https://github.com/rust-lang/crates.io-index)" = "4f610cb9664da38e417ea3225f23051f589851999535290e077939838ab7a595"
 "checksum rand 0.3.18 (registry+https://github.com/rust-lang/crates.io-index)" = "6475140dfd8655aeb72e1fd4b7a1cc1c202be65d71669476e392fe62532b9edd"

build/src/main.rs

@@ -115,7 +115,7 @@ fn main() {
 		if !matches.is_present("skip_optimization") {
 			wasm_utils::optimize(&mut ctor_module, vec![CREATE_SYMBOL]).expect("Optimizer to finish without errors");
 		}
-		wasm_utils::pack_instance(raw_module, &mut ctor_module);
+		let ctor_module = wasm_utils::pack_instance(raw_module, ctor_module).expect("Packing failed");
 		parity_wasm::serialize_to_file(&path, ctor_module).expect("Failed to serialize to file");
 	} else {
 		let mut file = fs::File::create(&path).expect("Failed to create file");

src/pack.rs

@@ -1,37 +1,64 @@
-use parity_wasm::{elements};
+use parity_wasm::elements::{self, Section, Opcode, DataSegment, InitExpr, Internal};
-use self::elements::{ External, Section, Opcode, DataSegment, InitExpr, Internal };
+use parity_wasm::builder;
-
 use super::{CREATE_SYMBOL, CALL_SYMBOL};
 
+/// Pack error.
+///
+/// Pack has number of assumptions of passed module structure.
+/// When they are violated, pack_instance returns one of these.
+#[derive(Debug)]
+pub enum Error {
+    MalformedModule,
+    NoTypeSection,
+    NoExportSection,
+    NoCodeSection,
+    InvalidCreateSignature,
+    NoCreateSymbol,
+    InvalidCreateMember,
+}
+
 /// If module has an exported "_create" function we want to pack it into "constructor".
 /// `raw_module` is the actual contract code
 /// `ctor_module` is the constructor which should return `raw_module`
-pub fn pack_instance(raw_module: Vec<u8>, ctor_module: &mut elements::Module) {
+pub fn pack_instance(raw_module: Vec<u8>, mut ctor_module: elements::Module) -> Result<elements::Module, Error> {
+
+    // Total number of constructor module import functions
+    let ctor_import_functions = ctor_module.import_section().map(|x| x.functions()).unwrap_or(0);
 
     // We need to find an internal ID of function witch is exported as "_create"
     // in order to find it in the Code section of the module
     let create_func_id = {
-        let found_entry = ctor_module.export_section().expect("No export section found").entries().iter()
+        let found_entry = ctor_module.export_section().ok_or(Error::NoExportSection)?.entries().iter()
-            .find(|entry| CREATE_SYMBOL == entry.field()).expect(&format!("No export with name {} found", CREATE_SYMBOL));
+            .find(|entry| CREATE_SYMBOL == entry.field()).ok_or(Error::NoCreateSymbol)?;
 
         let function_index: usize = match found_entry.internal() {
             &Internal::Function(index) => index as usize,
-            _ => panic!("export is not a function"),
+            _ => { return Err(Error::InvalidCreateMember) },
         };
 
-        let import_section_len: usize = match ctor_module.import_section() {
+        // Constructor should be of signature `func(i32)` (void), fail otherwise
-            Some(import) =>
+        let type_id = ctor_module.function_section().ok_or(Error::NoCodeSection)?
-                import.entries().iter().filter(|entry| match entry.external() {
+            .entries().get(function_index).ok_or(Error::MalformedModule)?
-                    &External::Function(_) => true,
+            .type_ref();
-                    _ => false,
+
-                    }).count(),
+        let &elements::Type::Function(ref func) = ctor_module.type_section().ok_or(Error::NoTypeSection)?
-            None => 0,
+            .types().get(type_id as usize).ok_or(Error::MalformedModule)?;
-        };
+
+        if func.params() != &[elements::ValueType::I32] {
+            return Err(Error::InvalidCreateSignature);
+        }
+        if func.return_type().is_some() {
+            return Err(Error::InvalidCreateSignature);
+        }
 
         // Calculates a function index within module's function section
-        function_index - import_section_len
+        function_index - ctor_import_functions
     };
 
+    // If new function is put in ctor module, it will have this callable index
+    let last_function_index = ctor_module.function_section().map(|x| x.entries().len()).unwrap_or(0)
+        + ctor_import_functions;
+
     // Code data address is an address where we put the contract's code (raw_module)
     let mut code_data_address = 0i32;
 
@@ -62,34 +89,40 @@ pub fn pack_instance(raw_module: Vec<u8>, ctor_module: &mut elements::Module) {
         }
     }
 
-    for section in ctor_module.sections_mut() {
+    let mut new_module = builder::from_module(ctor_module)
+        .function()
+        .signature().param().i32().build()
+        .body().with_opcodes(elements::Opcodes::new(
+            vec![
+                Opcode::GetLocal(0),
+                Opcode::Call(create_func_id as u32),
+                Opcode::GetLocal(0),
+                Opcode::I32Const(code_data_address),
+                Opcode::I32Store(0, 8),
+                Opcode::GetLocal(0),
+                Opcode::I32Const(raw_module.len() as i32),
+                Opcode::I32Store(0, 12),
+                Opcode::End,
+            ])).build()
+            .build()
+        .build();
+
+    for section in new_module.sections_mut() {
         match section {
             &mut Section::Export(ref mut export_section) => {
                 for entry in export_section.entries_mut().iter_mut() {
                     if CREATE_SYMBOL == entry.field() {
                         // change _create export name into default _call
                         *entry.field_mut() = CALL_SYMBOL.to_owned();
+                        *entry.internal_mut() = elements::Internal::Function(last_function_index as u32);
                     }
                 }
-            }
-
-            &mut Section::Code(ref mut code_section) => {
-                let code = code_section.bodies_mut()[create_func_id].code_mut().elements_mut();
-                let init_result_code = &[
-                    Opcode::GetLocal(0),
-                    Opcode::I32Const(code_data_address),
-                    Opcode::I32Store(0, 8),
-                    Opcode::GetLocal(0),
-                    Opcode::I32Const(raw_module.len() as i32),
-                    Opcode::I32Store(0, 12)];
-                let mut updated_func_code = Vec::with_capacity(init_result_code.len() + code.len());
-                updated_func_code.extend_from_slice(init_result_code);
-                updated_func_code.extend_from_slice(&code);
-                *code = updated_func_code;
             },
-            _ => {;},
+            _ => { },
         }
     };
+
+    Ok(new_module)
 }
 
 #[cfg(test)]
@@ -153,7 +186,7 @@ mod test {
         optimize(&mut ctor_module, vec![CREATE_SYMBOL]).expect("Optimizer to finish without errors");
 
         let raw_module = parity_wasm::serialize(module).unwrap();
-        pack_instance(raw_module.clone(), &mut ctor_module);
+        let ctor_module = pack_instance(raw_module.clone(), ctor_module).expect("Packing failed");
 
         let program = parity_wasm::DefaultProgramInstance::new().expect("Program instance failed to load");
         let env_instance = program.module("env").expect("Wasm program to contain env module");