Merge pull request #59 from paritytech/fix-pack

Fix pack with new dispatch/return methods

src/lib.rs

@@ -6,6 +6,7 @@ extern crate byteorder;
 
 pub static CREATE_SYMBOL: &'static str = "deploy";
 pub static CALL_SYMBOL: &'static str = "call";
+pub static RET_SYMBOL: &'static str = "ret";
 
 pub mod rules;
 

src/pack.rs

@@ -1,6 +1,10 @@
-use parity_wasm::elements::{self, Section, Opcode, DataSegment, InitExpr, Internal};
+use parity_wasm::elements::{
+    self, Section, Opcode, DataSegment, InitExpr, Internal, External,
+    ImportCountType,
+};
 use parity_wasm::builder;
-use super::{CREATE_SYMBOL, CALL_SYMBOL};
+use super::{CREATE_SYMBOL, CALL_SYMBOL, RET_SYMBOL};
+use super::gas::update_call_index;
 
 /// Pack error.
 ///
@@ -15,6 +19,8 @@ pub enum Error {
     InvalidCreateSignature,
     NoCreateSymbol,
     InvalidCreateMember,
+    NoRetImported,
+    NoImportSection,
 }
 
 /// If module has an exported "_create" function we want to pack it into "constructor".
@@ -27,7 +33,7 @@ pub fn pack_instance(raw_module: Vec<u8>, mut ctor_module: elements::Module) ->
 
     // 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 mut create_func_id = {
         let found_entry = ctor_module.export_section().ok_or(Error::NoExportSection)?.entries().iter()
             .find(|entry| CREATE_SYMBOL == entry.field()).ok_or(Error::NoCreateSymbol)?;
 
@@ -47,7 +53,8 @@ pub fn pack_instance(raw_module: Vec<u8>, mut ctor_module: elements::Module) ->
         let &elements::Type::Function(ref func) = ctor_module.type_section().ok_or(Error::NoTypeSection)?
             .types().get(type_id as usize).ok_or(Error::MalformedModule)?;
 
-        if func.params() != &[elements::ValueType::I32] {
+        // Deploy should have no arguments and also should return nothing
+        if !func.params().is_empty() {
             return Err(Error::InvalidCreateSignature);
         }
         if func.return_type().is_some() {
@@ -57,9 +64,72 @@ pub fn pack_instance(raw_module: Vec<u8>, mut ctor_module: elements::Module) ->
         function_internal_index
     };
 
+    let ret_function_id = {
+        let mut id = 0;
+        let mut found = false;
+        for entry in ctor_module.import_section().ok_or(Error::NoImportSection)?.entries().iter() {
+            if let External::Function(_) = *entry.external() {
+                if entry.field() == RET_SYMBOL { found = true; break; }
+                else { id += 1; }
+            }
+        }
+        if !found {
+            let mut mbuilder = builder::from_module(ctor_module);
+            let import_sig = mbuilder.push_signature(
+                builder::signature()
+                    .param().i32().param().i32()
+                    .build_sig()
+                );
+
+            mbuilder.push_import(
+                builder::import()
+                    .module("env")
+                    .field("ret")
+                    .external().func(import_sig)
+                    .build()
+                );
+
+            ctor_module = mbuilder.build();
+
+            let ret_func = ctor_module.import_count(ImportCountType::Function) as u32 - 1;
+
+            for section in ctor_module.sections_mut() {
+                match *section {
+                    elements::Section::Code(ref mut code_section) => {
+                        for ref mut func_body in code_section.bodies_mut() {
+                            update_call_index(func_body.code_mut(), ret_func);
+                        }
+                    },
+                    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 >= ret_func { *func_index += 1}
+                                },
+                                _ => {}
+                            }
+                        }
+                    },
+                    elements::Section::Element(ref mut elements_section) => {
+                        for ref mut segment in elements_section.entries_mut() {
+                            // update all indirect call addresses initial values
+                            for func_index in segment.members_mut() {
+                                if *func_index >= ret_func { *func_index += 1}
+                            }
+                        }
+                    },
+                    _ => { }
+                }
+            }
+
+            create_func_id += 1;
+            ret_func
+         }
+        else { id }
+    };
+
     // 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;
+    let last_function_index = ctor_module.functions_space();
 
     // Code data address is an address where we put the contract's code (raw_module)
     let mut code_data_address = 0i32;
@@ -93,17 +163,13 @@ pub fn pack_instance(raw_module: Vec<u8>, mut ctor_module: elements::Module) ->
 
     let mut new_module = builder::from_module(ctor_module)
         .function()
-        .signature().param().i32().build()
+        .signature().build()
         .body().with_opcodes(elements::Opcodes::new(
             vec![
-                Opcode::GetLocal(0),
                 Opcode::Call((create_func_id + ctor_import_functions) 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::Call(ret_function_id as u32),
                 Opcode::End,
             ])).build()
             .build()