use crate::prelude::*;
use core::error::Error;
use core::fmt;
use core::mem;
use wasmtime_environ::PrimaryMap;
use wasmtime_environ::component::{
ComponentTypes, RuntimeComponentInstanceIndex, TypeResourceTableIndex,
};
const MAX_RESOURCE_HANDLE: u32 = 1 << 30;
pub struct ResourceTables<'a> {
pub guest: Option<(
&'a mut PrimaryMap<RuntimeComponentInstanceIndex, ResourceTable>,
&'a ComponentTypes,
)>,
pub host_table: Option<&'a mut ResourceTable>,
pub calls: &'a mut CallContexts,
}
#[derive(Default)]
pub struct ResourceTable {
next: u32,
slots: Vec<Slot>,
}
#[derive(Debug)]
pub enum TypedResource {
Host(u32),
Component {
rep: u32,
ty: TypeResourceTableIndex,
},
}
impl TypedResource {
fn rep(&self, access_ty: &TypedResourceIndex) -> Result<u32> {
match (self, access_ty) {
(Self::Host(rep), TypedResourceIndex::Host(_)) => Ok(*rep),
(Self::Host(_), expected) => bail!(ResourceTypeMismatch {
expected: *expected,
found: "host resource",
}),
(Self::Component { rep, ty }, TypedResourceIndex::Component { ty: expected, .. }) => {
if ty == expected {
Ok(*rep)
} else {
bail!(ResourceTypeMismatch {
expected: *access_ty,
found: "a different guest-defined resource",
})
}
}
(Self::Component { .. }, expected) => bail!(ResourceTypeMismatch {
expected: *expected,
found: "guest-defined resource",
}),
}
}
}
#[derive(Debug, Copy, Clone)]
pub enum TypedResourceIndex {
Host(u32),
Component {
index: u32,
ty: TypeResourceTableIndex,
},
}
impl TypedResourceIndex {
fn raw_index(&self) -> u32 {
match self {
Self::Host(index) | Self::Component { index, .. } => *index,
}
}
fn desc(&self) -> &'static str {
match self {
Self::Host(_) => "host resource",
Self::Component { .. } => "guest-defined resource",
}
}
}
enum Slot {
Free { next: u32 },
Own {
resource: TypedResource,
lend_count: u32,
},
Borrow {
resource: TypedResource,
scope: usize,
},
}
#[derive(Default)]
pub struct CallContexts {
scopes: Vec<CallContext>,
}
#[derive(Default)]
struct CallContext {
lenders: Vec<TypedResourceIndex>,
borrow_count: u32,
}
impl ResourceTables<'_> {
fn table_for_resource(&mut self, resource: &TypedResource) -> &mut ResourceTable {
match resource {
TypedResource::Host(_) => self.host_table.as_mut().unwrap(),
TypedResource::Component { ty, .. } => {
let (tables, types) = self.guest.as_mut().unwrap();
&mut tables[types[*ty].instance]
}
}
}
fn table_for_index(&mut self, index: &TypedResourceIndex) -> &mut ResourceTable {
match index {
TypedResourceIndex::Host(_) => self.host_table.as_mut().unwrap(),
TypedResourceIndex::Component { ty, .. } => {
let (tables, types) = self.guest.as_mut().unwrap();
&mut tables[types[*ty].instance]
}
}
}
pub fn resource_new(&mut self, resource: TypedResource) -> Result<u32> {
self.table_for_resource(&resource).insert(Slot::Own {
resource,
lend_count: 0,
})
}
pub fn resource_rep(&mut self, index: TypedResourceIndex) -> Result<u32> {
self.table_for_index(&index).rep(index)
}
pub fn resource_drop(&mut self, index: TypedResourceIndex) -> Result<Option<u32>> {
match self.table_for_index(&index).remove(index)? {
Slot::Own {
resource,
lend_count: 0,
} => resource.rep(&index).map(Some),
Slot::Own { .. } => bail!("cannot remove owned resource while borrowed"),
Slot::Borrow {
scope, resource, ..
} => {
resource.rep(&index)?;
self.calls.scopes[scope].borrow_count -= 1;
Ok(None)
}
Slot::Free { .. } => unreachable!(),
}
}
pub fn resource_lower_own(&mut self, resource: TypedResource) -> Result<u32> {
self.table_for_resource(&resource).insert(Slot::Own {
resource,
lend_count: 0,
})
}
pub fn resource_lift_own(&mut self, index: TypedResourceIndex) -> Result<u32> {
match self.table_for_index(&index).remove(index)? {
Slot::Own {
resource,
lend_count: 0,
} => resource.rep(&index),
Slot::Own { .. } => bail!("cannot remove owned resource while borrowed"),
Slot::Borrow { .. } => bail!("cannot lift own resource from a borrow"),
Slot::Free { .. } => unreachable!(),
}
}
pub fn resource_lift_borrow(&mut self, index: TypedResourceIndex) -> Result<u32> {
match self.table_for_index(&index).get_mut(index)? {
Slot::Own {
resource,
lend_count,
} => {
let rep = resource.rep(&index)?;
*lend_count = lend_count.checked_add(1).unwrap();
let scope = self.calls.scopes.last_mut().unwrap();
scope.lenders.push(index);
Ok(rep)
}
Slot::Borrow { resource, .. } => resource.rep(&index),
Slot::Free { .. } => unreachable!(),
}
}
pub fn resource_lower_borrow(&mut self, resource: TypedResource) -> Result<u32> {
let scope = self.calls.scopes.len() - 1;
let borrow_count = &mut self.calls.scopes.last_mut().unwrap().borrow_count;
*borrow_count = borrow_count.checked_add(1).unwrap();
self.table_for_resource(&resource)
.insert(Slot::Borrow { resource, scope })
}
#[inline]
pub fn enter_call(&mut self) {
self.calls.scopes.push(CallContext::default());
}
#[inline]
pub fn exit_call(&mut self) -> Result<()> {
let cx = self.calls.scopes.pop().unwrap();
if cx.borrow_count > 0 {
bail!("borrow handles still remain at the end of the call")
}
for lender in cx.lenders.iter() {
match self.table_for_index(lender).get_mut(*lender).unwrap() {
Slot::Own { lend_count, .. } => {
*lend_count -= 1;
}
_ => unreachable!(),
}
}
Ok(())
}
}
impl ResourceTable {
fn insert(&mut self, new: Slot) -> Result<u32> {
let next = self.next as usize;
if next == self.slots.len() {
self.slots.push(Slot::Free {
next: self.next.checked_add(1).unwrap(),
});
}
let ret = self.next;
self.next = match mem::replace(&mut self.slots[next], new) {
Slot::Free { next } => next,
_ => unreachable!(),
};
let ret = ret + 1;
if ret >= MAX_RESOURCE_HANDLE {
bail!("cannot allocate another handle: index overflow");
}
Ok(ret)
}
fn handle_index_to_table_index(&self, idx: u32) -> Option<usize> {
let idx = idx.checked_sub(1)?;
usize::try_from(idx).ok()
}
fn rep(&self, idx: TypedResourceIndex) -> Result<u32> {
let slot = self
.handle_index_to_table_index(idx.raw_index())
.and_then(|i| self.slots.get(i));
match slot {
None | Some(Slot::Free { .. }) => bail!(UnknownHandleIndex(idx)),
Some(Slot::Own { resource, .. } | Slot::Borrow { resource, .. }) => resource.rep(&idx),
}
}
fn get_mut(&mut self, idx: TypedResourceIndex) -> Result<&mut Slot> {
let slot = self
.handle_index_to_table_index(idx.raw_index())
.and_then(|i| self.slots.get_mut(i));
match slot {
None | Some(Slot::Free { .. }) => bail!(UnknownHandleIndex(idx)),
Some(other) => Ok(other),
}
}
fn remove(&mut self, idx: TypedResourceIndex) -> Result<Slot> {
let to_fill = Slot::Free { next: self.next };
let ret = mem::replace(self.get_mut(idx)?, to_fill);
self.next = idx.raw_index() - 1;
Ok(ret)
}
}
#[derive(Debug)]
struct UnknownHandleIndex(TypedResourceIndex);
impl fmt::Display for UnknownHandleIndex {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "unknown handle index {}", self.0.raw_index())
}
}
impl Error for UnknownHandleIndex {}
#[derive(Debug)]
struct ResourceTypeMismatch {
expected: TypedResourceIndex,
found: &'static str,
}
impl fmt::Display for ResourceTypeMismatch {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"handle index {} used with the wrong type, \
expected {} but found {}",
self.expected.raw_index(),
self.expected.desc(),
self.found,
)
}
}
impl Error for ResourceTypeMismatch {}