use crate::func::HostFunc;
use crate::instance::InstancePre;
use crate::store::StoreOpaque;
use crate::{
AsContextMut, Caller, Engine, Extern, ExternType, Func, FuncType, ImportType, Instance,
IntoFunc, Module, Trap, Val,
};
use anyhow::{anyhow, bail, Context, Error, Result};
use log::warn;
use std::collections::hash_map::{Entry, HashMap};
#[cfg(feature = "async")]
use std::future::Future;
use std::marker;
#[cfg(feature = "async")]
use std::pin::Pin;
use std::sync::Arc;
pub struct Linker<T> {
engine: Engine,
string2idx: HashMap<Arc<str>, usize>,
strings: Vec<Arc<str>>,
map: HashMap<ImportKey, Definition>,
allow_shadowing: bool,
allow_unknown_exports: bool,
_marker: marker::PhantomData<fn() -> T>,
}
#[derive(Copy, Clone, Hash, PartialEq, Eq)]
struct ImportKey {
name: usize,
module: usize,
}
#[derive(Clone)]
pub(crate) enum Definition {
Extern(Extern),
HostFunc(Arc<HostFunc>),
Instance(Arc<indexmap::IndexMap<String, Definition>>),
}
macro_rules! generate_wrap_async_func {
($num:tt $($args:ident)*) => (paste::paste!{
#[allow(non_snake_case)]
#[cfg(feature = "async")]
#[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
pub fn [<func_wrap $num _async>]<$($args,)* R>(
&mut self,
module: &str,
name: &str,
func: impl for<'a> Fn(Caller<'a, T>, $($args),*) -> Box<dyn Future<Output = R> + Send + 'a> + Send + Sync + 'static,
) -> Result<&mut Self>
where
$($args: crate::WasmTy,)*
R: crate::WasmRet,
{
assert!(
self.engine.config().async_support,
concat!(
"cannot use `func_wrap",
$num,
"_async` without enabling async support on the config",
),
);
self.func_wrap(module, name, move |mut caller: Caller<'_, T>, $($args: $args),*| {
let async_cx = caller.store.as_context_mut().opaque().async_cx();
let mut future = Pin::from(func(caller, $($args),*));
match unsafe { async_cx.block_on(future.as_mut()) } {
Ok(ret) => ret.into_fallible(),
Err(e) => R::fallible_from_trap(e),
}
})
}
})
}
impl<T> Linker<T> {
pub fn new(engine: &Engine) -> Linker<T> {
Linker {
engine: engine.clone(),
map: HashMap::new(),
string2idx: HashMap::new(),
strings: Vec::new(),
allow_shadowing: false,
allow_unknown_exports: false,
_marker: marker::PhantomData,
}
}
pub fn engine(&self) -> &Engine {
&self.engine
}
pub fn allow_shadowing(&mut self, allow: bool) -> &mut Self {
self.allow_shadowing = allow;
self
}
pub fn allow_unknown_exports(&mut self, allow: bool) -> &mut Self {
self.allow_unknown_exports = allow;
self
}
pub fn define(
&mut self,
module: &str,
name: &str,
item: impl Into<Extern>,
) -> Result<&mut Self> {
let key = self.import_key(module, Some(name));
self.insert(key, Definition::Extern(item.into()))?;
Ok(self)
}
pub fn define_name(&mut self, name: &str, item: impl Into<Extern>) -> Result<&mut Self> {
let key = self.import_key(name, None);
self.insert(key, Definition::Extern(item.into()))?;
Ok(self)
}
pub fn func_new(
&mut self,
module: &str,
name: &str,
ty: FuncType,
func: impl Fn(Caller<'_, T>, &[Val], &mut [Val]) -> Result<(), Trap> + Send + Sync + 'static,
) -> Result<&mut Self> {
let func = HostFunc::new(&self.engine, ty, func);
let key = self.import_key(module, Some(name));
self.insert(key, Definition::HostFunc(Arc::new(func)))?;
Ok(self)
}
#[cfg(feature = "async")]
#[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
pub fn func_new_async<F>(
&mut self,
module: &str,
name: &str,
ty: FuncType,
func: F,
) -> Result<&mut Self>
where
F: for<'a> Fn(
Caller<'a, T>,
&'a [Val],
&'a mut [Val],
) -> Box<dyn Future<Output = Result<(), Trap>> + Send + 'a>
+ Send
+ Sync
+ 'static,
{
assert!(
self.engine.config().async_support,
"cannot use `func_new_async` without enabling async support in the config"
);
self.func_new(module, name, ty, move |mut caller, params, results| {
let async_cx = caller.store.as_context_mut().opaque().async_cx();
let mut future = Pin::from(func(caller, params, results));
match unsafe { async_cx.block_on(future.as_mut()) } {
Ok(Ok(())) => Ok(()),
Ok(Err(trap)) | Err(trap) => Err(trap),
}
})
}
pub fn func_wrap<Params, Args>(
&mut self,
module: &str,
name: &str,
func: impl IntoFunc<T, Params, Args>,
) -> Result<&mut Self> {
let func = HostFunc::wrap(&self.engine, func);
let key = self.import_key(module, Some(name));
self.insert(key, Definition::HostFunc(Arc::new(func)))?;
Ok(self)
}
for_each_function_signature!(generate_wrap_async_func);
pub fn instance(
&mut self,
mut store: impl AsContextMut<Data = T>,
module_name: &str,
instance: Instance,
) -> Result<&mut Self> {
for export in instance.exports(store.as_context_mut()) {
let key = self.import_key(module_name, Some(export.name()));
self.insert(key, Definition::Extern(export.into_extern()))?;
}
Ok(self)
}
pub fn module(
&mut self,
mut store: impl AsContextMut<Data = T>,
module_name: &str,
module: &Module,
) -> Result<&mut Self>
where
T: 'static,
{
match ModuleKind::categorize(module)? {
ModuleKind::Command => self.command(store, module_name, module),
ModuleKind::Reactor => {
let instance = self.instantiate(&mut store, &module)?;
if let Some(export) = instance.get_export(&mut store, "_initialize") {
if let Extern::Func(func) = export {
func.typed::<(), (), _>(&store)
.and_then(|f| f.call(&mut store, ()).map_err(Into::into))
.context("calling the Reactor initialization function")?;
}
}
self.instance(store, module_name, instance)
}
}
}
fn command(
&mut self,
mut store: impl AsContextMut<Data = T>,
module_name: &str,
module: &Module,
) -> Result<&mut Self>
where
T: 'static,
{
for export in module.exports() {
if let Some(func_ty) = export.ty().func() {
let instance_pre = self.instantiate_pre(&mut store, module)?;
let export_name = export.name().to_owned();
let func = Func::new(
&mut store,
func_ty.clone(),
move |mut caller, params, results| {
let instance = instance_pre.instantiate(&mut caller)?;
let command_results = instance
.get_export(&mut caller, &export_name)
.unwrap()
.into_func()
.unwrap()
.call(&mut caller, params)
.map_err(|error| error.downcast::<Trap>().unwrap())?;
for (result, command_result) in
results.iter_mut().zip(command_results.into_vec())
{
*result = command_result;
}
Ok(())
},
);
let key = self.import_key(module_name, Some(export.name()));
self.insert(key, Definition::Extern(func.into()))?;
} else if export.name() == "memory" && export.ty().memory().is_some() {
} else if export.name() == "__indirect_function_table" && export.ty().table().is_some()
{
} else if export.name() == "table" && export.ty().table().is_some() {
} else if export.name() == "__data_end" && export.ty().global().is_some() {
warn!("command module exporting '__data_end' is deprecated");
} else if export.name() == "__heap_base" && export.ty().global().is_some() {
warn!("command module exporting '__heap_base' is deprecated");
} else if export.name() == "__dso_handle" && export.ty().global().is_some() {
warn!("command module exporting '__dso_handle' is deprecated")
} else if export.name() == "__rtti_base" && export.ty().global().is_some() {
warn!("command module exporting '__rtti_base' is deprecated; pass `--runtime half` to the AssemblyScript compiler");
} else if !self.allow_unknown_exports {
bail!("command export '{}' is not a function", export.name());
}
}
Ok(self)
}
pub fn alias(
&mut self,
module: &str,
name: &str,
as_module: &str,
as_name: &str,
) -> Result<&mut Self> {
let src = self.import_key(module, Some(name));
let dst = self.import_key(as_module, Some(as_name));
match self.map.get(&src).cloned() {
Some(item) => self.insert(dst, item)?,
None => bail!("no item named `{}::{}` defined", module, name),
}
Ok(self)
}
pub fn alias_module(&mut self, module: &str, as_module: &str) -> Result<()> {
let module = self.intern_str(module);
let as_module = self.intern_str(as_module);
let items = self
.map
.iter()
.filter(|(key, _def)| key.module == module)
.map(|(key, def)| (key.name, def.clone()))
.collect::<Vec<_>>();
for (name, item) in items {
self.insert(
ImportKey {
module: as_module,
name,
},
item,
)?;
}
Ok(())
}
fn insert(&mut self, key: ImportKey, item: Definition) -> Result<()> {
match self.map.entry(key) {
Entry::Occupied(_) if !self.allow_shadowing => {
let module = &self.strings[key.module];
let desc = match self.strings.get(key.name) {
Some(name) => format!("{}::{}", module, name),
None => module.to_string(),
};
bail!("import of `{}` defined twice", desc)
}
Entry::Occupied(mut o) => {
o.insert(item);
}
Entry::Vacant(v) => {
v.insert(item);
}
}
Ok(())
}
fn import_key(&mut self, module: &str, name: Option<&str>) -> ImportKey {
ImportKey {
module: self.intern_str(module),
name: name
.map(|name| self.intern_str(name))
.unwrap_or(usize::max_value()),
}
}
fn intern_str(&mut self, string: &str) -> usize {
if let Some(idx) = self.string2idx.get(string) {
return *idx;
}
let string: Arc<str> = string.into();
let idx = self.strings.len();
self.strings.push(string.clone());
self.string2idx.insert(string, idx);
idx
}
pub fn instantiate(
&self,
mut store: impl AsContextMut<Data = T>,
module: &Module,
) -> Result<Instance> {
self.instantiate_pre(&mut store, module)?.instantiate(store)
}
#[cfg(feature = "async")]
#[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
pub async fn instantiate_async(
&self,
mut store: impl AsContextMut<Data = T>,
module: &Module,
) -> Result<Instance>
where
T: Send,
{
self.instantiate_pre(&mut store, module)?
.instantiate_async(store)
.await
}
pub fn instantiate_pre(
&self,
mut store: impl AsContextMut<Data = T>,
module: &Module,
) -> Result<InstancePre<T>> {
let imports = module
.imports()
.map(|import| {
self._get_by_import(&import)
.ok_or_else(|| self.link_error(&import))
})
.collect::<Result<_>>()?;
unsafe { InstancePre::new(&mut store.as_context_mut().opaque(), module, imports) }
}
fn link_error(&self, import: &ImportType) -> Error {
let desc = match import.name() {
Some(name) => format!("{}::{}", import.module(), name),
None => import.module().to_string(),
};
anyhow!("unknown import: `{}` has not been defined", desc)
}
pub fn iter<'a: 'p, 'p>(
&'a self,
mut store: impl AsContextMut<Data = T> + 'p,
) -> impl Iterator<Item = (&str, &str, Extern)> + 'p {
self.map.iter().map(move |(key, item)| {
let mut store = store.as_context_mut().opaque();
(
&*self.strings[key.module],
&*self.strings[key.name],
unsafe { item.to_extern(&mut store) },
)
})
}
pub fn get(
&self,
mut store: impl AsContextMut<Data = T>,
module: &str,
name: Option<&str>,
) -> Option<Extern> {
let mut store = store.as_context_mut().opaque();
Some(unsafe { self._get(module, name)?.to_extern(&mut store) })
}
fn _get(&self, module: &str, name: Option<&str>) -> Option<&Definition> {
let key = ImportKey {
module: *self.string2idx.get(module)?,
name: match name {
Some(name) => *self.string2idx.get(name)?,
None => usize::max_value(),
},
};
self.map.get(&key)
}
pub fn get_by_import(
&self,
mut store: impl AsContextMut<Data = T>,
import: &ImportType,
) -> Option<Extern> {
let mut store = store.as_context_mut().opaque();
Some(unsafe { self._get_by_import(import)?.to_extern(&mut store) })
}
fn _get_by_import(&self, import: &ImportType) -> Option<Definition> {
if let Some(item) = self._get(import.module(), import.name()) {
return Some(item.clone());
}
if import.name().is_some() {
return None;
}
if let ExternType::Instance(t) = import.ty() {
let mut map = indexmap::IndexMap::new();
for export in t.exports() {
let item = self._get(import.module(), Some(export.name()))?;
map.insert(export.name().to_string(), item.clone());
}
return Some(Definition::Instance(Arc::new(map)));
}
None
}
pub fn get_default(
&self,
mut store: impl AsContextMut<Data = T>,
module: &str,
) -> Result<Func> {
if let Some(external) = self.get(&mut store, module, Some("")) {
if let Extern::Func(func) = external {
return Ok(func.clone());
}
bail!("default export in '{}' is not a function", module);
}
if let Some(external) = self.get(&mut store, module, Some("_start")) {
if let Extern::Func(func) = external {
return Ok(func.clone());
}
bail!("`_start` in '{}' is not a function", module);
}
Ok(Func::wrap(store, || {}))
}
}
impl<T> Default for Linker<T> {
fn default() -> Linker<T> {
Linker::new(&Engine::default())
}
}
impl Definition {
pub(crate) unsafe fn to_extern(&self, store: &mut StoreOpaque) -> Extern {
match self {
Definition::Extern(e) => e.clone(),
Definition::HostFunc(func) => func.to_func(store).into(),
Definition::Instance(i) => {
let items = Arc::new(
i.iter()
.map(|(name, item)| (name.clone(), item.to_extern(store)))
.collect(),
);
Instance::from_wasmtime(items, store).into()
}
}
}
pub(crate) fn comes_from_same_store(&self, store: &StoreOpaque) -> bool {
match self {
Definition::Extern(e) => e.comes_from_same_store(store),
Definition::HostFunc(_func) => true,
Definition::Instance(i) => i.values().all(|e| e.comes_from_same_store(store)),
}
}
}
enum ModuleKind {
Command,
Reactor,
}
impl ModuleKind {
fn categorize(module: &Module) -> Result<ModuleKind> {
let command_start = module.get_export("_start");
let reactor_start = module.get_export("_initialize");
match (command_start, reactor_start) {
(Some(command_start), None) => {
if let Some(_) = command_start.func() {
Ok(ModuleKind::Command)
} else {
bail!("`_start` must be a function")
}
}
(None, Some(reactor_start)) => {
if let Some(_) = reactor_start.func() {
Ok(ModuleKind::Reactor)
} else {
bail!("`_initialize` must be a function")
}
}
(None, None) => {
Ok(ModuleKind::Reactor)
}
(Some(_), Some(_)) => {
bail!("Program cannot be both a Command and a Reactor")
}
}
}
}