separate the context intialization from the function execution

Geal · Geal · commit 6854deca1b84 · 2018-11-20T15:35:42.000+01:00
the start function is called from the initialization phase
diff --git a/lib/environ/src/lib.rs b/lib/environ/src/lib.rs
@@ -43,4 +43,4 @@ mod module;
 
 pub use compilation::{compile_module, Compilation, Relocation, RelocationTarget, Relocations};
 pub use environ::{ModuleEnvironment, ModuleTranslation};
-pub use module::{DataInitializer, Module, TableElements};
+pub use module::{DataInitializer, Export, Module, TableElements};
diff --git a/lib/execute/src/execute.rs b/lib/execute/src/execute.rs
@@ -9,7 +9,7 @@ use region::Protection;
 use std::mem::transmute;
 use std::ptr::{self, write_unaligned};
 use wasmtime_environ::{
-    compile_module, Compilation, Module, ModuleTranslation, Relocation, RelocationTarget,
+    compile_module, Compilation, Export, Module, ModuleTranslation, Relocation, RelocationTarget,
 };
 
 /// Executes a module that has been translated with the `wasmtime-environ` environment
@@ -125,15 +125,12 @@ fn make_vmctx(instance: &mut Instance, mem_base_addrs: &mut [*mut u8]) -> Vec<*m
     vmctx
 }
 
-/// Jumps to the code region of memory and execute the start function of the module.
-pub fn execute(
+/// prepares the execution context
+pub fn finish_instantiation(
     module: &Module,
     compilation: &Compilation,
     instance: &mut Instance,
-) -> Result<(), String> {
-    let start_index = module
-        .start_func
-        .ok_or_else(|| String::from("No start function defined, aborting execution"))?;
+) -> Result<Vec<*mut u8>, String> {
     // TODO: Put all the function bodies into a page-aligned memory region, and
     // then make them ReadExecute rather than ReadWriteExecute.
     for code_buf in compilation.functions.values() {
@@ -154,26 +151,62 @@ pub fn execute(
         }
     }
 
-    let code_buf =
-        &compilation.functions[module
-                                   .defined_func_index(start_index)
-                                   .expect("imported start functions not supported yet")];
-
     // Collect all memory base addresses and Vec.
     let mut mem_base_addrs = instance
         .memories
         .iter_mut()
         .map(LinearMemory::base_addr)
         .collect::<Vec<_>>();
+
     let vmctx = make_vmctx(instance, &mut mem_base_addrs);
 
+    if let Ok(start_index) = module
+        .start_func
+        .ok_or_else(|| String::from("No start function defined, aborting execution"))
+    {
+        let code_buf =
+            &compilation.functions[module
+                                       .defined_func_index(start_index)
+                                       .expect("imported start functions not supported yet")];
+
+        // Rather than writing inline assembly to jump to the code region, we use the fact that
+        // the Rust ABI for calling a function with no arguments and no return matches the one of
+        // the generated code. Thanks to this, we can transmute the code region into a first-class
+        // Rust function and call it.
+        unsafe {
+            let start_func = transmute::<_, fn(*const *mut u8)>(code_buf.as_ptr());
+            start_func(vmctx.as_ptr());
+        }
+    }
+
+    Ok(vmctx)
+}
+
+/// Jumps to the code region of memory and execute the exported function
+pub fn execute(
+    module: &Module,
+    compilation: &Compilation,
+    vmctx: &mut Vec<*mut u8>,
+    function: &str,
+) -> Result<(), String> {
+    let fn_index = match module.exports.get(function) {
+        Some(Export::Function(index)) => *index,
+        Some(_) => return Err(format!("exported item \"{}\" is not a function", function)),
+        None => return Err(format!("no export named \"{}\"", function)),
+    };
+
+    let code_buf =
+        &compilation.functions[module
+                                   .defined_func_index(fn_index)
+                                   .expect("imported start functions not supported yet")];
+
     // Rather than writing inline assembly to jump to the code region, we use the fact that
     // the Rust ABI for calling a function with no arguments and no return matches the one of
     // the generated code. Thanks to this, we can transmute the code region into a first-class
     // Rust function and call it.
     unsafe {
-        let start_func = transmute::<_, fn(*const *mut u8)>(code_buf.as_ptr());
-        start_func(vmctx.as_ptr());
+        let func = transmute::<_, fn(*const *mut u8)>(code_buf.as_ptr());
+        func(vmctx.as_ptr());
     }
     Ok(())
 }
diff --git a/lib/execute/src/lib.rs b/lib/execute/src/lib.rs
@@ -40,5 +40,5 @@ mod execute;
 mod instance;
 mod memory;
 
-pub use execute::{compile_and_link_module, execute};
+pub use execute::{compile_and_link_module, execute, finish_instantiation};
 pub use instance::Instance;
diff --git a/src/main.rs b/src/main.rs
@@ -56,7 +56,7 @@ use std::path::PathBuf;
 use std::process::{exit, Command};
 use tempdir::TempDir;
 use wasmtime_environ::{Module, ModuleEnvironment};
-use wasmtime_execute::{compile_and_link_module, execute, Instance};
+use wasmtime_execute::{compile_and_link_module, execute, finish_instantiation, Instance};
 
 const USAGE: &str = "
 Wasm to Cranelift IL translation utility.
@@ -65,11 +65,13 @@ The translation is dependent on the environment chosen.
 
 Usage:
     wasmtime [-mop] <file>...
+    wasmtime [-mop] <file>... --function=<fn>
     wasmtime --help | --version
 
 Options:
     -o, --optimize      runs optimization passes on the translated functions
     -m, --memory        interactive memory inspector after execution
+    --function=<fn>     name of function to run
     -h, --help          print this help message
     --version           print the Cranelift version
 ";
@@ -79,6 +81,7 @@ struct Args {
     arg_file: Vec<String>,
     flag_memory: bool,
     flag_optimize: bool,
+    flag_function: Option<String>,
 }
 
 fn read_to_end(path: PathBuf) -> Result<Vec<u8>, io::Error> {
@@ -157,7 +160,14 @@ fn handle_module(args: &Args, path: PathBuf, isa: &TargetIsa) -> Result<(), Stri
                 &compilation,
                 &translation.lazy.data_initializers,
             );
-            execute(&translation.module, &compilation, &mut instance)?;
+
+            let mut context =
+                finish_instantiation(&translation.module, &compilation, &mut instance)?;
+
+            if let Some(ref f) = args.flag_function {
+                execute(&translation.module, &compilation, &mut context, &f)?;
+            }
+
             instance
         }
         Err(s) => {
