Reimplement function search in rust to remove the need for CString re…

…allocation
wasm3 · Jan 26, 2020 · 6188b2a · 6188b2a
1 parent 7641acb
commit 6188b2a
Show file tree

Hide file tree

Showing 7 changed files with 60 additions and 22 deletions.
diff --git a/examples/call_wasm.rs b/examples/call_wasm.rs
@@ -6,7 +6,7 @@ fn main() {
     let rt = env.create_runtime(1024 * 60);
     let module = Module::parse(&env, &include_bytes!("wasm/wasm_add/wasm_add.wasm")[..]).unwrap();
 
-    assert!(rt.load_module(module).is_ok());
-    let func = rt.find_function::<(i64, i64), i64>("add").unwrap();
+    let module = rt.load_module(module).map_err(|(_, e)| e).unwrap();
+    let func = module.find_function::<(i64, i64), i64>("add").unwrap();
     println!("Wasm says that 3 + 6 is {}", func.call(3, 6).unwrap())
 }
diff --git a/examples/wasm_link.rs b/examples/wasm_link.rs
@@ -14,17 +14,17 @@ fn main() {
     module
         .link_function::<(), u64>("time", "millis", Some(millis))
         .unwrap();
-    let func = rt.find_function::<(), u64>("seconds").unwrap();
+    let func = module.find_function::<(), u64>("seconds").unwrap();
     println!("{}ms in seconds is {:?}s.", MILLIS, func.call());
     assert_eq!(func.call(), Ok(MILLIS / 1000));
 }
 
 const MILLIS: u64 = 500_000;
 
 unsafe extern "C" fn millis(
-    rt: ffi::IM3Runtime,
+    _rt: ffi::IM3Runtime,
     sp: *mut u64,
-    mem: *mut std::ffi::c_void,
+    _mem: *mut std::ffi::c_void,
 ) -> *const std::ffi::c_void {
     *sp = MILLIS;
     ffi::m3Err_none as _

diff --git a/src/environment.rs b/src/environment.rs
@@ -2,6 +2,7 @@ use crate::error::Result;
 use crate::module::ParsedModule;
 use crate::runtime::Runtime;
 
+#[derive(Debug)]
 pub struct Environment(ffi::IM3Environment);
 
 impl Environment {

diff --git a/src/function.rs b/src/function.rs
@@ -4,6 +4,7 @@ use crate::error::{Error, Result};
 use crate::runtime::Runtime;
 use crate::{WasmArgs, WasmType};
 
+#[derive(Debug)]
 pub struct Function<'env, 'rt, ARGS, RET> {
     raw: ffi::IM3Function,
     rt: &'rt Runtime<'env>,
@@ -28,16 +29,27 @@ where
     #[inline]
     pub(crate) fn from_raw(rt: &'rt Runtime<'env>, raw: ffi::IM3Function) -> Result<Self> {
         if Self::validate_sig(raw) {
-            Ok(Function {
+            let this = Function {
                 raw,
                 rt,
                 _pd: PhantomData,
-            })
+            };
+            this.compile()
         } else {
             Err(Error::InvalidFunctionSignature)
         }
     }
 
+    #[inline]
+    pub(crate) fn compile(self) -> Result<Self> {
+        unsafe {
+            if (*self.raw).compiled.is_null() {
+                Error::from_ffi_res(ffi::Compile_Function(self.raw))?;
+            }
+        };
+        Ok(self)
+    }
+
     pub fn import_module_name(&self) -> &str {
         unsafe {
             std::str::from_utf8_unchecked(

diff --git a/src/module.rs b/src/module.rs
@@ -68,7 +68,7 @@ impl<'env, 'rt> Module<'env, 'rt> {
         ARGS: crate::WasmArgs,
         RET: crate::WasmType,
     {
-        let func = self.find_function::<ARGS, RET>(module_name, function_name)?;
+        let func = self.find_import_function(module_name, function_name)?;
         if Function::<'_, '_, ARGS, RET>::validate_sig(func) {
             unsafe { self.link_func_impl(func, f) }
             Ok(())
@@ -94,15 +94,11 @@ impl<'env, 'rt> Module<'env, 'rt> {
         }
     }
 
-    fn find_function<ARGS, RET>(
+    fn find_import_function(
         &self,
         module_name: &str,
         function_name: &str,
-    ) -> Result<ffi::IM3Function>
-    where
-        ARGS: crate::WasmArgs,
-        RET: crate::WasmType,
-    {
+    ) -> Result<ffi::IM3Function> {
         if let Some(func) = unsafe {
             std::slice::from_raw_parts_mut((*self.raw).functions, (*self.raw).numFunctions as usize)
                 .iter_mut()
@@ -115,6 +111,33 @@ impl<'env, 'rt> Module<'env, 'rt> {
         }
     }
 
+    pub fn find_function<ARGS, RET>(
+        &self,
+        function_name: &str,
+    ) -> Result<Function<'env, 'rt, ARGS, RET>>
+    where
+        ARGS: crate::WasmArgs,
+        RET: crate::WasmType,
+    {
+        if let Some(func) = unsafe {
+            let functions_ptr = (*self.raw).functions;
+            std::slice::from_raw_parts_mut(
+                if functions_ptr.is_null() {
+                    std::ptr::NonNull::dangling().as_ptr()
+                } else {
+                    functions_ptr
+                },
+                (*self.raw).numFunctions as usize,
+            )
+            .iter_mut()
+            .find(|func| eq_cstr_str(CStr::from_ptr(func.name), function_name))
+        } {
+            Function::from_raw(self.rt, func).and_then(Function::compile)
+        } else {
+            Err(Error::FunctionNotFound)
+        }
+    }
+
     #[cfg(feature = "wasi")]
     pub fn link_wasi(&mut self) {
         unsafe { ffi::m3_LinkWASI(self.raw) };

diff --git a/src/runtime.rs b/src/runtime.rs
@@ -1,4 +1,3 @@
-use std::ffi::CString;
 use std::mem;
 use std::ptr;
 use std::slice;
@@ -8,6 +7,7 @@ use crate::error::{Error, Result};
 use crate::function::Function;
 use crate::module::{Module, ParsedModule};
 
+#[derive(Debug)]
 pub struct Runtime<'env> {
     raw: ffi::IM3Runtime,
     environment: &'env Environment,
@@ -59,13 +59,15 @@ impl<'env> Runtime<'env> {
         ARGS: crate::WasmArgs,
         RET: crate::WasmType,
     {
-        unsafe {
-            let mut function = ptr::null_mut();
-            let name = CString::new(name).unwrap();
-            Error::from_ffi_res(ffi::m3_FindFunction(&mut function, self.raw, name.as_ptr()))
-                .map_err(|_| Error::FunctionNotFound)
-                .and_then(|_| Function::from_raw(self, function))
+        let mut module = unsafe { (*self.raw).modules };
+        while !module.is_null() {
+            match Module::from_raw(self, module).find_function::<ARGS, RET>(name) {
+                res @ Ok(_) => return res,
+                res @ Err(Error::InvalidFunctionSignature) => return res,
+                _ => module = unsafe { (*module).next },
+            }
         }
+        Err(Error::FunctionNotFound)
     }
 
     #[inline]

diff --git a/wasm3-sys/build.rs b/wasm3-sys/build.rs
@@ -6,7 +6,7 @@ use std::path::PathBuf;
 
 fn gen_bindings() -> io::Result<()> {
     let whitelist_regex =
-        "((?:I|c_)?[Mm]3.*)|(.*Page.*)|(Module_.*)|EmitWord_impl|op_CallRawFunction";
+        "((?:I|c_)?[Mm]3.*)|.*Page.*|Module_.*|EmitWord_impl|op_CallRawFunction|Compile_Function";
     let bindgen = bindgen::builder()
         .layout_tests(false)
         .generate_comments(false)