use crate::func::HostFunc;
use crate::instance::InstancePre;
use crate::store::StoreOpaque;
use crate::{
AsContext, AsContextMut, Caller, Engine, Extern, ExternType, Func, FuncType, ImportType,
Instance, IntoFunc, Module, StoreContextMut, Val, ValRaw, ValType,
};
use anyhow::{bail, Context, 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>,
}
impl<T> Clone for Linker<T> {
fn clone(&self) -> Linker<T> {
Linker {
engine: self.engine.clone(),
string2idx: self.string2idx.clone(),
strings: self.strings.clone(),
map: self.map.clone(),
allow_shadowing: self.allow_shadowing,
allow_unknown_exports: self.allow_unknown_exports,
_marker: self._marker,
}
}
}
#[derive(Copy, Clone, Hash, PartialEq, Eq)]
struct ImportKey {
name: usize,
module: usize,
}
#[derive(Clone)]
pub(crate) enum Definition {
Extern(Extern, DefinitionType),
HostFunc(Arc<HostFunc>),
}
#[derive(Clone)]
pub(crate) enum DefinitionType {
Func(wasmtime_runtime::VMSharedSignatureIndex),
Global(wasmtime_environ::Global),
Table(wasmtime_environ::Table, u32),
Memory(wasmtime_environ::Memory, u64),
}
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().0.async_cx().expect("Attempt to start async function on dying fiber");
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_error(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
}
#[cfg(any(feature = "cranelift", feature = "winch"))]
#[cfg_attr(nightlydoc, doc(cfg(any(feature = "cranelift", feature = "winch"))))]
pub fn define_unknown_imports_as_traps(&mut self, module: &Module) -> anyhow::Result<()> {
for import in module.imports() {
if let Err(import_err) = self._get_by_import(&import) {
if let ExternType::Func(func_ty) = import_err.ty() {
self.func_new(import.module(), import.name(), func_ty, move |_, _, _| {
bail!(import_err.clone());
})?;
}
}
}
Ok(())
}
#[cfg(any(feature = "cranelift", feature = "winch"))]
#[cfg_attr(nightlydoc, doc(cfg(any(feature = "cranelift", feature = "winch"))))]
pub fn define_unknown_imports_as_default_values(
&mut self,
module: &Module,
) -> anyhow::Result<()> {
for import in module.imports() {
if let Err(import_err) = self._get_by_import(&import) {
if let ExternType::Func(func_ty) = import_err.ty() {
let result_tys: Vec<_> = func_ty.results().collect();
self.func_new(
import.module(),
import.name(),
func_ty,
move |_caller, _args, results| {
for (result, ty) in results.iter_mut().zip(&result_tys) {
*result = match ty {
ValType::I32 => Val::I32(0),
ValType::I64 => Val::I64(0),
ValType::F32 => Val::F32(0.0_f32.to_bits()),
ValType::F64 => Val::F64(0.0_f64.to_bits()),
ValType::V128 => Val::V128(0_u128.into()),
ValType::FuncRef => Val::FuncRef(None),
ValType::ExternRef => Val::ExternRef(None),
};
}
Ok(())
},
)?;
}
}
}
Ok(())
}
pub fn define(
&mut self,
store: impl AsContext<Data = T>,
module: &str,
name: &str,
item: impl Into<Extern>,
) -> Result<&mut Self> {
let store = store.as_context();
let key = self.import_key(module, Some(name));
self.insert(key, Definition::new(store.0, item.into()))?;
Ok(self)
}
pub fn define_name(
&mut self,
store: impl AsContext<Data = T>,
name: &str,
item: impl Into<Extern>,
) -> Result<&mut Self> {
let store = store.as_context();
let key = self.import_key(name, None);
self.insert(key, Definition::new(store.0, item.into()))?;
Ok(self)
}
#[cfg(any(feature = "cranelift", feature = "winch"))]
#[cfg_attr(nightlydoc, doc(cfg(any(feature = "cranelift", feature = "winch"))))]
pub fn func_new(
&mut self,
module: &str,
name: &str,
ty: FuncType,
func: impl Fn(Caller<'_, T>, &[Val], &mut [Val]) -> Result<()> + 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(any(feature = "cranelift", feature = "winch"))]
#[cfg_attr(nightlydoc, doc(cfg(any(feature = "cranelift", feature = "winch"))))]
pub unsafe fn func_new_unchecked(
&mut self,
module: &str,
name: &str,
ty: FuncType,
func: impl Fn(Caller<'_, T>, &mut [ValRaw]) -> Result<()> + Send + Sync + 'static,
) -> Result<&mut Self> {
let func = HostFunc::new_unchecked(&self.engine, ty, func);
let key = self.import_key(module, Some(name));
self.insert(key, Definition::HostFunc(Arc::new(func)))?;
Ok(self)
}
#[cfg(all(feature = "async", feature = "cranelift"))]
#[cfg_attr(nightlydoc, doc(cfg(all(feature = "async", feature = "cranelift"))))]
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<()>> + 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()
.0
.async_cx()
.expect("Attempt to spawn new function on dying fiber");
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> {
let mut store = store.as_context_mut();
let exports = instance
.exports(&mut store)
.map(|e| {
(
self.import_key(module_name, Some(e.name())),
e.into_extern(),
)
})
.collect::<Vec<_>>();
for (key, export) in exports {
self.insert(key, Definition::new(store.0, export))?;
}
Ok(self)
}
#[cfg(any(feature = "cranelift", feature = "winch"))]
#[cfg_attr(nightlydoc, doc(cfg(any(feature = "cranelift", feature = "winch"))))]
pub fn module(
&mut self,
mut store: impl AsContextMut<Data = T>,
module_name: &str,
module: &Module,
) -> Result<&mut Self>
where
T: 'static,
{
assert!(
Engine::same(&self.engine, store.as_context().engine()),
"different engines for this linker and the store provided"
);
match ModuleKind::categorize(module)? {
ModuleKind::Command => {
self.command(
store,
module_name,
module,
|store, func_ty, export_name, instance_pre| {
Func::new(
store,
func_ty.clone(),
move |mut caller, params, results| {
let instance = instance_pre.instantiate(&mut caller)?;
instance
.get_export(&mut caller, &export_name)
.unwrap()
.into_func()
.unwrap()
.call(&mut caller, params, results)?;
Ok(())
},
)
},
)
}
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)
}
}
}
#[cfg(all(feature = "async", feature = "cranelift"))]
#[cfg_attr(nightlydoc, doc(cfg(all(feature = "async", feature = "cranelift"))))]
pub async fn module_async(
&mut self,
mut store: impl AsContextMut<Data = T>,
module_name: &str,
module: &Module,
) -> Result<&mut Self>
where
T: Send + 'static,
{
assert!(
Engine::same(&self.engine, store.as_context().engine()),
"different engines for this linker and the store provided"
);
match ModuleKind::categorize(module)? {
ModuleKind::Command => self.command(
store,
module_name,
module,
|store, func_ty, export_name, instance_pre| {
let upvars = Arc::new((instance_pre, export_name));
Func::new_async(
store,
func_ty.clone(),
move |mut caller, params, results| {
let upvars = upvars.clone();
Box::new(async move {
let (instance_pre, export_name) = &*upvars;
let instance = instance_pre.instantiate_async(&mut caller).await?;
instance
.get_export(&mut caller, &export_name)
.unwrap()
.into_func()
.unwrap()
.call_async(&mut caller, params, results)
.await?;
Ok(())
})
},
)
},
),
ModuleKind::Reactor => {
let instance = self.instantiate_async(&mut store, &module).await?;
if let Some(export) = instance.get_export(&mut store, "_initialize") {
if let Extern::Func(func) = export {
let func = func
.typed::<(), ()>(&store)
.context("loading the Reactor initialization function")?;
func.call_async(&mut store, ())
.await
.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,
mk_func: impl Fn(&mut StoreContextMut<T>, &FuncType, String, InstancePre<T>) -> Func,
) -> Result<&mut Self>
where
T: 'static,
{
let mut store = store.as_context_mut();
for export in module.exports() {
if let Some(func_ty) = export.ty().func() {
let instance_pre = self.instantiate_pre(module)?;
let export_name = export.name().to_owned();
let func = mk_func(&mut store, func_ty, export_name, instance_pre);
let key = self.import_key(module_name, Some(export.name()));
self.insert(key, Definition::new(store.0, 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(module, Some(store.as_context_mut().0))?
.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(module, Some(store.as_context_mut().0))?
.instantiate_async(store)
.await
}
pub fn instantiate_pre(&self, module: &Module) -> Result<InstancePre<T>> {
self._instantiate_pre(module, None)
}
fn _instantiate_pre(
&self,
module: &Module,
store: Option<&StoreOpaque>,
) -> Result<InstancePre<T>> {
let mut imports = module
.imports()
.map(|import| self._get_by_import(&import))
.collect::<Result<Vec<_>, _>>()?;
if let Some(store) = store {
for import in imports.iter_mut() {
import.update_size(store);
}
}
unsafe { InstancePre::new(module, imports) }
}
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 store = store.as_context_mut();
(
&*self.strings[key.module],
&*self.strings[key.name],
unsafe { item.to_extern(store.0) },
)
})
}
pub fn get(
&self,
mut store: impl AsContextMut<Data = T>,
module: &str,
name: &str,
) -> Option<Extern> {
let store = store.as_context_mut().0;
Some(unsafe { self._get(module, name)?.to_extern(store) })
}
fn _get(&self, module: &str, name: &str) -> Option<&Definition> {
let key = ImportKey {
module: *self.string2idx.get(module)?,
name: *self.string2idx.get(name)?,
};
self.map.get(&key)
}
pub fn get_by_import(
&self,
mut store: impl AsContextMut<Data = T>,
import: &ImportType,
) -> Option<Extern> {
let store = store.as_context_mut().0;
Some(unsafe { self._get_by_import(import).ok()?.to_extern(store) })
}
fn _get_by_import(&self, import: &ImportType) -> Result<Definition, UnknownImportError> {
match self._get(import.module(), import.name()) {
Some(item) => Ok(item.clone()),
None => Err(UnknownImportError::new(import)),
}
}
pub fn get_default(
&self,
mut store: impl AsContextMut<Data = T>,
module: &str,
) -> Result<Func> {
if let Some(external) = self.get(&mut store, module, "") {
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, "_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 {
fn new(store: &StoreOpaque, item: Extern) -> Definition {
let ty = DefinitionType::from(store, &item);
Definition::Extern(item, ty)
}
pub(crate) fn ty(&self) -> DefinitionType {
match self {
Definition::Extern(_, ty) => ty.clone(),
Definition::HostFunc(func) => DefinitionType::Func(func.sig_index()),
}
}
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(),
}
}
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,
}
}
fn update_size(&mut self, store: &StoreOpaque) {
match self {
Definition::Extern(Extern::Memory(m), DefinitionType::Memory(_, size)) => {
*size = m.internal_size(store);
}
Definition::Extern(Extern::SharedMemory(m), DefinitionType::Memory(_, size)) => {
*size = m.size();
}
Definition::Extern(Extern::Table(m), DefinitionType::Table(_, size)) => {
*size = m.internal_size(store);
}
_ => {}
}
}
}
impl DefinitionType {
pub(crate) fn from(store: &StoreOpaque, item: &Extern) -> DefinitionType {
let data = store.store_data();
match item {
Extern::Func(f) => DefinitionType::Func(f.sig_index(data)),
Extern::Table(t) => DefinitionType::Table(*t.wasmtime_ty(data), t.internal_size(store)),
Extern::Global(t) => DefinitionType::Global(*t.wasmtime_ty(data)),
Extern::Memory(t) => {
DefinitionType::Memory(*t.wasmtime_ty(data), t.internal_size(store))
}
Extern::SharedMemory(t) => DefinitionType::Memory(*t.ty().wasmtime_memory(), t.size()),
}
}
pub(crate) fn desc(&self) -> &'static str {
match self {
DefinitionType::Func(_) => "function",
DefinitionType::Table(..) => "table",
DefinitionType::Memory(..) => "memory",
DefinitionType::Global(_) => "global",
}
}
}
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")
}
}
}
}
#[derive(Clone, Debug)]
pub struct UnknownImportError {
module: String,
name: String,
ty: ExternType,
}
impl UnknownImportError {
fn new(import: &ImportType) -> Self {
Self {
module: import.module().to_string(),
name: import.name().to_string(),
ty: import.ty(),
}
}
pub fn module(&self) -> &str {
&self.module
}
pub fn name(&self) -> &str {
&self.name
}
pub fn ty(&self) -> ExternType {
self.ty.clone()
}
}
impl std::fmt::Display for UnknownImportError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"unknown import: `{}::{}` has not been defined",
self.module, self.name,
)
}
}
impl std::error::Error for UnknownImportError {}