More module accessors for runtime

wasm3 · Jan 28, 2020 · af998de · af998de
1 parent ef7cda4
commit af998de
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 20 deletions.
diff --git a/src/error.rs b/src/error.rs
@@ -9,6 +9,7 @@ pub enum Error {
     Wasm3(&'static str),
     InvalidFunctionSignature,
     FunctionNotFound,
+    ModuleNotFound,
 }
 
 impl Error {
@@ -31,6 +32,7 @@ impl fmt::Display for Error {
                 write!(f, "the found function had an unexpected signature")
             }
             Error::FunctionNotFound => write!(f, "the function could not be found"),
+            Error::ModuleNotFound => write!(f, "the module could not be found"),
         }
     }
 }
diff --git a/src/macros.rs b/src/macros.rs
@@ -7,8 +7,8 @@ macro_rules! make_func_wrapper {
         $wis unsafe extern "C" fn $wrapper_name(
             _rt: ffi::IM3Runtime,
             _sp: *mut u64,
-            _mem: *mut std::ffi::c_void,
-        ) -> *const std::ffi::c_void {
+            _mem: *mut core::ffi::c_void,
+        ) -> *const core::ffi::c_void {
             let ssp = _sp;
             $(
                 let $pname = $crate::read_stack_param!(_sp -> $ptype);

diff --git a/src/module.rs b/src/module.rs
@@ -6,7 +6,7 @@ use crate::environment::Environment;
 use crate::error::{Error, Result};
 use crate::function::{Function, NNM3Function, RawCall};
 use crate::runtime::Runtime;
-use crate::utils::eq_cstr_str;
+use crate::utils::{cstr_to_str, eq_cstr_str};
 
 pub struct ParsedModule<'env> {
     raw: ffi::IM3Module,
@@ -101,12 +101,12 @@ impl<'env, 'rt> Module<'env, 'rt> {
                 },
                 (*self.raw).numFunctions as usize,
             )
+            .iter_mut()
+            .filter(|func| eq_cstr_str(func.import.moduleUtf8, module_name))
+            .find(|func| eq_cstr_str(func.import.fieldUtf8, function_name))
+            .map(NonNull::from)
+            .ok_or(Error::FunctionNotFound)
         }
-        .iter_mut()
-        .filter(|func| eq_cstr_str(func.import.moduleUtf8, module_name))
-        .find(|func| eq_cstr_str(func.import.fieldUtf8, function_name))
-        .map(NonNull::from)
-        .ok_or(Error::FunctionNotFound)
     }
 
     pub fn find_function<ARGS, RET>(
@@ -134,6 +134,34 @@ impl<'env, 'rt> Module<'env, 'rt> {
         Function::from_raw(self.rt, func).and_then(Function::compile)
     }
 
+    pub fn function<ARGS, RET>(
+        &self,
+        function_index: usize,
+    ) -> Result<Function<'env, 'rt, ARGS, RET>>
+    where
+        ARGS: crate::WasmArgs,
+        RET: crate::WasmType,
+    {
+        let func = unsafe {
+            slice::from_raw_parts_mut(
+                if (*self.raw).functions.is_null() {
+                    NonNull::dangling().as_ptr()
+                } else {
+                    (*self.raw).functions
+                },
+                (*self.raw).numFunctions as usize,
+            )
+            .get(function_index)
+            .map(NonNull::from)
+            .ok_or(Error::FunctionNotFound)?
+        };
+        Function::from_raw(self.rt, func).and_then(Function::compile)
+    }
+
+    pub fn name(&self) -> &str {
+        unsafe { cstr_to_str((*self.raw).name) }
+    }
+
     #[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
@@ -6,6 +6,7 @@ use crate::environment::Environment;
 use crate::error::{Error, Result};
 use crate::function::Function;
 use crate::module::{Module, ParsedModule};
+use crate::utils::eq_cstr_str;
 
 #[derive(Debug)]
 pub struct Runtime<'env> {
@@ -59,15 +60,36 @@ impl<'env> Runtime<'env> {
         ARGS: crate::WasmArgs,
         RET: crate::WasmType,
     {
-        let mut module = unsafe { ptr::NonNull::new((*self.raw).modules) };
-        while let Some(raw_mod) = module {
-            match Module::from_raw(self, raw_mod.as_ptr()).find_function::<ARGS, RET>(name) {
-                res @ Ok(_) => return res,
-                res @ Err(Error::InvalidFunctionSignature) => return res,
-                _ => module = unsafe { ptr::NonNull::new(raw_mod.as_ref().next) },
+        self.modules()
+            .find_map(|module| match module.find_function::<ARGS, RET>(name) {
+                res @ Ok(_) => Some(res),
+                res @ Err(Error::InvalidFunctionSignature) => Some(res),
+                _ => None,
+            })
+            .unwrap_or(Err(Error::FunctionNotFound))
+    }
+
+    /// Using this over searching through [`modules`] is a bit more efficient.
+    pub fn find_module<'rt>(&'rt self, name: &str) -> Result<Module<'env, 'rt>> {
+        unsafe {
+            let mut module = ptr::NonNull::new((*self.raw).modules);
+            while let Some(raw_mod) = module {
+                if eq_cstr_str(raw_mod.as_ref().name, name) {
+                    return Ok(Module::from_raw(self, raw_mod.as_ptr()));
+                }
+                module = ptr::NonNull::new(raw_mod.as_ref().next);
             }
+            Err(Error::ModuleNotFound)
         }
-        Err(Error::FunctionNotFound)
+    }
+
+    pub fn modules<'rt>(&'rt self) -> impl Iterator<Item = Module<'env, 'rt>> + 'rt {
+        // pointer could get invalidated if modules can become unloaded
+        let mut module = unsafe { ptr::NonNull::new((*self.raw).modules) };
+        core::iter::from_fn(move || {
+            let next = unsafe { module.and_then(|module| ptr::NonNull::new(module.as_ref().next)) };
+            mem::replace(&mut module, next).map(|raw| Module::from_raw(self, raw.as_ptr()))
+        })
     }
 
     #[inline]

diff --git a/src/utils.rs b/src/utils.rs
@@ -16,19 +16,17 @@ pub unsafe fn cstr_to_str<'a>(ptr: *const libc::c_char) -> &'a str {
     core::str::from_utf8_unchecked(bytes_till_null(ptr))
 }
 
-pub fn eq_cstr_str(cstr: *const libc::c_char, str: &str) -> bool {
+pub unsafe fn eq_cstr_str(cstr: *const libc::c_char, str: &str) -> bool {
     if cstr.is_null() {
         return false;
     }
     let mut bytes = str.as_bytes().iter();
     let mut cstr = cstr.cast::<u8>();
     loop {
-        match (bytes.next(), unsafe { *cstr }) {
+        match (bytes.next(), *cstr) {
             (None, 0) => break true,
             (Some(_), 0) => break false,
-            (Some(&byte), cbyte) if cbyte == byte => unsafe {
-                cstr = cstr.add(1);
-            },
+            (Some(&byte), cbyte) if cbyte == byte => cstr = cstr.add(1),
             _ => break false,
         }
     }