wasmtime 50.0.0

//! Implementation of [`FutureAny`] and [`StreamAny`].

use crate::component::concurrent::futures_and_streams::{self, TransmitOrigin};
use crate::component::concurrent::{TableId, TransmitHandle};
use crate::component::func::{LiftContext, LowerContext, desc};
use crate::component::matching::InstanceType;
use crate::component::types::{FutureType, StreamType};
use crate::component::{
    ComponentInstanceId, ComponentType, FutureReader, Lift, Lower, StreamReader,
};
use crate::store::StoreOpaque;
use crate::{AsContextMut, Result, bail};
use core::any::TypeId;
use core::mem::MaybeUninit;
use wasmtime_environ::component::{
    CanonicalAbiInfo, InterfaceType, TypeFutureTableIndex, TypeStreamTableIndex,
};

/// Represents a type-erased component model `future`.
///
/// This type is similar to [`ResourceAny`](crate::component::ResourceAny)
/// where it's a static guarantee that it represents a component model
/// `future`, but it does not contain any information about the underlying type
/// that is associated with this future. This is intended to be used in
/// "dynamically typed" situations where embedders may not know ahead of time
/// the type of a `future` being used by component that is loaded.
///
/// # Closing futures
///
/// A [`FutureAny`] represents a resource that is owned by a [`Store`]. Proper
/// disposal of a future requires invoking the [`FutureAny::close`] method to
/// ensure that this handle does not leak. If [`FutureAny::close`] is not
/// called then memory will not be leaked once the owning [`Store`] is dropped,
/// but the resource handle will be leaked until the [`Store`] is dropped.
///
/// [`Store`]: crate::Store
#[derive(Debug, Clone, PartialEq)]
pub struct FutureAny {
    id: TableId<TransmitHandle>,
    ty: PayloadType<FutureType>,
}

impl FutureAny {
    /// Attempts to convert this [`FutureAny`] to a [`FutureReader<T>`]
    /// with a statically known type.
    ///
    /// # Errors
    ///
    /// This function will return an error if `T` does not match the type of
    /// value on this future.
    pub fn try_into_future_reader<T>(self) -> Result<FutureReader<T>>
    where
        T: ComponentType + 'static,
    {
        self.ty
            .typecheck_host::<T>(FutureType::equivalent_payload_host::<T>)?;
        Ok(FutureReader::new_(self.id))
    }

    /// Attempts to convert `reader` to a [`FutureAny`], erasing its statically
    /// known type.
    ///
    /// # Errors
    ///
    /// This function will return an error if `reader` does not belong to
    /// `store`.
    pub fn try_from_future_reader<T>(
        mut store: impl AsContextMut,
        reader: FutureReader<T>,
    ) -> Result<Self>
    where
        T: ComponentType + 'static,
    {
        let store = store.as_context_mut();
        let ty = match store.0.transmit_origin(reader.id())? {
            TransmitOrigin::Host { .. } => PayloadType::new_host::<T>(),
            TransmitOrigin::GuestFuture(id, ty) => PayloadType::new_guest_future(store.0, id, ty),
            TransmitOrigin::GuestStream(..) => bail!("not a future"),
        };
        Ok(FutureAny {
            id: reader.id(),
            ty,
        })
    }

    fn lift_from_index(cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32) -> Result<Self> {
        let id = futures_and_streams::lift_index_to_future(cx, ty, index)?;
        let InterfaceType::Future(ty) = ty else {
            unreachable!()
        };
        let ty = cx.types[ty].ty;
        Ok(FutureAny {
            id,
            // Note that this future might actually be a host-originating
            // future which means that this ascription of "the type is the
            // guest" may be slightly in accurate. The guest, however, has the
            // most accurate view of what type this future has so that should
            // be reasonable to ascribe as the type here regardless.
            ty: PayloadType::Guest(FutureType::from(ty, &cx.instance_type())),
        })
    }

    /// Close this `FutureAny`.
    ///
    /// This will close this future and cause any write that happens later to
    /// returned `DROPPED`.
    ///
    /// # Errors
    ///
    /// Returns an error if this future has already been closed.
    ///
    /// # Panics
    ///
    /// Panics if the `store` does not own this future.
    pub fn close(&mut self, mut store: impl AsContextMut) -> Result<()> {
        futures_and_streams::future_close(store.as_context_mut().0, &mut self.id)
    }
}

unsafe impl ComponentType for FutureAny {
    const ABI: CanonicalAbiInfo = CanonicalAbiInfo::SCALAR4;

    type Lower = <u32 as ComponentType>::Lower;

    fn typecheck(ty: &InterfaceType, _types: &InstanceType<'_>) -> Result<()> {
        match ty {
            InterfaceType::Future(_) => Ok(()),
            other => bail!("expected `future`, found `{}`", desc(other)),
        }
    }
}

unsafe impl Lower for FutureAny {
    fn linear_lower_to_flat<T>(
        &self,
        cx: &mut LowerContext<'_, T>,
        ty: InterfaceType,
        dst: &mut MaybeUninit<Self::Lower>,
    ) -> Result<()> {
        futures_and_streams::lower_future_to_index(self.id, cx, ty)?.linear_lower_to_flat(
            cx,
            InterfaceType::U32,
            dst,
        )
    }

    fn linear_lower_to_memory<T>(
        &self,
        cx: &mut LowerContext<'_, T>,
        ty: InterfaceType,
        offset: usize,
    ) -> Result<()> {
        futures_and_streams::lower_future_to_index(self.id, cx, ty)?.linear_lower_to_memory(
            cx,
            InterfaceType::U32,
            offset,
        )
    }
}

unsafe impl Lift for FutureAny {
    fn linear_lift_from_flat(
        cx: &mut LiftContext<'_>,
        ty: InterfaceType,
        src: &Self::Lower,
    ) -> Result<Self> {
        let index = u32::linear_lift_from_flat(cx, InterfaceType::U32, src)?;
        Self::lift_from_index(cx, ty, index)
    }

    fn linear_lift_from_memory(
        cx: &mut LiftContext<'_>,
        ty: InterfaceType,
        bytes: &[u8],
    ) -> Result<Self> {
        let index = u32::linear_lift_from_memory(cx, InterfaceType::U32, bytes)?;
        Self::lift_from_index(cx, ty, index)
    }
}

/// Represents a type-erased component model `stream`.
///
/// This type is similar to [`ResourceAny`](crate::component::ResourceAny)
/// where it's a static guarantee that it represents a component model
/// `stream`, but it does not contain any information about the underlying type
/// that is associated with this stream. This is intended to be used in
/// "dynamically typed" situations where embedders may not know ahead of time
/// the type of a `stream` being used by component that is loaded.
///
/// # Closing streams
///
/// A [`StreamAny`] represents a resource that is owned by a [`Store`]. Proper
/// disposal of a stream requires invoking the [`StreamAny::close`] method to
/// ensure that this handle does not leak. If [`StreamAny::close`] is not
/// called then memory will not be leaked once the owning [`Store`] is dropped,
/// but the resource handle will be leaked until the [`Store`] is dropped.
///
/// [`Store`]: crate::Store
#[derive(Debug, Clone, PartialEq)]
pub struct StreamAny {
    id: TableId<TransmitHandle>,
    ty: PayloadType<StreamType>,
}

impl StreamAny {
    /// Attempts to convert this [`StreamAny`] to a [`StreamReader<T>`]
    /// with a statically known type.
    ///
    /// # Errors
    ///
    /// This function will return an error if `T` does not match the type of
    /// value on this stream.
    pub fn try_into_stream_reader<T>(self) -> Result<StreamReader<T>>
    where
        T: ComponentType + 'static,
    {
        self.ty
            .typecheck_host::<T>(StreamType::equivalent_payload_host::<T>)?;
        Ok(StreamReader::new_(self.id))
    }

    /// Attempts to convert `reader` to a [`StreamAny`], erasing its statically
    /// known type.
    ///
    /// # Errors
    ///
    /// This function will return an error if `reader` does not belong to
    /// `store`.
    pub fn try_from_stream_reader<T>(
        mut store: impl AsContextMut,
        reader: StreamReader<T>,
    ) -> Result<Self>
    where
        T: ComponentType + 'static,
    {
        let store = store.as_context_mut();
        let ty = match store.0.transmit_origin(reader.id())? {
            TransmitOrigin::Host { .. } => PayloadType::new_host::<T>(),
            TransmitOrigin::GuestStream(id, ty) => PayloadType::new_guest_stream(store.0, id, ty),
            TransmitOrigin::GuestFuture(..) => bail!("not a stream"),
        };
        Ok(StreamAny {
            id: reader.id(),
            ty,
        })
    }

    fn lift_from_index(cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32) -> Result<Self> {
        let id = futures_and_streams::lift_index_to_stream(cx, ty, index)?;
        let InterfaceType::Stream(ty) = ty else {
            unreachable!()
        };
        let ty = cx.types[ty].ty;
        Ok(StreamAny {
            id,
            // Note that this stream might actually be a host-originating, but
            // see the documentation in `FutureAny::lift_from_index` for why
            // this should be ok.
            ty: PayloadType::Guest(StreamType::from(ty, &cx.instance_type())),
        })
    }

    /// Close this `StreamAny`.
    ///
    /// This will close this stream and cause any write that happens later to
    /// returned `DROPPED`.
    ///
    /// # Errors
    ///
    /// Returns an error if this stream has already been closed.
    ///
    /// # Panics
    ///
    /// Panics if the `store` does not own this stream.
    pub fn close(&mut self, mut store: impl AsContextMut) -> Result<()> {
        futures_and_streams::stream_close(store.as_context_mut().0, &mut self.id)
    }
}

unsafe impl ComponentType for StreamAny {
    const ABI: CanonicalAbiInfo = CanonicalAbiInfo::SCALAR4;

    type Lower = <u32 as ComponentType>::Lower;

    fn typecheck(ty: &InterfaceType, _types: &InstanceType<'_>) -> Result<()> {
        match ty {
            InterfaceType::Stream(_) => Ok(()),
            other => bail!("expected `stream`, found `{}`", desc(other)),
        }
    }
}

unsafe impl Lower for StreamAny {
    fn linear_lower_to_flat<T>(
        &self,
        cx: &mut LowerContext<'_, T>,
        ty: InterfaceType,
        dst: &mut MaybeUninit<Self::Lower>,
    ) -> Result<()> {
        futures_and_streams::lower_stream_to_index(self.id, cx, ty)?.linear_lower_to_flat(
            cx,
            InterfaceType::U32,
            dst,
        )
    }

    fn linear_lower_to_memory<T>(
        &self,
        cx: &mut LowerContext<'_, T>,
        ty: InterfaceType,
        offset: usize,
    ) -> Result<()> {
        futures_and_streams::lower_stream_to_index(self.id, cx, ty)?.linear_lower_to_memory(
            cx,
            InterfaceType::U32,
            offset,
        )
    }
}

unsafe impl Lift for StreamAny {
    fn linear_lift_from_flat(
        cx: &mut LiftContext<'_>,
        ty: InterfaceType,
        src: &Self::Lower,
    ) -> Result<Self> {
        let index = u32::linear_lift_from_flat(cx, InterfaceType::U32, src)?;
        Self::lift_from_index(cx, ty, index)
    }

    fn linear_lift_from_memory(
        cx: &mut LiftContext<'_>,
        ty: InterfaceType,
        bytes: &[u8],
    ) -> Result<Self> {
        let index = u32::linear_lift_from_memory(cx, InterfaceType::U32, bytes)?;
        Self::lift_from_index(cx, ty, index)
    }
}

#[derive(Debug, Clone, PartialEq)]
enum PayloadType<T> {
    Guest(T),
    Host(TypeId),
}

impl PayloadType<FutureType> {
    fn new_guest_future(
        store: &StoreOpaque,
        id: ComponentInstanceId,
        ty: TypeFutureTableIndex,
    ) -> Self {
        let types = InstanceType::new(&store.component_instance(id));
        let ty = types.types[ty].ty;
        PayloadType::Guest(FutureType::from(ty, &types))
    }
}

impl PayloadType<StreamType> {
    fn new_guest_stream(
        store: &StoreOpaque,
        id: ComponentInstanceId,
        ty: TypeStreamTableIndex,
    ) -> Self {
        let types = InstanceType::new(&store.component_instance(id));
        let ty = types.types[ty].ty;
        PayloadType::Guest(StreamType::from(ty, &types))
    }
}

impl<T> PayloadType<T> {
    fn new_host<P>() -> Self
    where
        P: ComponentType + 'static,
    {
        PayloadType::Host(TypeId::of::<P>())
    }

    fn typecheck_host<P>(&self, equivalent: fn(&T) -> Result<()>) -> Result<()>
    where
        P: ComponentType + 'static,
    {
        match self {
            Self::Guest(ty) => equivalent(ty),
            Self::Host(id) => {
                if *id == TypeId::of::<P>() {
                    Ok(())
                } else {
                    bail!("future payload types differ")
                }
            }
        }
    }
}