Skip to main content

wasmtime/runtime/component/concurrent/
futures_and_streams.rs

1use super::table::{TableDebug, TableId};
2use super::{Event, GlobalErrorContextRefCount, Waitable, WaitableCommon};
3use crate::component::concurrent::{ConcurrentState, CurrentThread, WorkItem, tls};
4use crate::component::func::{self, LiftContext, LowerContext};
5use crate::component::matching::InstanceType;
6use crate::component::resources::{HostResourceIndex, HostResourceTables};
7use crate::component::types;
8use crate::component::values::ErrorContextAny;
9use crate::component::{
10    AsAccessor, ComponentInstanceId, ComponentType, FutureAny, Instance, Lift, Lower,
11    RuntimeInstance, StreamAny, Val, WasmList,
12};
13use crate::prelude::*;
14use crate::store::{StoreOpaque, StoreToken};
15use crate::try_mutex::{TryMutex, TryMutexGuard};
16use crate::vm::component::{ComponentInstance, HandleCounter, HandleTable, TransmitLocalState};
17use crate::vm::{AlwaysMut, VMStore};
18use crate::{AsContext, AsContextMut, StoreContextMut, ValRaw};
19use crate::{
20    Error, Result, Trap, bail, bail_bug, ensure,
21    error::{Context as _, format_err},
22};
23use alloc::sync::Arc;
24use buffers::{Extender, SliceBuffer, UntypedWriteBuffer};
25use core::any::{Any, TypeId};
26use core::fmt;
27use core::future;
28use core::iter;
29use core::marker::PhantomData;
30use core::mem::{self, ManuallyDrop, MaybeUninit};
31use core::ops::{Deref, DerefMut};
32use core::pin::Pin;
33use core::task::{Context, Poll, Waker, ready};
34use futures::channel::oneshot;
35use futures::{FutureExt as _, stream};
36use wasmtime_environ::component::{
37    CanonicalAbiInfo, ComponentTypes, InterfaceType, OptionsIndex, RuntimeComponentInstanceIndex,
38    TypeComponentGlobalErrorContextTableIndex, TypeComponentLocalErrorContextTableIndex,
39    TypeFutureTableIndex, TypeStreamTableIndex,
40};
41
42pub use buffers::{ReadBuffer, VecBuffer, WriteBuffer};
43
44mod buffers;
45
46/// Enum for distinguishing between a stream or future in functions that handle
47/// both.
48#[derive(Copy, Clone, Debug)]
49pub enum TransmitKind {
50    Stream,
51    Future,
52}
53
54/// Represents `{stream,future}.{read,write}` results.
55#[derive(Copy, Clone, Debug, PartialEq)]
56pub enum ReturnCode {
57    Blocked,
58    Completed(ItemCount),
59    Dropped(ItemCount),
60    Cancelled(ItemCount),
61}
62
63impl ReturnCode {
64    /// Pack `self` into a single 32-bit integer that may be returned to the
65    /// guest.
66    ///
67    /// This corresponds to `pack_copy_result` in the Component Model spec.
68    pub fn encode(&self) -> u32 {
69        const BLOCKED: u32 = 0xffff_ffff;
70        const COMPLETED: u32 = 0x0;
71        const DROPPED: u32 = 0x1;
72        const CANCELLED: u32 = 0x2;
73        match self {
74            ReturnCode::Blocked => BLOCKED,
75            ReturnCode::Completed(n) => (n.as_u32() << 4) | COMPLETED,
76            ReturnCode::Dropped(n) => (n.as_u32() << 4) | DROPPED,
77            ReturnCode::Cancelled(n) => (n.as_u32() << 4) | CANCELLED,
78        }
79    }
80
81    /// Returns `Self::Completed` with the specified count (or zero if
82    /// `matches!(kind, TransmitKind::Future)`)
83    fn completed(kind: TransmitKind, count: ItemCount) -> Self {
84        Self::Completed(if let TransmitKind::Future = kind {
85            ItemCount::ZERO
86        } else {
87            count
88        })
89    }
90}
91
92/// Representation of how many items are being operated on in a stream read or
93/// write.
94///
95/// The component model requires that stream operations are limited to `1<<28`
96/// items in one go. This type is a newtype wrapper around `u32` with the
97/// invariant that the internal value is limited by this amount.
98#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
99#[repr(transparent)]
100pub struct ItemCount {
101    raw: u32,
102}
103
104impl ItemCount {
105    const MAX: u32 = 1 << 28;
106    const ZERO: ItemCount = ItemCount { raw: 0 };
107
108    /// Creates a new `ItemCount` with the specified count, or a trap if it's
109    /// too large.
110    fn new(count: u32) -> Result<Self, Trap> {
111        if count < Self::MAX {
112            Ok(Self { raw: count })
113        } else {
114            Err(Trap::StreamOpTooBig)
115        }
116    }
117
118    /// Same as `Self::new` but takes a `usize`.
119    fn new_usize(count: usize) -> Result<Self, Trap> {
120        let count = u32::try_from(count).map_err(|_| Trap::StreamOpTooBig)?;
121        Self::new(count)
122    }
123
124    fn as_u32(&self) -> u32 {
125        self.raw
126    }
127
128    fn as_usize(&self) -> usize {
129        usize::try_from(self.raw).unwrap()
130    }
131
132    /// Increments `self` by `amt`, returning a trap if the amount would exceed
133    /// the maximum item count.
134    fn inc(&mut self, amt: usize) -> Result<(), Trap> {
135        let amt = u32::try_from(amt).map_err(|_| Trap::StreamOpTooBig)?;
136        let new_raw = self.raw.checked_add(amt).ok_or(Trap::StreamOpTooBig)?;
137        if new_raw < Self::MAX {
138            self.raw = new_raw;
139            Ok(())
140        } else {
141            Err(Trap::StreamOpTooBig)
142        }
143    }
144
145    /// Helper to add two `ItemCount`s together, fallibly.
146    ///
147    /// It's considered a bug if this overflows, so this is only suitable in
148    /// situations where overflow and/or exceeding the total item count is known
149    /// that it may be possible.
150    fn add(&self, other: ItemCount) -> Result<ItemCount> {
151        match self.raw.checked_add(other.raw) {
152            Some(raw) => Ok(ItemCount::new(raw)?),
153            None => bail_bug!("overflow in `ItemCount::add`"),
154        }
155    }
156
157    /// Same as `add`, but for subtraction.
158    ///
159    /// Like with `add` this is only suitable for situations where the result is
160    /// known to not underflow.
161    fn sub(&self, other: ItemCount) -> Result<ItemCount> {
162        match self.raw.checked_sub(other.raw) {
163            Some(raw) => Ok(ItemCount { raw }),
164            None => bail_bug!("underflow in `ItemCount::sub`"),
165        }
166    }
167}
168
169impl fmt::Display for ItemCount {
170    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
171        self.raw.fmt(f)
172    }
173}
174
175impl fmt::Debug for ItemCount {
176    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
177        self.raw.fmt(f)
178    }
179}
180
181impl PartialEq<u32> for ItemCount {
182    fn eq(&self, other: &u32) -> bool {
183        self.raw == *other
184    }
185}
186
187impl PartialOrd<u32> for ItemCount {
188    fn partial_cmp(&self, other: &u32) -> Option<core::cmp::Ordering> {
189        self.raw.partial_cmp(other)
190    }
191}
192
193/// Represents a stream or future type index.
194///
195/// This is useful as a parameter type for functions which operate on either a
196/// future or a stream.
197#[derive(Copy, Clone, Debug)]
198pub enum TransmitIndex {
199    Stream(TypeStreamTableIndex),
200    Future(TypeFutureTableIndex),
201}
202
203impl TransmitIndex {
204    pub fn kind(&self) -> TransmitKind {
205        match self {
206            TransmitIndex::Stream(_) => TransmitKind::Stream,
207            TransmitIndex::Future(_) => TransmitKind::Future,
208        }
209    }
210
211    /// Retrieve the payload type of the specified stream or future, or `None`
212    /// if it has no payload type.
213    fn payload<'a>(&self, types: &'a ComponentTypes) -> Option<&'a InterfaceType> {
214        match self {
215            TransmitIndex::Stream(i) => {
216                let ty = types[*i].ty;
217                types[ty].payload.as_ref()
218            }
219            TransmitIndex::Future(i) => {
220                let ty = types[*i].ty;
221                types[ty].payload.as_ref()
222            }
223        }
224    }
225}
226
227/// Retrieve the host rep and state for the specified guest-visible waitable
228/// handle.
229fn get_mut_by_index_from(
230    handle_table: &mut HandleTable,
231    ty: TransmitIndex,
232    index: u32,
233) -> Result<(u32, &mut TransmitLocalState)> {
234    match ty {
235        TransmitIndex::Stream(ty) => handle_table.stream_rep(ty, index),
236        TransmitIndex::Future(ty) => handle_table.future_rep(ty, index),
237    }
238}
239
240fn lower<T: func::Lower + Send + 'static, B: WriteBuffer<T>, U: 'static>(
241    mut store: StoreContextMut<U>,
242    instance: Instance,
243    options: OptionsIndex,
244    ty: TransmitIndex,
245    address: usize,
246    count: usize,
247    buffer: &mut B,
248) -> Result<()> {
249    let count = buffer.remaining().len().min(count);
250
251    let lower = &mut if T::MAY_REQUIRE_REALLOC {
252        LowerContext::new(store.as_context_mut(), options, instance)
253    } else {
254        LowerContext::new_without_realloc(store.as_context_mut(), options, instance)
255    };
256
257    // Note that validation only happens with operations that have a payload.
258    if let Some(ty) = ty.payload(lower.types).copied() {
259        if address % usize::try_from(T::ALIGN32)? != 0 {
260            bail!("read pointer not aligned");
261        }
262        lower
263            .as_slice_mut()
264            .get_mut(address..)
265            .and_then(|b| b.get_mut(..T::SIZE32 * count))
266            .ok_or_else(|| crate::format_err!("read pointer out of bounds of memory"))?;
267
268        T::linear_store_list_to_memory(lower, ty, address, &buffer.remaining()[..count])?;
269    }
270
271    buffer.skip(count);
272
273    Ok(())
274}
275
276fn lift<T: func::Lift + Send + 'static, B: ReadBuffer<T>>(
277    lift: &mut LiftContext<'_>,
278    ty: Option<InterfaceType>,
279    buffer: &mut B,
280    address: usize,
281    count: usize,
282) -> Result<()> {
283    let count = count.min(buffer.remaining_capacity());
284    if T::IS_RUST_UNIT_TYPE {
285        // SAFETY: `T::IS_RUST_UNIT_TYPE` is only true for `()`, a
286        // zero-sized type, so `MaybeUninit::uninit().assume_init()`
287        // is a valid way to populate the zero-sized buffer.
288        buffer.extend(
289            iter::repeat_with(|| unsafe { MaybeUninit::uninit().assume_init() }).take(count),
290        )
291    } else {
292        let ty = match ty {
293            Some(ty) => ty,
294            None => bail_bug!("type required for non-unit lift"),
295        };
296        if address % usize::try_from(T::ALIGN32)? != 0 {
297            bail!("write pointer not aligned");
298        }
299        lift.memory()
300            .get(address..)
301            .and_then(|b| b.get(..T::SIZE32 * count))
302            .ok_or_else(|| crate::format_err!("write pointer out of bounds of memory"))?;
303
304        let list = &WasmList::new(address, count, lift, ty)?;
305        T::linear_lift_into_from_memory(lift, list, &mut Extender(buffer))?
306    }
307    Ok(())
308}
309
310/// Represents the state associated with an error context
311#[derive(Debug, PartialEq, Eq, PartialOrd)]
312pub(super) struct ErrorContextState {
313    /// Debug message associated with the error context
314    pub(crate) debug_msg: String,
315}
316
317/// Represents the size and alignment for a "flat" Component Model type,
318/// i.e. one containing no pointers or handles.
319#[derive(Debug, Clone, Copy, PartialEq, Eq)]
320pub(super) struct FlatAbi {
321    pub(super) size: u32,
322    pub(super) align: u32,
323}
324
325struct HostBuffer<'a> {
326    dst: &'a mut Vec<u8>,
327    marked_written: &'a mut usize,
328}
329
330impl HostBuffer<'_> {
331    fn reborrow(&mut self) -> HostBuffer<'_> {
332        HostBuffer {
333            dst: &mut *self.dst,
334            marked_written: &mut *self.marked_written,
335        }
336    }
337}
338
339/// Represents the buffer for a host- or guest-initiated stream read.
340pub struct Destination<'a, T, B> {
341    id: TableId<TransmitState>,
342    buffer: &'a mut B,
343    host_buffer: Option<HostBuffer<'a>>,
344    _phantom: PhantomData<fn() -> T>,
345}
346
347impl<'a, T, B> Destination<'a, T, B> {
348    /// Reborrow `self` so it can be used again later.
349    pub fn reborrow(&mut self) -> Destination<'_, T, B> {
350        Destination {
351            id: self.id,
352            buffer: &mut *self.buffer,
353            host_buffer: self.host_buffer.as_mut().map(|b| b.reborrow()),
354            _phantom: PhantomData,
355        }
356    }
357
358    /// Take the buffer out of `self`, leaving a default-initialized one in its
359    /// place.
360    ///
361    /// This can be useful for reusing the previously-stored buffer's capacity
362    /// instead of allocating a fresh one.
363    pub fn take_buffer(&mut self) -> B
364    where
365        B: Default,
366    {
367        mem::take(self.buffer)
368    }
369
370    /// Store the specified buffer in `self`.
371    ///
372    /// Any items contained in the buffer will be delivered to the reader after
373    /// the `StreamProducer::poll_produce` call to which this `Destination` was
374    /// passed returns (unless overwritten by another call to `set_buffer`).
375    ///
376    /// If items are stored via this buffer _and_ written via a
377    /// `DirectDestination` view of `self`, then the items in the buffer will be
378    /// delivered after the ones written using `DirectDestination`.
379    pub fn set_buffer(&mut self, buffer: B) {
380        *self.buffer = buffer;
381    }
382
383    /// Return the remaining number of items the current read has capacity to
384    /// accept, if known.
385    ///
386    /// This will return `Some(_)` if the reader is a guest; it will return
387    /// `None` if the reader is the host.
388    ///
389    /// Note that this can return `Some(0)`. This means that the guest is
390    /// attempting to perform a zero-length read which typically means that it's
391    /// trying to wait for this stream to be ready-to-read but is not actually
392    /// ready to receive the items yet. The host in this case is allowed to
393    /// either block waiting for readiness or immediately complete the
394    /// operation. The guest is expected to handle both cases. Some more
395    /// discussion about this case can be found in the discussion of ["Stream
396    /// Readiness" in the component-model repo][docs].
397    ///
398    /// [docs]: https://github.com/WebAssembly/component-model/blob/main/design/mvp/Concurrency.md#stream-readiness
399    pub fn remaining(&self, mut store: impl AsContextMut) -> Result<Option<usize>> {
400        self.remaining_(store.as_context_mut().0)
401    }
402
403    fn remaining_(&self, store: &mut StoreOpaque) -> Result<Option<usize>> {
404        let transmit = store.concurrent_state_mut()?.get_mut(self.id)?;
405
406        if let &ReadState::GuestReady { count, .. } = &transmit.read {
407            let &WriteState::HostReady { guest_offset, .. } = &transmit.write else {
408                bail_bug!("expected WriteState::HostReady")
409            };
410
411            Ok(Some(count.as_usize() - guest_offset.as_usize()))
412        } else {
413            Ok(None)
414        }
415    }
416}
417
418impl<'a, B> Destination<'a, u8, B> {
419    /// Return a `DirectDestination` view of `self`.
420    ///
421    /// If the reader is a guest, this will provide direct access to the guest's
422    /// read buffer.  If the reader is a host, this will provide access to a
423    /// buffer which will be delivered to the host before any items stored using
424    /// `Destination::set_buffer`.
425    ///
426    /// `capacity` will only be used if the reader is a host, in which case it
427    /// will update the length of the buffer, possibly zero-initializing the new
428    /// elements if the new length is larger than the old length.
429    pub fn as_direct<D>(
430        mut self,
431        store: StoreContextMut<'a, D>,
432        capacity: usize,
433    ) -> DirectDestination<'a, D> {
434        if let Some(buffer) = &mut self.host_buffer {
435            *buffer.marked_written = 0;
436            buffer.dst.resize(capacity, 0);
437        }
438
439        DirectDestination {
440            id: self.id,
441            host_buffer: self.host_buffer,
442            store,
443        }
444    }
445}
446
447/// Represents a read from a `stream<u8>`, providing direct access to the
448/// writer's buffer.
449pub struct DirectDestination<'a, D: 'static> {
450    id: TableId<TransmitState>,
451    host_buffer: Option<HostBuffer<'a>>,
452    store: StoreContextMut<'a, D>,
453}
454
455#[cfg(feature = "std")]
456impl<D: 'static> std::io::Write for DirectDestination<'_, D> {
457    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
458        let rem = self.remaining().map_err(std::io::Error::other)?;
459        let n = rem.len().min(buf.len());
460        rem[..n].copy_from_slice(&buf[..n]);
461        self.mark_written(n);
462        Ok(n)
463    }
464
465    fn flush(&mut self) -> std::io::Result<()> {
466        Ok(())
467    }
468}
469
470impl<D: 'static> DirectDestination<'_, D> {
471    /// Provide direct access to the writer's buffer.
472    pub fn remaining(&mut self) -> Result<&mut [u8]> {
473        if let Some(buffer) = self.host_buffer.as_mut() {
474            return Ok(buffer.dst);
475        }
476        let transmit = self
477            .store
478            .as_context_mut()
479            .0
480            .concurrent_state_mut()?
481            .get_mut(self.id)?;
482
483        let &ReadState::GuestReady {
484            address,
485            count,
486            options,
487            instance,
488            ..
489        } = &transmit.read
490        else {
491            bail_bug!("expected ReadState::GuestReady")
492        };
493
494        let &WriteState::HostReady { guest_offset, .. } = &transmit.write else {
495            bail_bug!("expected WriteState::HostReady")
496        };
497
498        let memory = instance
499            .options_memory_mut(self.store.0, options)
500            .get_mut(
501                if count == 0 {
502                    0
503                } else {
504                    address + guest_offset.as_usize()
505                }..,
506            )
507            .and_then(|b| b.get_mut(..(count.as_usize() - guest_offset.as_usize())));
508        match memory {
509            Some(memory) => Ok(memory),
510            None => bail_bug!("guest buffer unexpectedly out of bounds"),
511        }
512    }
513
514    /// Mark the specified number of bytes as written to the writer's buffer.
515    ///
516    /// # Panics
517    ///
518    /// This will panic if the count is larger than the size of the
519    /// buffer returned by `Self::remaining`.
520    pub fn mark_written(&mut self, count: usize) {
521        // Note that this unwrap should only trigger for bugs in Wasmtime, and
522        // this is modeled here to centralize the `.unwrap()` for this method in
523        // one location.
524        self.mark_written_(count).unwrap()
525    }
526
527    fn mark_written_(&mut self, count: usize) -> Result<()> {
528        if let Some(buffer) = self.host_buffer.as_mut() {
529            // Note that this `.unwrap` is a documented panic condition of
530            // `mark_written`.
531            *buffer.marked_written = buffer.marked_written.checked_add(count).unwrap();
532        } else {
533            let transmit = self
534                .store
535                .as_context_mut()
536                .0
537                .concurrent_state_mut()?
538                .get_mut(self.id)?;
539
540            let ReadState::GuestReady {
541                count: read_count, ..
542            } = &transmit.read
543            else {
544                bail_bug!("expected ReadState::GuestReady")
545            };
546
547            let WriteState::HostReady { guest_offset, .. } = &mut transmit.write else {
548                bail_bug!("expected WriteState::HostReady");
549            };
550
551            if guest_offset.as_usize() + count > read_count.as_usize() {
552                // Note that this `panic` is a documented panic condition of
553                // `mark_written`.
554                panic!(
555                    "write count ({count}) must be less than or equal to read count ({read_count})"
556                )
557            } else {
558                guest_offset.inc(count)?;
559            }
560        }
561        Ok(())
562    }
563}
564
565/// Represents the state of a `Stream{Producer,Consumer}`.
566#[derive(Copy, Clone, Debug)]
567pub enum StreamResult {
568    /// The operation completed normally, and the producer or consumer may be
569    /// able to produce or consume more items, respectively.
570    Completed,
571    /// The operation was interrupted (i.e. it wrapped up early after receiving
572    /// a `finish` parameter value of true in a call to `poll_produce` or
573    /// `poll_consume`), and the producer or consumer may be able to produce or
574    /// consume more items, respectively.
575    Cancelled,
576    /// The operation completed normally, but the producer or consumer will
577    /// _not_ able to produce or consume more items, respectively.
578    Dropped,
579}
580
581/// Represents the host-owned write end of a stream.
582pub trait StreamProducer<D>: Send + 'static {
583    /// The payload type of this stream.
584    type Item;
585
586    /// The `WriteBuffer` type to use when delivering items.
587    type Buffer: WriteBuffer<Self::Item> + Default;
588
589    /// Handle a host- or guest-initiated read by delivering zero or more items
590    /// to the specified destination.
591    ///
592    /// This will be called whenever the reader starts a read.
593    ///
594    /// # Arguments
595    ///
596    /// * `self` - a `Pin`'d version of self to perform Rust-level
597    ///   future-related operations on.
598    /// * `cx` - a Rust-related [`Context`] which is passed to other
599    ///   future-related operations or used to acquire a waker.
600    /// * `store` - the Wasmtime store that this operation is happening within.
601    ///   Used, for example, to consult the state `D` associated with the store.
602    /// * `destination` - the location that items are to be written to.
603    /// * `finish` - a flag indicating whether the host should strive to
604    ///   immediately complete/cancel any pending operation. See below for more
605    ///   details.
606    ///
607    /// # Behavior
608    ///
609    /// If the implementation is able to produce one or more items immediately,
610    /// it should write them to `destination` and return either
611    /// `Poll::Ready(Ok(StreamResult::Completed))` if it expects to produce more
612    /// items, or `Poll::Ready(Ok(StreamResult::Dropped))` if it cannot produce
613    /// any more items.
614    ///
615    /// If the implementation is unable to produce any items immediately, but
616    /// expects to do so later, and `finish` is _false_, it should store the
617    /// waker from `cx` for later and return `Poll::Pending` without writing
618    /// anything to `destination`.  Later, it should alert the waker when either
619    /// the items arrive, the stream has ended, or an error occurs.
620    ///
621    /// If more items are written to `destination` than the reader has immediate
622    /// capacity to accept, they will be retained in memory by the caller and
623    /// used to satisfy future reads, in which case `poll_produce` will only be
624    /// called again once all those items have been delivered.
625    ///
626    /// # Zero-length reads
627    ///
628    /// This function may be called with a zero-length capacity buffer
629    /// (i.e. `Destination::remaining` returns `Some(0)`). This indicates that
630    /// the guest wants to wait to see if an item is ready without actually
631    /// reading the item. For example think of a UNIX `poll` function run on a
632    /// TCP stream, seeing if it's readable without actually reading it.
633    ///
634    /// In this situation the host is allowed to either return immediately or
635    /// wait for readiness. Note that waiting for readiness is not always
636    /// possible. For example it's impossible to test if a Rust-native `Future`
637    /// is ready without actually reading the item. Stream-specific
638    /// optimizations, such as testing if a TCP stream is readable, may be
639    /// possible however.
640    ///
641    /// For a zero-length read, the host is allowed to:
642    ///
643    /// - Return `Poll::Ready(Ok(StreamResult::Completed))` without writing
644    ///   anything if it expects to be able to produce items immediately (i.e.
645    ///   without first returning `Poll::Pending`) the next time `poll_produce`
646    ///   is called with non-zero capacity. This is the best-case scenario of
647    ///   fulfilling the guest's desire -- items aren't read/buffered but the
648    ///   host is saying it's ready when the guest is.
649    ///
650    /// - Return `Poll::Ready(Ok(StreamResult::Completed))` without actually
651    ///   testing for readiness. The guest doesn't know this yet, but the guest
652    ///   will realize that zero-length reads won't work on this stream when a
653    ///   subsequent nonzero read attempt is made which returns `Poll::Pending`
654    ///   here.
655    ///
656    /// - Return `Poll::Pending` if the host has performed necessary async work
657    ///   to wait for this stream to be readable without actually reading
658    ///   anything. This is also a best-case scenario where the host is letting
659    ///   the guest know that nothing is ready yet. Later the zero-length read
660    ///   will complete and then the guest will attempt a nonzero-length read to
661    ///   actually read some bytes.
662    ///
663    /// - Return `Poll::Ready(Ok(StreamResult::Completed))` after calling
664    ///   `Destination::set_buffer` with one more more items. Note, however,
665    ///   that this creates the hazard that the items will never be received by
666    ///   the guest if it decides not to do another non-zero-length read before
667    ///   closing the stream.  Moreover, if `Self::Item` is e.g. a
668    ///   `Resource<_>`, they may end up leaking in that scenario. It is not
669    ///   recommended to do this and it's better to return
670    ///   `StreamResult::Completed` without buffering anything instead.
671    ///
672    /// For more discussion on zero-length reads see the [documentation in the
673    /// component-model repo itself][docs].
674    ///
675    /// [docs]: https://github.com/WebAssembly/component-model/blob/main/design/mvp/Concurrency.md#stream-readiness
676    ///
677    /// # Return
678    ///
679    /// This function can return a number of possible cases from this function:
680    ///
681    /// * `Poll::Pending` - this operation cannot complete at this time. The
682    ///   Rust-level `Future::poll` contract applies here where a waker should
683    ///   be stored from the `cx` argument and be arranged to receive a
684    ///   notification when this implementation can make progress. For example
685    ///   if you call `Future::poll` on a sub-future, that's enough. If items
686    ///   were written to `destination` then a trap in the guest will be raised.
687    ///
688    ///   Note that implementations should strive to avoid this return value
689    ///   when `finish` is `true`. In such a situation the guest is attempting
690    ///   to, for example, cancel a previous operation. By returning
691    ///   `Poll::Pending` the guest will be blocked during the cancellation
692    ///   request. If `finish` is `true` then `StreamResult::Cancelled` is
693    ///   favored to indicate that no items were read. If a short read happened,
694    ///   however, it's ok to return `StreamResult::Completed` indicating some
695    ///   items were read.
696    ///
697    /// * `Poll::Ok(StreamResult::Completed)` - items, if applicable, were
698    ///   written to the `destination`. Only zero-length reads (see above) may
699    ///   complete without writing anything; if nothing was written for a
700    ///   nonzero-length read, or for any read by the host, then a trap in the
701    ///   guest will be raised.
702    ///
703    /// * `Poll::Ok(StreamResult::Cancelled)` - used when `finish` is `true` and
704    ///   the implementation was able to successfully cancel any async work that
705    ///   a previous read kicked off, if any. The host should not buffer values
706    ///   received after returning `Cancelled` because the guest will not be
707    ///   aware of these values and the guest could close the stream after
708    ///   cancelling a read. Hosts should only return `Cancelled` when there are
709    ///   no more async operations in flight for a previous read.
710    ///
711    ///   If items were written to `destination` then a trap in the guest will
712    ///   be raised. If `finish` is `false` then this return value will raise a
713    ///   trap in the guest.
714    ///
715    /// * `Poll::Ok(StreamResult::Dropped)` - end-of-stream marker, indicating
716    ///   that this producer should not be polled again. Note that items may
717    ///   still be written to `destination`.
718    ///
719    /// # Errors
720    ///
721    /// The implementation may alternatively choose to return `Err(_)` to
722    /// indicate an unrecoverable error. This will cause the guest (if any) to
723    /// trap and render the component instance (if any) unusable. The
724    /// implementation should report errors that _are_ recoverable by other
725    /// means (e.g. by writing to a `future`) and return
726    /// `Poll::Ready(Ok(StreamResult::Dropped))`.
727    fn poll_produce<'a>(
728        self: Pin<&mut Self>,
729        cx: &mut Context<'_>,
730        store: StoreContextMut<'a, D>,
731        destination: Destination<'a, Self::Item, Self::Buffer>,
732        finish: bool,
733    ) -> Poll<Result<StreamResult>>;
734
735    /// Attempt to convert the specified object into a `Box<dyn Any>` which may
736    /// be downcast to the specified type.
737    ///
738    /// The implementation must ensure that, if it returns `Ok(_)`, a downcast
739    /// to the specified type is guaranteed to succeed.
740    fn try_into(me: Pin<Box<Self>>, _ty: TypeId) -> Result<Box<dyn Any>, Pin<Box<Self>>> {
741        Err(me)
742    }
743}
744
745impl<T, D> StreamProducer<D> for iter::Empty<T>
746where
747    T: Send + Sync + 'static,
748{
749    type Item = T;
750    type Buffer = Option<Self::Item>;
751
752    fn poll_produce<'a>(
753        self: Pin<&mut Self>,
754        _: &mut Context<'_>,
755        _: StoreContextMut<'a, D>,
756        _: Destination<'a, Self::Item, Self::Buffer>,
757        _: bool,
758    ) -> Poll<Result<StreamResult>> {
759        Poll::Ready(Ok(StreamResult::Dropped))
760    }
761}
762
763impl<T, D> StreamProducer<D> for stream::Empty<T>
764where
765    T: Send + Sync + 'static,
766{
767    type Item = T;
768    type Buffer = Option<Self::Item>;
769
770    fn poll_produce<'a>(
771        self: Pin<&mut Self>,
772        _: &mut Context<'_>,
773        _: StoreContextMut<'a, D>,
774        _: Destination<'a, Self::Item, Self::Buffer>,
775        _: bool,
776    ) -> Poll<Result<StreamResult>> {
777        Poll::Ready(Ok(StreamResult::Dropped))
778    }
779}
780
781impl<T, D> StreamProducer<D> for Vec<T>
782where
783    T: Unpin + Send + Sync + 'static,
784{
785    type Item = T;
786    type Buffer = VecBuffer<T>;
787
788    fn poll_produce<'a>(
789        self: Pin<&mut Self>,
790        _: &mut Context<'_>,
791        _: StoreContextMut<'a, D>,
792        mut dst: Destination<'a, Self::Item, Self::Buffer>,
793        _: bool,
794    ) -> Poll<Result<StreamResult>> {
795        dst.set_buffer(mem::take(self.get_mut()).into());
796        Poll::Ready(Ok(StreamResult::Dropped))
797    }
798}
799
800impl<T, D> StreamProducer<D> for Box<[T]>
801where
802    T: Unpin + Send + Sync + 'static,
803{
804    type Item = T;
805    type Buffer = VecBuffer<T>;
806
807    fn poll_produce<'a>(
808        self: Pin<&mut Self>,
809        _: &mut Context<'_>,
810        _: StoreContextMut<'a, D>,
811        mut dst: Destination<'a, Self::Item, Self::Buffer>,
812        _: bool,
813    ) -> Poll<Result<StreamResult>> {
814        dst.set_buffer(mem::take(self.get_mut()).into_vec().into());
815        Poll::Ready(Ok(StreamResult::Dropped))
816    }
817}
818
819#[cfg(feature = "component-model-bytes")]
820impl<D> StreamProducer<D> for bytes::Bytes {
821    type Item = u8;
822    type Buffer = Self;
823
824    fn poll_produce<'a>(
825        self: Pin<&mut Self>,
826        _: &mut Context<'_>,
827        _store: StoreContextMut<'a, D>,
828        mut dst: Destination<'a, Self::Item, Self::Buffer>,
829        _: bool,
830    ) -> Poll<Result<StreamResult>> {
831        dst.set_buffer(mem::take(self.get_mut()));
832        Poll::Ready(Ok(StreamResult::Dropped))
833    }
834}
835
836#[cfg(feature = "component-model-bytes")]
837impl<D> StreamProducer<D> for bytes::BytesMut {
838    type Item = u8;
839    type Buffer = Self;
840
841    fn poll_produce<'a>(
842        self: Pin<&mut Self>,
843        _: &mut Context<'_>,
844        _store: StoreContextMut<'a, D>,
845        mut dst: Destination<'a, Self::Item, Self::Buffer>,
846        _: bool,
847    ) -> Poll<Result<StreamResult>> {
848        dst.set_buffer(mem::take(self.get_mut()));
849        Poll::Ready(Ok(StreamResult::Dropped))
850    }
851}
852
853/// Represents the buffer for a host- or guest-initiated stream write.
854pub struct Source<'a, T> {
855    id: TableId<TransmitState>,
856    host_buffer: Option<&'a mut dyn WriteBuffer<T>>,
857}
858
859impl<'a, T> Source<'a, T> {
860    /// Reborrow `self` so it can be used again later.
861    pub fn reborrow(&mut self) -> Source<'_, T> {
862        Source {
863            id: self.id,
864            host_buffer: self.host_buffer.as_deref_mut(),
865        }
866    }
867
868    /// Accept zero or more items from the writer.
869    pub fn read<B, S: AsContextMut>(&mut self, mut store: S, buffer: &mut B) -> Result<()>
870    where
871        T: func::Lift + 'static,
872        B: ReadBuffer<T>,
873    {
874        if let Some(input) = &mut self.host_buffer {
875            let count = input.remaining().len().min(buffer.remaining_capacity());
876            buffer.move_from(*input, count);
877        } else {
878            let store = store.as_context_mut();
879            let transmit = store.0.concurrent_state_mut()?.get_mut(self.id)?;
880
881            let &ReadState::HostReady { guest_offset, .. } = &transmit.read else {
882                bail_bug!("expected ReadState::HostReady");
883            };
884
885            let &WriteState::GuestReady {
886                ty,
887                address,
888                count,
889                options,
890                instance,
891                ..
892            } = &transmit.write
893            else {
894                bail_bug!("expected WriteState::GuestReady");
895            };
896
897            // Zero-length writes are exempt from bounds/alignment checks, so
898            // there's nothing to lift.
899            if guest_offset >= count {
900                return Ok(());
901            }
902
903            let cx = &mut LiftContext::new(store.0.store_opaque_mut(), options, instance)?;
904            let ty = ty.payload(cx.types);
905            let old_remaining = buffer.remaining_capacity();
906            lift::<T, B>(
907                cx,
908                ty.copied(),
909                buffer,
910                address + (T::SIZE32 * guest_offset.as_usize()),
911                count.as_usize() - guest_offset.as_usize(),
912            )?;
913
914            let transmit = store.0.concurrent_state_mut()?.get_mut(self.id)?;
915
916            let ReadState::HostReady { guest_offset, .. } = &mut transmit.read else {
917                bail_bug!("expected ReadState::HostReady");
918            };
919
920            guest_offset.inc(old_remaining - buffer.remaining_capacity())?;
921        }
922
923        Ok(())
924    }
925
926    /// Return the number of items remaining to be read from the current write
927    /// operation.
928    pub fn remaining(&self, mut store: impl AsContextMut) -> Result<usize>
929    where
930        T: 'static,
931    {
932        self.remaining_(store.as_context_mut().0)
933    }
934
935    fn remaining_(&self, store: &mut StoreOpaque) -> Result<usize>
936    where
937        T: 'static,
938    {
939        let transmit = store.concurrent_state_mut()?.get_mut(self.id)?;
940
941        if let &WriteState::GuestReady { count, .. } = &transmit.write {
942            let &ReadState::HostReady { guest_offset, .. } = &transmit.read else {
943                bail_bug!("expected ReadState::HostReady")
944            };
945
946            Ok(count.as_usize() - guest_offset.as_usize())
947        } else if let Some(host_buffer) = &self.host_buffer {
948            Ok(host_buffer.remaining().len())
949        } else {
950            bail_bug!("expected either WriteState::GuestReady or host buffer")
951        }
952    }
953}
954
955impl<'a> Source<'a, u8> {
956    /// Return a `DirectSource` view of `self`.
957    pub fn as_direct<D>(self, store: StoreContextMut<'a, D>) -> DirectSource<'a, D> {
958        DirectSource {
959            id: self.id,
960            host_buffer: self.host_buffer,
961            store,
962        }
963    }
964}
965
966/// Represents a write to a `stream<u8>`, providing direct access to the
967/// writer's buffer.
968pub struct DirectSource<'a, D: 'static> {
969    id: TableId<TransmitState>,
970    host_buffer: Option<&'a mut dyn WriteBuffer<u8>>,
971    store: StoreContextMut<'a, D>,
972}
973
974#[cfg(feature = "std")]
975impl<D: 'static> std::io::Read for DirectSource<'_, D> {
976    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
977        let rem = self.remaining().map_err(std::io::Error::other)?;
978        let n = rem.len().min(buf.len());
979        buf[..n].copy_from_slice(&rem[..n]);
980        self.mark_read(n);
981        Ok(n)
982    }
983}
984
985impl<D: 'static> DirectSource<'_, D> {
986    /// Provide direct access to the writer's buffer.
987    pub fn remaining(&mut self) -> Result<&[u8]> {
988        if let Some(buffer) = self.host_buffer.as_deref_mut() {
989            return Ok(buffer.remaining());
990        }
991        let transmit = self
992            .store
993            .as_context_mut()
994            .0
995            .concurrent_state_mut()?
996            .get_mut(self.id)?;
997
998        let &WriteState::GuestReady {
999            address,
1000            count,
1001            options,
1002            instance,
1003            ..
1004        } = &transmit.write
1005        else {
1006            bail_bug!("expected WriteState::GuestReady")
1007        };
1008
1009        let &ReadState::HostReady { guest_offset, .. } = &transmit.read else {
1010            bail_bug!("expected ReadState::HostReady")
1011        };
1012
1013        let memory = instance
1014            .options_memory(self.store.0, options)
1015            .get(
1016                if count == 0 {
1017                    0
1018                } else {
1019                    address + guest_offset.as_usize()
1020                }..,
1021            )
1022            .and_then(|b| b.get(..(count.as_usize() - guest_offset.as_usize())));
1023        match memory {
1024            Some(memory) => Ok(memory),
1025            None => bail_bug!("guest buffer unexpectedly out of bounds"),
1026        }
1027    }
1028
1029    /// Mark the specified number of bytes as read from the writer's buffer.
1030    ///
1031    /// # Panics
1032    ///
1033    /// This will panic if the count is larger than the size of the buffer
1034    /// returned by `Self::remaining`.
1035    pub fn mark_read(&mut self, count: usize) {
1036        // Note that this unwrap should only trigger for bugs in Wasmtime, and
1037        // this is modeled here to centralize the `.unwrap()` for this method in
1038        // one location.
1039        self.mark_read_(count).unwrap()
1040    }
1041
1042    fn mark_read_(&mut self, count: usize) -> Result<()> {
1043        if let Some(buffer) = self.host_buffer.as_deref_mut() {
1044            buffer.skip(count);
1045            return Ok(());
1046        }
1047
1048        let transmit = self
1049            .store
1050            .as_context_mut()
1051            .0
1052            .concurrent_state_mut()?
1053            .get_mut(self.id)?;
1054
1055        let WriteState::GuestReady {
1056            count: write_count, ..
1057        } = &transmit.write
1058        else {
1059            bail_bug!("expected WriteState::GuestReady");
1060        };
1061
1062        let ReadState::HostReady { guest_offset, .. } = &mut transmit.read else {
1063            bail_bug!("expected ReadState::HostReady");
1064        };
1065
1066        if guest_offset.as_usize() + count > write_count.as_usize() {
1067            // Note that this is a documented panic condition of `mark_read`.
1068            panic!("read count ({count}) must be less than or equal to write count ({write_count})")
1069        } else {
1070            guest_offset.inc(count)?;
1071        }
1072        Ok(())
1073    }
1074}
1075
1076/// Represents the host-owned read end of a stream.
1077pub trait StreamConsumer<D>: Send + 'static {
1078    /// The payload type of this stream.
1079    type Item;
1080
1081    /// Handle a host- or guest-initiated write by accepting zero or more items
1082    /// from the specified source.
1083    ///
1084    /// This will be called whenever the writer starts a write.
1085    ///
1086    /// If the implementation is able to consume one or more items immediately,
1087    /// it should take them from `source` and return either
1088    /// `Poll::Ready(Ok(StreamResult::Completed))` if it expects to be able to consume
1089    /// more items, or `Poll::Ready(Ok(StreamResult::Dropped))` if it cannot
1090    /// accept any more items.  Alternatively, it may return `Poll::Pending` to
1091    /// indicate that the caller should delay sending a `COMPLETED` event to the
1092    /// writer until a later call to this function returns `Poll::Ready(_)`.
1093    /// For more about that, see the `Backpressure` section below.
1094    ///
1095    /// If the implementation cannot consume any items immediately and `finish`
1096    /// is _false_, it should store the waker from `cx` for later and return
1097    /// `Poll::Pending` without writing anything to `destination`.  Later, it
1098    /// should alert the waker when either (1) the items arrive, (2) the stream
1099    /// has ended, or (3) an error occurs.
1100    ///
1101    /// If the implementation cannot consume any items immediately and `finish`
1102    /// is _true_, it should, if possible, return
1103    /// `Poll::Ready(Ok(StreamResult::Cancelled))` immediately without taking
1104    /// anything from `source`.  However, that might not be possible if an
1105    /// earlier call to `poll_consume` kicked off an asynchronous operation
1106    /// which needs to be completed (and possibly interrupted) gracefully, in
1107    /// which case the implementation may return `Poll::Pending` and later alert
1108    /// the waker as described above.  In other words, when `finish` is true,
1109    /// the implementation should prioritize returning a result to the reader
1110    /// (even if no items can be consumed) rather than wait indefinitely for at
1111    /// capacity to free up.
1112    ///
1113    /// In all of the above cases, the implementation may alternatively choose
1114    /// to return `Err(_)` to indicate an unrecoverable error.  This will cause
1115    /// the guest (if any) to trap and render the component instance (if any)
1116    /// unusable.  The implementation should report errors that _are_
1117    /// recoverable by other means (e.g. by writing to a `future`) and return
1118    /// `Poll::Ready(Ok(StreamResult::Dropped))`.
1119    ///
1120    /// Note that the implementation should only return
1121    /// `Poll::Ready(Ok(StreamResult::Cancelled))` without having taken any
1122    /// items from `source` if called with `finish` set to true.  If it does so
1123    /// when `finish` is false, the caller will trap.  Additionally, it should
1124    /// only return `Poll::Ready(Ok(StreamResult::Completed))` after taking at
1125    /// least one item from `source` if there is an item available; otherwise,
1126    /// the caller will trap.  Items taken during earlier calls for the same
1127    /// write which returned `Poll::Pending` count towards this, so it's fine
1128    /// to return `Completed` without taking more items once those have been
1129    /// forwarded (see "Backpressure" below).  If `poll_consume` is called with
1130    /// no items in
1131    /// `source`, it should only return `Poll::Ready(_)` once it is able to
1132    /// accept at least one item during the next call to `poll_consume`.
1133    ///
1134    /// Note that any items which the implementation of this trait takes from
1135    /// `source` become the responsibility of that implementation.  For that
1136    /// reason, an implementation which forwards items to an upstream sink
1137    /// should reserve capacity in that sink before taking items out of
1138    /// `source`, if possible.  Alternatively, it might buffer items which can't
1139    /// be forwarded immediately and send them once capacity is freed up.
1140    ///
1141    /// ## Backpressure
1142    ///
1143    /// As mentioned above, an implementation might choose to return
1144    /// `Poll::Pending` after taking items from `source`, which tells the caller
1145    /// to delay sending a `COMPLETED` event to the writer.  This can be used as
1146    /// a form of backpressure when the items are forwarded to an upstream sink
1147    /// asynchronously.  Note, however, that it's not possible to "put back"
1148    /// items into `source` once they've been taken out, so if the upstream sink
1149    /// is unable to accept all the items, that cannot be communicated to the
1150    /// writer at this level of abstraction.  Just as with application-specific,
1151    /// recoverable errors, information about which items could be forwarded and
1152    /// which could not must be communicated out-of-band, e.g. by writing to an
1153    /// application-specific `future`.
1154    ///
1155    /// Similarly, if the writer cancels the write after items have been taken
1156    /// from `source` but before the items have all been forwarded to an
1157    /// upstream sink, `poll_consume` will be called with `finish` set to true,
1158    /// and the implementation may either:
1159    ///
1160    /// - Interrupt the forwarding process gracefully.  This may be preferable
1161    /// if there is an out-of-band channel for communicating to the writer how
1162    /// many items were forwarded before being interrupted.
1163    ///
1164    /// - Allow the forwarding to complete without interrupting it.  This is
1165    /// usually preferable if there's no out-of-band channel for reporting back
1166    /// to the writer how many items were forwarded.
1167    fn poll_consume(
1168        self: Pin<&mut Self>,
1169        cx: &mut Context<'_>,
1170        store: StoreContextMut<D>,
1171        source: Source<'_, Self::Item>,
1172        finish: bool,
1173    ) -> Poll<Result<StreamResult>>;
1174}
1175
1176/// Represents a host-owned write end of a future.
1177pub trait FutureProducer<D>: Send + 'static {
1178    /// The payload type of this future.
1179    type Item;
1180
1181    /// Handle a host- or guest-initiated read by producing a value.
1182    ///
1183    /// This is equivalent to `StreamProducer::poll_produce`, but with a
1184    /// simplified interface for futures.
1185    ///
1186    /// If `finish` is true, the implementation may return
1187    /// `Poll::Ready(Ok(None))` to indicate the operation was canceled before it
1188    /// could produce a value.  Otherwise, it must either return
1189    /// `Poll::Ready(Ok(Some(_)))`, `Poll::Ready(Err(_))`, or `Poll::Pending`.
1190    fn poll_produce(
1191        self: Pin<&mut Self>,
1192        cx: &mut Context<'_>,
1193        store: StoreContextMut<D>,
1194        finish: bool,
1195    ) -> Poll<Result<Option<Self::Item>>>;
1196}
1197
1198impl<T, E, D, Fut> FutureProducer<D> for Fut
1199where
1200    E: Into<Error>,
1201    Fut: Future<Output = Result<T, E>> + ?Sized + Send + 'static,
1202{
1203    type Item = T;
1204
1205    fn poll_produce<'a>(
1206        self: Pin<&mut Self>,
1207        cx: &mut Context<'_>,
1208        _: StoreContextMut<'a, D>,
1209        finish: bool,
1210    ) -> Poll<Result<Option<T>>> {
1211        match self.poll(cx) {
1212            Poll::Ready(Ok(v)) => Poll::Ready(Ok(Some(v))),
1213            Poll::Ready(Err(err)) => Poll::Ready(Err(err.into())),
1214            Poll::Pending if finish => Poll::Ready(Ok(None)),
1215            Poll::Pending => Poll::Pending,
1216        }
1217    }
1218}
1219
1220/// Represents a host-owned read end of a future.
1221pub trait FutureConsumer<D>: Send + 'static {
1222    /// The payload type of this future.
1223    type Item;
1224
1225    /// Handle a host- or guest-initiated write by consuming a value.
1226    ///
1227    /// This is equivalent to `StreamProducer::poll_produce`, but with a
1228    /// simplified interface for futures.
1229    ///
1230    /// If `finish` is true, the implementation may return `Poll::Ready(Ok(()))`
1231    /// without taking the item from `source`, which indicates the operation was
1232    /// canceled before it could consume the value.  Otherwise, it must either
1233    /// take the item from `source` and return `Poll::Ready(Ok(()))`, or else
1234    /// return `Poll::Ready(Err(_))` or `Poll::Pending` (with or without taking
1235    /// the item).
1236    fn poll_consume(
1237        self: Pin<&mut Self>,
1238        cx: &mut Context<'_>,
1239        store: StoreContextMut<D>,
1240        source: Source<'_, Self::Item>,
1241        finish: bool,
1242    ) -> Poll<Result<()>>;
1243}
1244
1245/// Represents the readable end of a Component Model `future`.
1246///
1247/// Note that `FutureReader` instances must be disposed of using either `pipe`
1248/// or `close`; otherwise the in-store representation will leak and the writer
1249/// end will hang indefinitely.  Consider using [`GuardedFutureReader`] to
1250/// ensure that disposal happens automatically.
1251pub struct FutureReader<T> {
1252    idx: HostResourceIndex,
1253    _phantom: PhantomData<T>,
1254}
1255
1256impl<T> FutureReader<T> {
1257    /// Create a new future with the specified producer.
1258    ///
1259    /// # Errors
1260    ///
1261    /// Returns an error if the resource table for this store is full or if
1262    /// [`Config::concurrency_support`] is not enabled.
1263    ///
1264    /// [`Config::concurrency_support`]: crate::Config::concurrency_support
1265    pub fn new<S: AsContextMut>(
1266        mut store: S,
1267        producer: impl FutureProducer<S::Data, Item = T>,
1268    ) -> Result<Self>
1269    where
1270        T: func::Lower + func::Lift + Send + Sync + 'static,
1271    {
1272        ensure!(
1273            store.as_context().0.concurrency_support(),
1274            "concurrency support is not enabled"
1275        );
1276
1277        struct Producer<P>(P);
1278
1279        impl<D, T: func::Lower + 'static, P: FutureProducer<D, Item = T>> StreamProducer<D>
1280            for Producer<P>
1281        {
1282            type Item = P::Item;
1283            type Buffer = Option<P::Item>;
1284
1285            fn poll_produce<'a>(
1286                self: Pin<&mut Self>,
1287                cx: &mut Context<'_>,
1288                store: StoreContextMut<D>,
1289                mut destination: Destination<'a, Self::Item, Self::Buffer>,
1290                finish: bool,
1291            ) -> Poll<Result<StreamResult>> {
1292                // SAFETY: This is a standard pin-projection, and we never move
1293                // out of `self`.
1294                let producer = unsafe { self.map_unchecked_mut(|v| &mut v.0) };
1295
1296                Poll::Ready(Ok(
1297                    if let Some(value) = ready!(producer.poll_produce(cx, store, finish))? {
1298                        destination.set_buffer(Some(value));
1299
1300                        // Here we return `StreamResult::Completed` even though
1301                        // we've produced the last item we'll ever produce.
1302                        // That's because the ABI expects
1303                        // `ReturnCode::Completed(1)` rather than
1304                        // `ReturnCode::Dropped(1)`.  In any case, we won't be
1305                        // called again since the future will have resolved.
1306                        StreamResult::Completed
1307                    } else {
1308                        StreamResult::Cancelled
1309                    },
1310                ))
1311            }
1312        }
1313
1314        Ok(Self::new_(
1315            store
1316                .as_context_mut()
1317                .new_transmit(TransmitKind::Future, Producer(producer))?,
1318        ))
1319    }
1320
1321    pub(super) fn new_(idx: HostResourceIndex) -> Self {
1322        Self {
1323            idx,
1324            _phantom: PhantomData,
1325        }
1326    }
1327
1328    pub(super) fn idx(&self) -> HostResourceIndex {
1329        self.idx
1330    }
1331
1332    /// Set the consumer that accepts the result of this future.
1333    ///
1334    /// # Errors
1335    ///
1336    /// Returns an error if this future has already been closed or otherwise
1337    /// transferred (e.g. lowered into a guest), or if it does not belong to
1338    /// `store`.
1339    pub fn pipe<S: AsContextMut>(
1340        self,
1341        mut store: S,
1342        consumer: impl FutureConsumer<S::Data, Item = T> + Unpin,
1343    ) -> Result<()>
1344    where
1345        T: func::Lift + 'static,
1346    {
1347        struct Consumer<C>(C);
1348
1349        impl<D: 'static, T: func::Lift + 'static, C: FutureConsumer<D, Item = T>> StreamConsumer<D>
1350            for Consumer<C>
1351        {
1352            type Item = T;
1353
1354            fn poll_consume(
1355                self: Pin<&mut Self>,
1356                cx: &mut Context<'_>,
1357                mut store: StoreContextMut<D>,
1358                mut source: Source<Self::Item>,
1359                finish: bool,
1360            ) -> Poll<Result<StreamResult>> {
1361                // SAFETY: This is a standard pin-projection, and we never move
1362                // out of `self`.
1363                let consumer = unsafe { self.map_unchecked_mut(|v| &mut v.0) };
1364
1365                ready!(consumer.poll_consume(
1366                    cx,
1367                    store.as_context_mut(),
1368                    source.reborrow(),
1369                    finish
1370                ))?;
1371
1372                Poll::Ready(Ok(if source.remaining(store)? == 0 {
1373                    // Here we return `StreamResult::Completed` even though
1374                    // we've consumed the last item we'll ever consume.  That's
1375                    // because the ABI expects `ReturnCode::Completed(1)` rather
1376                    // than `ReturnCode::Dropped(1)`.  In any case, we won't be
1377                    // called again since the future will have resolved.
1378                    StreamResult::Completed
1379                } else {
1380                    StreamResult::Cancelled
1381                }))
1382            }
1383        }
1384
1385        store
1386            .as_context_mut()
1387            .set_consumer(self.idx, TransmitKind::Future, Consumer(consumer))
1388    }
1389
1390    /// Transfer ownership of the read end of a future from a guest to the host.
1391    fn lift_from_index(cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32) -> Result<Self> {
1392        Ok(Self::new_(lift_index_to_future(cx, ty, index)?))
1393    }
1394
1395    /// Close this `FutureReader`.
1396    ///
1397    /// This will close this half of the future which will signal to a pending
1398    /// write, if any, that the reader side is dropped. If the writer half has
1399    /// not yet written a value then when it attempts to write a value it will
1400    /// see that this end is closed.
1401    ///
1402    /// # Errors
1403    ///
1404    /// Returns an error if this future has already been closed or otherwise
1405    /// transferred (e.g. lowered into a guest), or if it does not belong to
1406    /// `store`.
1407    pub fn close(&mut self, mut store: impl AsContextMut) -> Result<()> {
1408        future_close(store.as_context_mut().0, self.idx)
1409    }
1410
1411    /// Convenience method around [`Self::close`].
1412    pub fn close_with(&mut self, accessor: impl AsAccessor) -> Result<()> {
1413        accessor.as_accessor().with(|access| self.close(access))
1414    }
1415
1416    /// Returns a [`GuardedFutureReader`] which will auto-close this future on
1417    /// drop and clean it up from the store.
1418    ///
1419    /// Note that the `accessor` provided must own this future and is
1420    /// additionally transferred to the `GuardedFutureReader` return value.
1421    pub fn guard<A>(self, accessor: A) -> GuardedFutureReader<T, A>
1422    where
1423        A: AsAccessor,
1424    {
1425        GuardedFutureReader::new(accessor, self)
1426    }
1427
1428    /// Attempts to convert this [`FutureReader<T>`] to a [`FutureAny`].
1429    ///
1430    /// # Errors
1431    ///
1432    /// This function will return an error if `self` does not belong to
1433    /// `store`.
1434    pub fn try_into_future_any(self, store: impl AsContextMut) -> Result<FutureAny>
1435    where
1436        T: ComponentType + 'static,
1437    {
1438        FutureAny::try_from_future_reader(store, self)
1439    }
1440
1441    /// Attempts to convert a [`FutureAny`] into a [`FutureReader<T>`].
1442    ///
1443    /// # Errors
1444    ///
1445    /// This function will fail if `T` doesn't match the type of the value that
1446    /// `future` is sending.
1447    pub fn try_from_future_any(future: FutureAny) -> Result<Self>
1448    where
1449        T: ComponentType + 'static,
1450    {
1451        future.try_into_future_reader()
1452    }
1453}
1454
1455impl<T> fmt::Debug for FutureReader<T> {
1456    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1457        f.debug_struct("FutureReader")
1458            .field("idx", &self.idx)
1459            .finish()
1460    }
1461}
1462
1463pub(super) fn future_close(store: &mut StoreOpaque, idx: HostResourceIndex) -> Result<()> {
1464    let id = store.host_reader_take(idx)?;
1465    store.host_drop_reader(id, TransmitKind::Future)
1466}
1467
1468/// Transfer ownership of the read end of a future from a guest to the host.
1469pub(super) fn lift_index_to_future(
1470    cx: &mut LiftContext<'_>,
1471    ty: InterfaceType,
1472    index: u32,
1473) -> Result<HostResourceIndex> {
1474    match ty {
1475        InterfaceType::Future(src) => {
1476            let id = {
1477                let (state, instance, counter) = cx.concurrent_state_and_instance_mut();
1478                lift_index_to_transmit(instance, counter, state, TransmitIndex::Future(src), index)?
1479            };
1480            cx.host_resource_lower_own(id.rep(), None, None)
1481        }
1482        _ => func::bad_type_info(),
1483    }
1484}
1485
1486/// Checks that the host-owned read end `idx` may be lowered into a guest as a
1487/// future or stream with the payload type `payload`.
1488///
1489/// Note that this is technically redundant for `StreamReader<T>` where its
1490/// static typecheck already verified this, but this is otherwise required in
1491/// the face of the `{Stream,Future}Any` abstraction to handle various use cases
1492/// of types dynamically being labeled or possibly switched.
1493fn typecheck_lowered_transmit<U>(
1494    idx: HostResourceIndex,
1495    cx: &mut LowerContext<'_, U>,
1496    payload: Option<&InterfaceType>,
1497) -> Result<()> {
1498    let origin = cx.store.0.transmit_origin(idx)?;
1499    let dst = cx.instance_type();
1500    let equivalent = match origin {
1501        TransmitOrigin::Host { typecheck } => {
1502            return typecheck(payload, &dst).context("future or stream payload types differ");
1503        }
1504        TransmitOrigin::GuestFuture(id, ty) => {
1505            let src = InstanceType::new(cx.store.0.component_instance(id));
1506            types::FutureType::from(src.types[ty].ty, &src).equivalent_payload_guest(&dst, payload)
1507        }
1508        TransmitOrigin::GuestStream(id, ty) => {
1509            let src = InstanceType::new(cx.store.0.component_instance(id));
1510            types::StreamType::from(src.types[ty].ty, &src).equivalent_payload_guest(&dst, payload)
1511        }
1512    };
1513    if !equivalent {
1514        bail!("future or stream payload types differ");
1515    }
1516    Ok(())
1517}
1518
1519/// Transfer ownership of the read end of a future from the host to a guest.
1520pub(super) fn lower_future_to_index<U>(
1521    idx: HostResourceIndex,
1522    cx: &mut LowerContext<'_, U>,
1523    ty: InterfaceType,
1524) -> Result<u32> {
1525    match ty {
1526        InterfaceType::Future(dst) => {
1527            let types = cx.types;
1528            typecheck_lowered_transmit(idx, cx, types[types[dst].ty].payload.as_ref())?;
1529            let id = cx.store.0.host_reader_take(idx)?;
1530            cx.instance_handle()
1531                .lower_transmit_to_index(cx.store.0, TransmitIndex::Future(dst), id)
1532        }
1533        _ => func::bad_type_info(),
1534    }
1535}
1536
1537// SAFETY: This relies on the `ComponentType` implementation for `u32` being
1538// safe and correct since we lift and lower future handles as `u32`s.
1539unsafe impl<T: ComponentType> ComponentType for FutureReader<T> {
1540    const ABI: CanonicalAbiInfo = CanonicalAbiInfo::SCALAR4;
1541
1542    type Lower = <u32 as func::ComponentType>::Lower;
1543
1544    fn typecheck(ty: &InterfaceType, types: &InstanceType<'_>) -> Result<()> {
1545        match ty {
1546            InterfaceType::Future(ty) => {
1547                let ty = types.types[*ty].ty;
1548                types::typecheck_payload::<T>(types.types[ty].payload.as_ref(), types)
1549            }
1550            other => bail!("expected `future`, found `{}`", func::desc(other)),
1551        }
1552    }
1553}
1554
1555// SAFETY: See the comment on the `ComponentType` `impl` for this type.
1556unsafe impl<T: ComponentType> func::Lower for FutureReader<T> {
1557    fn linear_lower_to_flat<U>(
1558        &self,
1559        cx: &mut LowerContext<'_, U>,
1560        ty: InterfaceType,
1561        dst: &mut MaybeUninit<Self::Lower>,
1562    ) -> Result<()> {
1563        lower_future_to_index(self.idx, cx, ty)?.linear_lower_to_flat(cx, InterfaceType::U32, dst)
1564    }
1565
1566    fn linear_lower_to_memory<U>(
1567        &self,
1568        cx: &mut LowerContext<'_, U>,
1569        ty: InterfaceType,
1570        offset: usize,
1571    ) -> Result<()> {
1572        lower_future_to_index(self.idx, cx, ty)?.linear_lower_to_memory(
1573            cx,
1574            InterfaceType::U32,
1575            offset,
1576        )
1577    }
1578}
1579
1580// SAFETY: See the comment on the `ComponentType` `impl` for this type.
1581unsafe impl<T: ComponentType> func::Lift for FutureReader<T> {
1582    fn linear_lift_from_flat(
1583        cx: &mut LiftContext<'_>,
1584        ty: InterfaceType,
1585        src: &Self::Lower,
1586    ) -> Result<Self> {
1587        let index = u32::linear_lift_from_flat(cx, InterfaceType::U32, src)?;
1588        Self::lift_from_index(cx, ty, index)
1589    }
1590
1591    fn linear_lift_from_memory(
1592        cx: &mut LiftContext<'_>,
1593        ty: InterfaceType,
1594        bytes: &[u8],
1595    ) -> Result<Self> {
1596        let index = u32::linear_lift_from_memory(cx, InterfaceType::U32, bytes)?;
1597        Self::lift_from_index(cx, ty, index)
1598    }
1599}
1600
1601/// A [`FutureReader`] paired with an [`Accessor`].
1602///
1603/// This is an RAII wrapper around [`FutureReader`] that ensures it is closed
1604/// when dropped. This can be created through [`GuardedFutureReader::new`] or
1605/// [`FutureReader::guard`].
1606///
1607/// [`Accessor`]: crate::component::Accessor
1608pub struct GuardedFutureReader<T, A>
1609where
1610    A: AsAccessor,
1611{
1612    // This field is `None` to implement the conversion from this guard back to
1613    // `FutureReader`. When `None` is seen in the destructor it will cause the
1614    // destructor to do nothing.
1615    reader: Option<FutureReader<T>>,
1616    accessor: A,
1617}
1618
1619impl<T, A> GuardedFutureReader<T, A>
1620where
1621    A: AsAccessor,
1622{
1623    /// Create a new `GuardedFutureReader` with the specified `accessor` and `reader`.
1624    ///
1625    /// # Panics
1626    ///
1627    /// Panics if [`Config::concurrency_support`] is not enabled.
1628    ///
1629    /// [`Config::concurrency_support`]: crate::Config::concurrency_support
1630    pub fn new(accessor: A, reader: FutureReader<T>) -> Self {
1631        assert!(
1632            accessor
1633                .as_accessor()
1634                .with(|a| a.as_context().0.concurrency_support())
1635        );
1636        Self {
1637            reader: Some(reader),
1638            accessor,
1639        }
1640    }
1641
1642    /// Extracts the underlying [`FutureReader`] from this guard, returning it
1643    /// back.
1644    pub fn into_future(self) -> FutureReader<T> {
1645        self.into()
1646    }
1647}
1648
1649impl<T, A> From<GuardedFutureReader<T, A>> for FutureReader<T>
1650where
1651    A: AsAccessor,
1652{
1653    fn from(mut guard: GuardedFutureReader<T, A>) -> Self {
1654        guard.reader.take().unwrap()
1655    }
1656}
1657
1658impl<T, A> Drop for GuardedFutureReader<T, A>
1659where
1660    A: AsAccessor,
1661{
1662    fn drop(&mut self) {
1663        if let Some(reader) = &mut self.reader {
1664            // Currently this can only fail if the future is closed twice or
1665            // the handle was cloned/closed elsewhere, all of which the embedder
1666            // is responsible for when using this guard. Hint during
1667            // development that's a bad idea, but otherwise ignore the error in
1668            // release builds.
1669            let result = reader.close_with(&self.accessor);
1670            debug_assert!(result.is_ok());
1671        }
1672    }
1673}
1674
1675/// Represents the readable end of a Component Model `stream`.
1676///
1677/// Note that `StreamReader` instances must be disposed of using `close`;
1678/// otherwise the in-store representation will leak and the writer end will hang
1679/// indefinitely.  Consider using [`GuardedStreamReader`] to ensure that
1680/// disposal happens automatically.
1681pub struct StreamReader<T> {
1682    idx: HostResourceIndex,
1683    _phantom: PhantomData<T>,
1684}
1685
1686impl<T> StreamReader<T> {
1687    /// Create a new stream with the specified producer.
1688    ///
1689    /// # Errors
1690    ///
1691    /// Returns an error if the resource table for this store is full or if
1692    /// [`Config::concurrency_support`] is not enabled.
1693    ///
1694    /// [`Config::concurrency_support`]: crate::Config::concurrency_support
1695    pub fn new<S: AsContextMut>(
1696        mut store: S,
1697        producer: impl StreamProducer<S::Data, Item = T>,
1698    ) -> Result<Self>
1699    where
1700        T: func::Lower + func::Lift + Send + Sync + 'static,
1701    {
1702        ensure!(
1703            store.as_context().0.concurrency_support(),
1704            "concurrency support is not enabled",
1705        );
1706        Ok(Self::new_(
1707            store
1708                .as_context_mut()
1709                .new_transmit(TransmitKind::Stream, producer)?,
1710        ))
1711    }
1712
1713    pub(super) fn new_(idx: HostResourceIndex) -> Self {
1714        Self {
1715            idx,
1716            _phantom: PhantomData,
1717        }
1718    }
1719
1720    pub(super) fn idx(&self) -> HostResourceIndex {
1721        self.idx
1722    }
1723
1724    /// Attempt to consume this object by converting it into the specified type.
1725    ///
1726    /// This can be useful for "short-circuiting" host-to-host streams,
1727    /// bypassing the guest entirely.  For example, if a guest task returns a
1728    /// host-created stream and then exits, this function may be used to
1729    /// retrieve the write end, after which the guest instance and store may be
1730    /// disposed of if no longer needed.
1731    ///
1732    /// This will return `Ok(_)` if and only if the following conditions are
1733    /// met:
1734    ///
1735    /// - The stream was created by the host (i.e. not by the guest).
1736    ///
1737    /// - The `StreamProducer::try_into` function returns `Ok(_)` when given the
1738    /// producer provided to `StreamReader::new` when the stream was created,
1739    /// along with `TypeId::of::<V>()`.
1740    ///
1741    /// # Panics
1742    ///
1743    /// Panics if this stream has already been closed, or if this stream doesn't
1744    /// belong to the specified `store`.
1745    pub fn try_into<V: 'static>(mut self, mut store: impl AsContextMut) -> Result<V, Self> {
1746        let store = store.as_context_mut();
1747        let id = store.0.host_reader_rep(self.idx).unwrap();
1748        let state = store.0.concurrent_state_mut_already_forced_current_thread();
1749        let id = state.get_mut(id).unwrap().state;
1750        if let WriteState::HostReady { try_into, .. } = &state.get_mut(id).unwrap().write {
1751            match try_into(TypeId::of::<V>()) {
1752                Some(result) => {
1753                    self.close(store).unwrap();
1754                    Ok(*result.downcast::<V>().unwrap())
1755                }
1756                None => Err(self),
1757            }
1758        } else {
1759            Err(self)
1760        }
1761    }
1762
1763    /// Set the consumer that accepts the items delivered to this stream.
1764    ///
1765    /// # Errors
1766    ///
1767    /// Returns an error if this stream has already been closed or otherwise
1768    /// transferred (e.g. lowered into a guest), or if it does not belong to
1769    /// `store`.
1770    pub fn pipe<S: AsContextMut>(
1771        self,
1772        mut store: S,
1773        consumer: impl StreamConsumer<S::Data, Item = T>,
1774    ) -> Result<()>
1775    where
1776        T: 'static,
1777    {
1778        store
1779            .as_context_mut()
1780            .set_consumer(self.idx, TransmitKind::Stream, consumer)
1781    }
1782
1783    /// Transfer ownership of the read end of a stream from a guest to the host.
1784    fn lift_from_index(cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32) -> Result<Self> {
1785        Ok(Self::new_(lift_index_to_stream(cx, ty, index)?))
1786    }
1787
1788    /// Close this `StreamReader`.
1789    ///
1790    /// This will signal that this portion of the stream is closed causing all
1791    /// future writes to return immediately with "DROPPED".
1792    ///
1793    /// # Errors
1794    ///
1795    /// Returns an error if this stream has already been closed or otherwise
1796    /// transferred (e.g. lowered into a guest), or if it does not belong to
1797    /// `store`.
1798    pub fn close(&mut self, mut store: impl AsContextMut) -> Result<()> {
1799        stream_close(store.as_context_mut().0, self.idx)
1800    }
1801
1802    /// Convenience method around [`Self::close`].
1803    pub fn close_with(&mut self, accessor: impl AsAccessor) -> Result<()> {
1804        accessor.as_accessor().with(|access| self.close(access))
1805    }
1806
1807    /// Returns a [`GuardedStreamReader`] which will auto-close this stream on
1808    /// drop and clean it up from the store.
1809    ///
1810    /// Note that the `accessor` provided must own this future and is
1811    /// additionally transferred to the `GuardedStreamReader` return value.
1812    pub fn guard<A>(self, accessor: A) -> GuardedStreamReader<T, A>
1813    where
1814        A: AsAccessor,
1815    {
1816        GuardedStreamReader::new(accessor, self)
1817    }
1818
1819    /// Attempts to convert this [`StreamReader<T>`] to a [`StreamAny`].
1820    ///
1821    /// # Errors
1822    ///
1823    /// This function will return an error if `self` does not belong to
1824    /// `store`.
1825    pub fn try_into_stream_any(self, store: impl AsContextMut) -> Result<StreamAny>
1826    where
1827        T: ComponentType + 'static,
1828    {
1829        StreamAny::try_from_stream_reader(store, self)
1830    }
1831
1832    /// Attempts to convert a [`StreamAny`] into a [`StreamReader<T>`].
1833    ///
1834    /// # Errors
1835    ///
1836    /// This function will fail if `T` doesn't match the type of the value that
1837    /// `stream` is sending.
1838    pub fn try_from_stream_any(stream: StreamAny) -> Result<Self>
1839    where
1840        T: ComponentType + 'static,
1841    {
1842        stream.try_into_stream_reader()
1843    }
1844}
1845
1846impl<T> fmt::Debug for StreamReader<T> {
1847    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1848        f.debug_struct("StreamReader")
1849            .field("idx", &self.idx)
1850            .finish()
1851    }
1852}
1853
1854pub(super) fn stream_close(store: &mut StoreOpaque, idx: HostResourceIndex) -> Result<()> {
1855    let id = store.host_reader_take(idx)?;
1856    store.host_drop_reader(id, TransmitKind::Stream)
1857}
1858
1859/// Transfer ownership of the read end of a stream from a guest to the host.
1860pub(super) fn lift_index_to_stream(
1861    cx: &mut LiftContext<'_>,
1862    ty: InterfaceType,
1863    index: u32,
1864) -> Result<HostResourceIndex> {
1865    match ty {
1866        InterfaceType::Stream(src) => {
1867            let id = {
1868                let (state, instance, counter) = cx.concurrent_state_and_instance_mut();
1869                lift_index_to_transmit(instance, counter, state, TransmitIndex::Stream(src), index)?
1870            };
1871            cx.host_resource_lower_own(id.rep(), None, None)
1872        }
1873        _ => func::bad_type_info(),
1874    }
1875}
1876
1877/// Transfer ownership of the read end of a stream from the host to a guest.
1878pub(super) fn lower_stream_to_index<U>(
1879    idx: HostResourceIndex,
1880    cx: &mut LowerContext<'_, U>,
1881    ty: InterfaceType,
1882) -> Result<u32> {
1883    match ty {
1884        InterfaceType::Stream(dst) => {
1885            let types = cx.types;
1886            typecheck_lowered_transmit(idx, cx, types[types[dst].ty].payload.as_ref())?;
1887            let id = cx.store.0.host_reader_take(idx)?;
1888            cx.instance_handle()
1889                .lower_transmit_to_index(cx.store.0, TransmitIndex::Stream(dst), id)
1890        }
1891        _ => func::bad_type_info(),
1892    }
1893}
1894
1895// SAFETY: This relies on the `ComponentType` implementation for `u32` being
1896// safe and correct since we lift and lower stream handles as `u32`s.
1897unsafe impl<T: ComponentType> ComponentType for StreamReader<T> {
1898    const ABI: CanonicalAbiInfo = CanonicalAbiInfo::SCALAR4;
1899
1900    type Lower = <u32 as func::ComponentType>::Lower;
1901
1902    fn typecheck(ty: &InterfaceType, types: &InstanceType<'_>) -> Result<()> {
1903        match ty {
1904            InterfaceType::Stream(ty) => {
1905                let ty = types.types[*ty].ty;
1906                types::typecheck_payload::<T>(types.types[ty].payload.as_ref(), types)
1907            }
1908            other => bail!("expected `stream`, found `{}`", func::desc(other)),
1909        }
1910    }
1911}
1912
1913// SAFETY: See the comment on the `ComponentType` `impl` for this type.
1914unsafe impl<T: ComponentType> func::Lower for StreamReader<T> {
1915    fn linear_lower_to_flat<U>(
1916        &self,
1917        cx: &mut LowerContext<'_, U>,
1918        ty: InterfaceType,
1919        dst: &mut MaybeUninit<Self::Lower>,
1920    ) -> Result<()> {
1921        lower_stream_to_index(self.idx, cx, ty)?.linear_lower_to_flat(cx, InterfaceType::U32, dst)
1922    }
1923
1924    fn linear_lower_to_memory<U>(
1925        &self,
1926        cx: &mut LowerContext<'_, U>,
1927        ty: InterfaceType,
1928        offset: usize,
1929    ) -> Result<()> {
1930        lower_stream_to_index(self.idx, cx, ty)?.linear_lower_to_memory(
1931            cx,
1932            InterfaceType::U32,
1933            offset,
1934        )
1935    }
1936}
1937
1938// SAFETY: See the comment on the `ComponentType` `impl` for this type.
1939unsafe impl<T: ComponentType> func::Lift for StreamReader<T> {
1940    fn linear_lift_from_flat(
1941        cx: &mut LiftContext<'_>,
1942        ty: InterfaceType,
1943        src: &Self::Lower,
1944    ) -> Result<Self> {
1945        let index = u32::linear_lift_from_flat(cx, InterfaceType::U32, src)?;
1946        Self::lift_from_index(cx, ty, index)
1947    }
1948
1949    fn linear_lift_from_memory(
1950        cx: &mut LiftContext<'_>,
1951        ty: InterfaceType,
1952        bytes: &[u8],
1953    ) -> Result<Self> {
1954        let index = u32::linear_lift_from_memory(cx, InterfaceType::U32, bytes)?;
1955        Self::lift_from_index(cx, ty, index)
1956    }
1957}
1958
1959/// A [`StreamReader`] paired with an [`Accessor`].
1960///
1961/// This is an RAII wrapper around [`StreamReader`] that ensures it is closed
1962/// when dropped. This can be created through [`GuardedStreamReader::new`] or
1963/// [`StreamReader::guard`].
1964///
1965/// [`Accessor`]: crate::component::Accessor
1966pub struct GuardedStreamReader<T, A>
1967where
1968    A: AsAccessor,
1969{
1970    // This field is `None` to implement the conversion from this guard back to
1971    // `StreamReader`. When `None` is seen in the destructor it will cause the
1972    // destructor to do nothing.
1973    reader: Option<StreamReader<T>>,
1974    accessor: A,
1975}
1976
1977impl<T, A> GuardedStreamReader<T, A>
1978where
1979    A: AsAccessor,
1980{
1981    /// Create a new `GuardedStreamReader` with the specified `accessor` and
1982    /// `reader`.
1983    ///
1984    /// # Panics
1985    ///
1986    /// Panics if [`Config::concurrency_support`] is not enabled.
1987    ///
1988    /// [`Config::concurrency_support`]: crate::Config::concurrency_support
1989    pub fn new(accessor: A, reader: StreamReader<T>) -> Self {
1990        assert!(
1991            accessor
1992                .as_accessor()
1993                .with(|a| a.as_context().0.concurrency_support())
1994        );
1995        Self {
1996            reader: Some(reader),
1997            accessor,
1998        }
1999    }
2000
2001    /// Extracts the underlying [`StreamReader`] from this guard, returning it
2002    /// back.
2003    pub fn into_stream(self) -> StreamReader<T> {
2004        self.into()
2005    }
2006}
2007
2008impl<T, A> From<GuardedStreamReader<T, A>> for StreamReader<T>
2009where
2010    A: AsAccessor,
2011{
2012    fn from(mut guard: GuardedStreamReader<T, A>) -> Self {
2013        guard.reader.take().unwrap()
2014    }
2015}
2016
2017impl<T, A> Drop for GuardedStreamReader<T, A>
2018where
2019    A: AsAccessor,
2020{
2021    fn drop(&mut self) {
2022        if let Some(reader) = &mut self.reader {
2023            // See comments in `Drop for GuardedFutureReader` for this assert.
2024            let result = reader.close_with(&self.accessor);
2025            debug_assert!(result.is_ok());
2026        }
2027    }
2028}
2029
2030/// Represents a Component Model `error-context`.
2031pub struct ErrorContext {
2032    rep: u32,
2033}
2034
2035impl ErrorContext {
2036    pub(crate) fn new(rep: u32) -> Self {
2037        Self { rep }
2038    }
2039
2040    /// Convert this `ErrorContext` into a [`Val`].
2041    pub fn into_val(self) -> Val {
2042        Val::ErrorContext(ErrorContextAny(self.rep))
2043    }
2044
2045    /// Attempt to convert the specified [`Val`] to a `ErrorContext`.
2046    pub fn from_val(_: impl AsContextMut, value: &Val) -> Result<Self> {
2047        let Val::ErrorContext(ErrorContextAny(rep)) = value else {
2048            bail!("expected `error-context`; got `{}`", value.desc());
2049        };
2050        Ok(Self::new(*rep))
2051    }
2052
2053    fn lift_from_index(cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32) -> Result<Self> {
2054        match ty {
2055            InterfaceType::ErrorContext(src) => {
2056                let rep = cx
2057                    .instance_mut()
2058                    .table_for_error_context(src)
2059                    .error_context_rep(index)?;
2060
2061                Ok(Self { rep })
2062            }
2063            _ => func::bad_type_info(),
2064        }
2065    }
2066}
2067
2068pub(crate) fn lower_error_context_to_index<U>(
2069    rep: u32,
2070    cx: &mut LowerContext<'_, U>,
2071    ty: InterfaceType,
2072) -> Result<u32> {
2073    match ty {
2074        InterfaceType::ErrorContext(dst) => {
2075            let (instance, counter) = cx
2076                .instance_handle()
2077                .id()
2078                .get_mut_and_handle_counter(cx.store.0);
2079            instance
2080                .table_for_error_context(dst)
2081                .error_context_insert(counter, rep)
2082        }
2083        _ => func::bad_type_info(),
2084    }
2085}
2086// SAFETY: This relies on the `ComponentType` implementation for `u32` being
2087// safe and correct since we lift and lower future handles as `u32`s.
2088unsafe impl func::ComponentType for ErrorContext {
2089    const ABI: CanonicalAbiInfo = CanonicalAbiInfo::SCALAR4;
2090
2091    type Lower = <u32 as func::ComponentType>::Lower;
2092
2093    fn typecheck(ty: &InterfaceType, _types: &InstanceType<'_>) -> Result<()> {
2094        match ty {
2095            InterfaceType::ErrorContext(_) => Ok(()),
2096            other => bail!("expected `error`, found `{}`", func::desc(other)),
2097        }
2098    }
2099}
2100
2101// SAFETY: See the comment on the `ComponentType` `impl` for this type.
2102unsafe impl func::Lower for ErrorContext {
2103    fn linear_lower_to_flat<T>(
2104        &self,
2105        cx: &mut LowerContext<'_, T>,
2106        ty: InterfaceType,
2107        dst: &mut MaybeUninit<Self::Lower>,
2108    ) -> Result<()> {
2109        lower_error_context_to_index(self.rep, cx, ty)?.linear_lower_to_flat(
2110            cx,
2111            InterfaceType::U32,
2112            dst,
2113        )
2114    }
2115
2116    fn linear_lower_to_memory<T>(
2117        &self,
2118        cx: &mut LowerContext<'_, T>,
2119        ty: InterfaceType,
2120        offset: usize,
2121    ) -> Result<()> {
2122        lower_error_context_to_index(self.rep, cx, ty)?.linear_lower_to_memory(
2123            cx,
2124            InterfaceType::U32,
2125            offset,
2126        )
2127    }
2128}
2129
2130// SAFETY: See the comment on the `ComponentType` `impl` for this type.
2131unsafe impl func::Lift for ErrorContext {
2132    fn linear_lift_from_flat(
2133        cx: &mut LiftContext<'_>,
2134        ty: InterfaceType,
2135        src: &Self::Lower,
2136    ) -> Result<Self> {
2137        let index = u32::linear_lift_from_flat(cx, InterfaceType::U32, src)?;
2138        Self::lift_from_index(cx, ty, index)
2139    }
2140
2141    fn linear_lift_from_memory(
2142        cx: &mut LiftContext<'_>,
2143        ty: InterfaceType,
2144        bytes: &[u8],
2145    ) -> Result<Self> {
2146        let index = u32::linear_lift_from_memory(cx, InterfaceType::U32, bytes)?;
2147        Self::lift_from_index(cx, ty, index)
2148    }
2149}
2150
2151/// Represents the read or write end of a stream or future.
2152pub(super) struct TransmitHandle {
2153    pub(super) common: WaitableCommon,
2154    /// See `TransmitState`
2155    state: TableId<TransmitState>,
2156}
2157
2158impl TransmitHandle {
2159    fn new(state: TableId<TransmitState>) -> Self {
2160        Self {
2161            common: WaitableCommon::default(),
2162            state,
2163        }
2164    }
2165}
2166
2167impl TableDebug for TransmitHandle {
2168    fn type_name() -> &'static str {
2169        "TransmitHandle"
2170    }
2171}
2172
2173/// Represents the state of a stream or future.
2174struct TransmitState {
2175    /// The write end of the stream or future.
2176    write_handle: TableId<TransmitHandle>,
2177    /// The read end of the stream or future.
2178    read_handle: TableId<TransmitHandle>,
2179    /// See `WriteState`
2180    write: WriteState,
2181    /// See `ReadState`
2182    read: ReadState,
2183    /// Whether further values may be transmitted via this stream or future.
2184    done: bool,
2185    /// The original creator of this stream or future, whose payload type is
2186    /// checked against a guest's whenever the host lowers the read end into
2187    /// that guest, and which `{Future,Stream}Any` use as their initial type.
2188    pub(super) origin: TransmitOrigin,
2189}
2190
2191#[derive(Copy, Clone)]
2192pub(super) enum TransmitOrigin {
2193    Host {
2194        /// Type-checks the producer's items against the guest-supplied payload.
2195        typecheck: fn(Option<&InterfaceType>, &InstanceType<'_>) -> Result<()>,
2196    },
2197    GuestFuture(ComponentInstanceId, TypeFutureTableIndex),
2198    GuestStream(ComponentInstanceId, TypeStreamTableIndex),
2199}
2200
2201impl TransmitState {
2202    fn new(origin: TransmitOrigin) -> Self {
2203        Self {
2204            write_handle: TableId::new(u32::MAX),
2205            read_handle: TableId::new(u32::MAX),
2206            read: ReadState::Open,
2207            write: WriteState::Open,
2208            done: false,
2209            origin,
2210        }
2211    }
2212}
2213
2214impl TableDebug for TransmitState {
2215    fn type_name() -> &'static str {
2216        "TransmitState"
2217    }
2218}
2219
2220impl TransmitOrigin {
2221    fn guest(id: ComponentInstanceId, index: TransmitIndex) -> Self {
2222        match index {
2223            TransmitIndex::Future(ty) => TransmitOrigin::GuestFuture(id, ty),
2224            TransmitIndex::Stream(ty) => TransmitOrigin::GuestStream(id, ty),
2225        }
2226    }
2227}
2228
2229type PollStream = Box<
2230    dyn Fn() -> Pin<Box<dyn Future<Output = Result<StreamResult>> + Send + 'static>> + Send + Sync,
2231>;
2232
2233type TryInto = Box<dyn Fn(TypeId) -> Option<Box<dyn Any>> + Send + Sync>;
2234
2235/// Represents the state of the write end of a stream or future.
2236enum WriteState {
2237    /// The write end is open, but no write is pending.
2238    Open,
2239    /// The write end is owned by a guest task and a write is pending.
2240    GuestReady {
2241        instance: Instance,
2242        caller: RuntimeComponentInstanceIndex,
2243        ty: TransmitIndex,
2244        flat_abi: Option<FlatAbi>,
2245        options: OptionsIndex,
2246        address: usize,
2247        count: ItemCount,
2248        handle: u32,
2249    },
2250    /// The write end is owned by the host, which is ready to produce items.
2251    HostReady {
2252        produce: PollStream,
2253        try_into: TryInto,
2254        guest_offset: ItemCount,
2255        cancel: bool,
2256        cancel_waker: Option<Waker>,
2257    },
2258    /// The write end has been dropped.
2259    Dropped,
2260}
2261
2262impl fmt::Debug for WriteState {
2263    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
2264        match self {
2265            Self::Open => f.debug_tuple("Open").finish(),
2266            Self::GuestReady { .. } => f.debug_tuple("GuestReady").finish(),
2267            Self::HostReady { .. } => f.debug_tuple("HostReady").finish(),
2268            Self::Dropped => f.debug_tuple("Dropped").finish(),
2269        }
2270    }
2271}
2272
2273/// Represents the state of the read end of a stream or future.
2274enum ReadState {
2275    /// The read end is open, but no read is pending.
2276    Open,
2277    /// The read end is owned by a guest task and a read is pending.
2278    GuestReady {
2279        ty: TransmitIndex,
2280        caller_instance: RuntimeComponentInstanceIndex,
2281        flat_abi: Option<FlatAbi>,
2282        instance: Instance,
2283        options: OptionsIndex,
2284        address: usize,
2285        count: ItemCount,
2286        handle: u32,
2287    },
2288    /// The read end is owned by a host task, and it is ready to consume items.
2289    HostReady {
2290        consume: PollStream,
2291        guest_offset: ItemCount,
2292        cancel: bool,
2293        cancel_waker: Option<Waker>,
2294    },
2295    /// Both the read and write ends are owned by the host.
2296    HostToHost {
2297        accept: Box<
2298            dyn for<'a> Fn(
2299                    &'a mut UntypedWriteBuffer<'a>,
2300                )
2301                    -> Pin<Box<dyn Future<Output = Result<StreamResult>> + Send + 'a>>
2302                + Send
2303                + Sync,
2304        >,
2305        buffer: Vec<u8>,
2306        limit: usize,
2307    },
2308    /// The read end has been dropped.
2309    Dropped,
2310}
2311
2312impl fmt::Debug for ReadState {
2313    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
2314        match self {
2315            Self::Open => f.debug_tuple("Open").finish(),
2316            Self::GuestReady { .. } => f.debug_tuple("GuestReady").finish(),
2317            Self::HostReady { .. } => f.debug_tuple("HostReady").finish(),
2318            Self::HostToHost { .. } => f.debug_tuple("HostToHost").finish(),
2319            Self::Dropped => f.debug_tuple("Dropped").finish(),
2320        }
2321    }
2322}
2323
2324fn return_code(kind: TransmitKind, state: StreamResult, count: ItemCount) -> Result<ReturnCode> {
2325    Ok(match state {
2326        StreamResult::Dropped => ReturnCode::Dropped(count),
2327        StreamResult::Completed => ReturnCode::completed(kind, count),
2328        StreamResult::Cancelled => ReturnCode::Cancelled(count),
2329    })
2330}
2331
2332fn settle_host_read(
2333    transmit: &mut TransmitState,
2334    kind: TransmitKind,
2335    state: StreamResult,
2336) -> Result<ReturnCode> {
2337    let ReadState::HostReady {
2338        consume,
2339        guest_offset,
2340        ..
2341    } = mem::replace(&mut transmit.read, ReadState::Open)
2342    else {
2343        bail_bug!("expected ReadState::HostReady")
2344    };
2345    let code = return_code(kind, state, guest_offset)?;
2346    transmit.read = match state {
2347        StreamResult::Dropped => ReadState::Dropped,
2348        StreamResult::Completed | StreamResult::Cancelled => ReadState::HostReady {
2349            consume,
2350            guest_offset: ItemCount::ZERO,
2351            cancel: false,
2352            cancel_waker: None,
2353        },
2354    };
2355    Ok(code)
2356}
2357
2358fn settle_host_write(
2359    transmit: &mut TransmitState,
2360    kind: TransmitKind,
2361    state: StreamResult,
2362) -> Result<ReturnCode> {
2363    let WriteState::HostReady {
2364        produce,
2365        try_into,
2366        guest_offset,
2367        ..
2368    } = mem::replace(&mut transmit.write, WriteState::Open)
2369    else {
2370        bail_bug!("expected WriteState::HostReady")
2371    };
2372    let code = return_code(kind, state, guest_offset)?;
2373    transmit.write = match state {
2374        StreamResult::Dropped => WriteState::Dropped,
2375        StreamResult::Completed | StreamResult::Cancelled => WriteState::HostReady {
2376            produce,
2377            try_into,
2378            guest_offset: ItemCount::ZERO,
2379            cancel: false,
2380            cancel_waker: None,
2381        },
2382    };
2383    Ok(code)
2384}
2385
2386impl StoreOpaque {
2387    fn pipe_from_guest(
2388        &mut self,
2389        kind: TransmitKind,
2390        id: TableId<TransmitState>,
2391        future: Pin<Box<dyn Future<Output = Result<StreamResult>> + Send + 'static>>,
2392    ) {
2393        let future = async move {
2394            let stream_state = future.await?;
2395            tls::get(|store| {
2396                let state = store.concurrent_state_mut()?;
2397                let transmit = state.get_mut(id)?;
2398                let code = settle_host_read(transmit, kind, stream_state)?;
2399                let WriteState::GuestReady { ty, handle, .. } =
2400                    mem::replace(&mut transmit.write, WriteState::Open)
2401                else {
2402                    bail_bug!("expected WriteState::GuestReady")
2403                };
2404                state.send_write_result(ty, id, handle, code)?;
2405                Ok(())
2406            })
2407        };
2408
2409        self.concurrent_state_mut_already_forced_current_thread()
2410            .push_future(future.boxed());
2411    }
2412
2413    fn pipe_to_guest(
2414        &mut self,
2415        kind: TransmitKind,
2416        id: TableId<TransmitState>,
2417        future: Pin<Box<dyn Future<Output = Result<StreamResult>> + Send + 'static>>,
2418    ) {
2419        let future = async move {
2420            let stream_state = future.await?;
2421            tls::get(|store| {
2422                let state = store.concurrent_state_mut()?;
2423                let transmit = state.get_mut(id)?;
2424                let code = settle_host_write(transmit, kind, stream_state)?;
2425                let ReadState::GuestReady { ty, handle, .. } =
2426                    mem::replace(&mut transmit.read, ReadState::Open)
2427                else {
2428                    bail_bug!("expected ReadState::GuestReady")
2429                };
2430                state.send_read_result(ty, id, handle, code)?;
2431                Ok(())
2432            })
2433        };
2434
2435        self.concurrent_state_mut_already_forced_current_thread()
2436            .push_future(future.boxed());
2437    }
2438
2439    /// Drop the read end of a stream or future read from the host.
2440    fn host_drop_reader(&mut self, id: TableId<TransmitHandle>, kind: TransmitKind) -> Result<()> {
2441        let state = self.concurrent_state_mut()?;
2442        Waitable::Transmit(id).join(state, None)?;
2443        let transmit_id = state.get_mut(id)?.state;
2444        let transmit = state
2445            .get_mut(transmit_id)
2446            .with_context(|| format!("error closing reader {transmit_id:?}"))?;
2447        log::trace!(
2448            "host_drop_reader state {transmit_id:?}; read state {:?} write state {:?}",
2449            transmit.read,
2450            transmit.write
2451        );
2452
2453        transmit.read = ReadState::Dropped;
2454
2455        // If the write end is already dropped, it should stay dropped,
2456        // otherwise, it should be opened.
2457        let new_state = if let WriteState::Dropped = &transmit.write {
2458            WriteState::Dropped
2459        } else {
2460            WriteState::Open
2461        };
2462
2463        let write_handle = transmit.write_handle;
2464        let done = transmit.done;
2465
2466        match mem::replace(&mut transmit.write, new_state) {
2467            // If a guest is waiting to write, notify it that the read end has
2468            // been dropped.
2469            WriteState::GuestReady { ty, handle, .. } => {
2470                state.update_event(
2471                    write_handle.rep(),
2472                    match ty {
2473                        TransmitIndex::Future(ty) => Event::FutureWrite {
2474                            code: ReturnCode::Dropped(ItemCount::ZERO),
2475                            pending: Some((ty, handle)),
2476                        },
2477                        TransmitIndex::Stream(ty) => Event::StreamWrite {
2478                            code: ReturnCode::Dropped(ItemCount::ZERO),
2479                            pending: Some((ty, handle)),
2480                        },
2481                    },
2482                )?;
2483            }
2484
2485            // Otherwise notify the idle write end that the read end has been
2486            // dropped, unless it's already done.
2487            WriteState::Open if done => {}
2488            WriteState::Open => {
2489                state.update_event(
2490                    write_handle.rep(),
2491                    match kind {
2492                        TransmitKind::Future => Event::FutureWrite {
2493                            code: ReturnCode::Dropped(ItemCount::ZERO),
2494                            pending: None,
2495                        },
2496                        TransmitKind::Stream => Event::StreamWrite {
2497                            code: ReturnCode::Dropped(ItemCount::ZERO),
2498                            pending: None,
2499                        },
2500                    },
2501                )?;
2502            }
2503
2504            // If the writer has already been dropped, then this cleans out the
2505            // state that the reader is using. If the write is host-owned then
2506            // by cleaning this out we run the host's `Drop` implementation
2507            // which notifies it of this drop.
2508            WriteState::Dropped | WriteState::HostReady { .. } => {
2509                log::trace!("host_drop_reader delete {transmit_id:?}");
2510                state.delete_transmit(transmit_id)?;
2511            }
2512        }
2513        Ok(())
2514    }
2515
2516    /// Drop the write end of a stream or future read from the host.
2517    fn host_drop_writer(
2518        &mut self,
2519        id: TableId<TransmitHandle>,
2520        on_drop_open: Option<fn() -> Result<()>>,
2521    ) -> Result<()> {
2522        let state = self.concurrent_state_mut()?;
2523        Waitable::Transmit(id).join(state, None)?;
2524        let transmit_id = state.get_mut(id)?.state;
2525        let transmit = state
2526            .get_mut(transmit_id)
2527            .with_context(|| format!("error closing writer {transmit_id:?}"))?;
2528        log::trace!(
2529            "host_drop_writer state {transmit_id:?}; read state {:?} writer state {:?}",
2530            transmit.read,
2531            transmit.write
2532        );
2533
2534        // Existing queued transmits must be updated with information for the impending writer closure
2535        match &mut transmit.write {
2536            WriteState::GuestReady { .. } => {
2537                bail_bug!("can't call `host_drop_writer` on a guest-owned writer");
2538            }
2539            WriteState::HostReady { .. } => {}
2540            v @ WriteState::Open => {
2541                if let (Some(on_drop_open), false) = (on_drop_open, transmit.done) {
2542                    on_drop_open()?;
2543                } else {
2544                    *v = WriteState::Dropped;
2545                }
2546            }
2547            WriteState::Dropped => bail_bug!("write state is already dropped"),
2548        }
2549
2550        let transmit = self.concurrent_state_mut()?.get_mut(transmit_id)?;
2551
2552        // If the existing read state is dropped, then there's nothing to read
2553        // and we can keep it that way.
2554        //
2555        // If the read state was any other state, then we must set the new state to open
2556        // to indicate that there *is* data to be read
2557        let new_state = if let ReadState::Dropped = &transmit.read {
2558            ReadState::Dropped
2559        } else {
2560            ReadState::Open
2561        };
2562
2563        let read_handle = transmit.read_handle;
2564        let done = transmit.done;
2565
2566        // Swap in the new read state
2567        match mem::replace(&mut transmit.read, new_state) {
2568            // If the guest was ready to read, then we cannot drop the reader (or writer);
2569            // we must deliver the event, and update the state associated with the handle to
2570            // represent that a read must be performed
2571            ReadState::GuestReady { ty, handle, .. } => {
2572                // Ensure the final read of the guest is queued, with appropriate closure indicator
2573                self.concurrent_state_mut()?.update_event(
2574                    read_handle.rep(),
2575                    match ty {
2576                        TransmitIndex::Future(ty) => Event::FutureRead {
2577                            code: ReturnCode::Dropped(ItemCount::ZERO),
2578                            pending: Some((ty, handle)),
2579                        },
2580                        TransmitIndex::Stream(ty) => Event::StreamRead {
2581                            code: ReturnCode::Dropped(ItemCount::ZERO),
2582                            pending: Some((ty, handle)),
2583                        },
2584                    },
2585                )?;
2586            }
2587
2588            // Notify the read-end of the drop, unless we're already done.
2589            ReadState::Open if done => {}
2590            ReadState::Open => {
2591                self.concurrent_state_mut()?.update_event(
2592                    read_handle.rep(),
2593                    match on_drop_open {
2594                        Some(_) => Event::FutureRead {
2595                            code: ReturnCode::Dropped(ItemCount::ZERO),
2596                            pending: None,
2597                        },
2598                        None => Event::StreamRead {
2599                            code: ReturnCode::Dropped(ItemCount::ZERO),
2600                            pending: None,
2601                        },
2602                    },
2603                )?;
2604            }
2605
2606            // If the read state was already dropped, then we can remove the
2607            // transmit state completely (both writer and reader have been
2608            // dropped). If the read state is host-owned then it's additionally
2609            // deleted here as a notification that the read end has gone away.
2610            // Running the host's `Drop` implementation is what notifies it of
2611            // this event.
2612            ReadState::Dropped | ReadState::HostReady { .. } | ReadState::HostToHost { .. } => {
2613                log::trace!("host_drop_writer delete {transmit_id:?}");
2614                self.concurrent_state_mut()?.delete_transmit(transmit_id)?;
2615            }
2616        }
2617        Ok(())
2618    }
2619
2620    pub(super) fn transmit_origin(&mut self, idx: HostResourceIndex) -> Result<TransmitOrigin> {
2621        let id = self.host_reader_rep(idx)?;
2622        let state = self.concurrent_state_mut()?;
2623        let state_id = state.get_mut(id)?.state;
2624        Ok(state.get_mut(state_id)?.origin)
2625    }
2626
2627    /// Insert `id` within the store, returning a `HostResourceIndex` suitable
2628    /// to re-acquire it via the `host_reader_take` and `host_reader_rep`
2629    /// methods below.
2630    fn host_reader_insert(&mut self, id: TableId<TransmitHandle>) -> Result<HostResourceIndex> {
2631        HostResourceTables::new_host(self)?.host_resource_lower_own(id.rep(), None, None)
2632    }
2633
2634    /// Removes `idx` from the store, returning the internal `id` it was
2635    /// created with.
2636    fn host_reader_take(&mut self, idx: HostResourceIndex) -> Result<TableId<TransmitHandle>> {
2637        let rep = HostResourceTables::new_host(self)?.host_resource_lift_own(idx)?;
2638        Ok(TableId::new(rep))
2639    }
2640
2641    /// Looks up `idx` within the store, returning the internal `id` it was
2642    /// created with.
2643    fn host_reader_rep(&mut self, idx: HostResourceIndex) -> Result<TableId<TransmitHandle>> {
2644        let rep = HostResourceTables::new_host(self)?.host_resource_rep(idx)?;
2645        Ok(TableId::new(rep))
2646    }
2647}
2648
2649impl<T> StoreContextMut<'_, T> {
2650    fn new_transmit<P: StreamProducer<T>>(
2651        mut self,
2652        kind: TransmitKind,
2653        producer: P,
2654    ) -> Result<HostResourceIndex>
2655    where
2656        P::Item: func::Lower,
2657    {
2658        let token = StoreToken::new(self.as_context_mut());
2659        let state = self.0.concurrent_state_mut()?;
2660        let (_, read) = state.new_transmit(TransmitOrigin::Host {
2661            typecheck: types::typecheck_payload::<P::Item>,
2662        })?;
2663        let producer = Arc::new(LockedState::new((Box::pin(producer), P::Buffer::default())));
2664        let id = state.get_mut(read)?.state;
2665        let mut dropped = false;
2666        let produce = Box::new({
2667            let producer = producer.clone();
2668            move || {
2669                let producer = producer.clone();
2670                async move {
2671                    let mut state = producer.take()?;
2672                    let (mine, buffer) = &mut *state;
2673
2674                    let (result, cancelled) = if buffer.remaining().is_empty() {
2675                        future::poll_fn(|cx| {
2676                            tls::get(|store| {
2677                                let transmit = store.concurrent_state_mut()?.get_mut(id)?;
2678
2679                                let &WriteState::HostReady { cancel, .. } = &transmit.write else {
2680                                    bail_bug!("expected WriteState::HostReady")
2681                                };
2682
2683                                let mut host_written = 0;
2684                                let mut host_buffer =
2685                                    if let ReadState::HostToHost { buffer, .. } = &mut transmit.read {
2686                                        Some(mem::take(buffer))
2687                                    } else {
2688                                        None
2689                                    };
2690
2691                                let poll = mine.as_mut().poll_produce(
2692                                    cx,
2693                                    token.as_context_mut(store),
2694                                    Destination {
2695                                        id,
2696                                        buffer,
2697                                        host_buffer: host_buffer.as_mut().map(|b| {
2698                                            HostBuffer {
2699                                                dst: b,
2700                                                marked_written: &mut host_written,
2701                                            }
2702                                        }),
2703                                        _phantom: PhantomData,
2704                                    },
2705                                    cancel,
2706                                );
2707
2708                                let transmit = store.concurrent_state_mut()?.get_mut(id)?;
2709
2710                                let host_offset = if let (
2711                                    Some(host_buffer),
2712                                    ReadState::HostToHost { buffer, limit, .. },
2713                                ) = (host_buffer, &mut transmit.read)
2714                                {
2715                                    *limit = host_written;
2716                                    *buffer = host_buffer;
2717                                    *limit
2718                                } else {
2719                                    0
2720                                };
2721
2722                                {
2723                                    let WriteState::HostReady {
2724                                        guest_offset,
2725                                        cancel,
2726                                        cancel_waker,
2727                                        ..
2728                                    } = &mut transmit.write
2729                                    else {
2730                                        bail_bug!("expected WriteState::HostReady")
2731                                    };
2732
2733                                    if poll.is_pending() {
2734                                        if !buffer.remaining().is_empty()
2735                                            || *guest_offset > 0
2736                                            || host_offset > 0
2737                                        {
2738                                            bail!(
2739                                                "StreamProducer::poll_produce returned Poll::Pending \
2740                                                 after producing at least one item"
2741                                            )
2742                                        }
2743                                        *cancel_waker = Some(cx.waker().clone());
2744                                    } else {
2745                                        *cancel_waker = None;
2746                                        *cancel = false;
2747                                    }
2748                                }
2749
2750                                Ok(poll.map(|v| v.map(|result| (result, cancel))))
2751                            })?
2752                        })
2753                            .await?
2754                    } else {
2755                        (StreamResult::Completed, false)
2756                    };
2757
2758                    let (guest_offset, host_offset, count) = tls::get(|store| {
2759                        let transmit = store.concurrent_state_mut()?.get_mut(id)?;
2760                        let (count, host_offset) = match &transmit.read {
2761                            &ReadState::GuestReady { count, .. } => (count.as_u32(), 0),
2762                            &ReadState::HostToHost { limit, .. } => (1, limit),
2763                            _ => bail_bug!("invalid read state"),
2764                        };
2765                        let guest_offset = match &transmit.write {
2766                            &WriteState::HostReady { guest_offset, .. } => guest_offset,
2767                            _ => bail_bug!("invalid write state"),
2768                        };
2769                        Ok((guest_offset, host_offset, count))
2770                    })?;
2771
2772                    let produced =
2773                        !buffer.remaining().is_empty() || guest_offset > 0 || host_offset > 0;
2774
2775                    match result {
2776                        StreamResult::Completed => {
2777                            // Only zero-length reads may complete without
2778                            // producing anything. Note that host-to-host reads
2779                            // always have a nonzero `count`.
2780                            if count > 0 && !produced {
2781                                bail!(
2782                                    "StreamProducer::poll_produce returned StreamResult::Completed \
2783                                     without producing any items"
2784                                );
2785                            }
2786                        }
2787                        StreamResult::Cancelled => {
2788                            if !cancelled {
2789                                bail!(
2790                                    "StreamProducer::poll_produce returned StreamResult::Cancelled \
2791                                     without being given a `finish` parameter value of true"
2792                                );
2793                            }
2794                            if produced {
2795                                bail!(
2796                                    "StreamProducer::poll_produce returned StreamResult::Cancelled \
2797                                     after producing at least one item"
2798                                );
2799                            }
2800                        }
2801                        StreamResult::Dropped => {
2802                            dropped = true;
2803                        }
2804                    }
2805
2806                    let write_buffer = !buffer.remaining().is_empty() || host_offset > 0;
2807
2808                    drop(state);
2809
2810                    if write_buffer {
2811                        write(token, id, producer.clone(), kind).await?;
2812                    }
2813
2814                    Ok(if dropped {
2815                        if producer.with(|p| p.1.remaining().is_empty())?  {
2816                            StreamResult::Dropped
2817                        } else {
2818                            StreamResult::Completed
2819                        }
2820                    } else {
2821                        result
2822                    })
2823                }
2824                .boxed()
2825            }
2826        });
2827        let try_into = Box::new(move |ty| {
2828            let (mine, buffer) = producer.try_lock().ok()?.take()?;
2829            match P::try_into(mine, ty) {
2830                Ok(value) => Some(value),
2831                Err(mine) => {
2832                    *producer.try_lock().ok()? = Some((mine, buffer));
2833                    None
2834                }
2835            }
2836        });
2837        state.get_mut(id)?.write = WriteState::HostReady {
2838            produce,
2839            try_into,
2840            guest_offset: ItemCount::ZERO,
2841            cancel: false,
2842            cancel_waker: None,
2843        };
2844        self.0.host_reader_insert(read)
2845    }
2846
2847    fn set_consumer<C: StreamConsumer<T>>(
2848        mut self,
2849        idx: HostResourceIndex,
2850        kind: TransmitKind,
2851        consumer: C,
2852    ) -> Result<()> {
2853        let token = StoreToken::new(self.as_context_mut());
2854        let id = self.0.host_reader_take(idx)?;
2855        let state = self.0.concurrent_state_mut()?;
2856        let id = state.get_mut(id)?.state;
2857        let transmit = state.get_mut(id)?;
2858        let consumer = Arc::new(LockedState::new(Box::pin(consumer)));
2859        let consume_with_buffer = {
2860            let consumer = consumer.clone();
2861            async move |mut host_buffer: Option<&mut dyn WriteBuffer<C::Item>>| {
2862                let mut mine = consumer.take()?;
2863
2864                let host_buffer_remaining_before =
2865                    host_buffer.as_deref_mut().map(|v| v.remaining().len());
2866
2867                let (result, cancelled) = future::poll_fn(|cx| {
2868                    tls::get(|store| {
2869                        let cancel = match &store.concurrent_state_mut()?.get_mut(id)?.read {
2870                            &ReadState::HostReady { cancel, .. } => cancel,
2871                            ReadState::Open => false,
2872                            _ => bail_bug!("unexpected read state"),
2873                        };
2874
2875                        let poll = mine.as_mut().poll_consume(
2876                            cx,
2877                            token.as_context_mut(store),
2878                            Source {
2879                                id,
2880                                host_buffer: host_buffer.as_deref_mut(),
2881                            },
2882                            cancel,
2883                        );
2884
2885                        if let ReadState::HostReady {
2886                            cancel_waker,
2887                            cancel,
2888                            ..
2889                        } = &mut store.concurrent_state_mut()?.get_mut(id)?.read
2890                        {
2891                            if poll.is_pending() {
2892                                *cancel_waker = Some(cx.waker().clone());
2893                            } else {
2894                                *cancel_waker = None;
2895                                *cancel = false;
2896                            }
2897                        }
2898
2899                        Ok(poll.map(|v| v.map(|result| (result, cancel))))
2900                    })?
2901                })
2902                .await?;
2903
2904                let (guest_offset, count) = tls::get(|store| {
2905                    let transmit = store.concurrent_state_mut()?.get_mut(id)?;
2906                    Ok((
2907                        match &transmit.read {
2908                            &ReadState::HostReady { guest_offset, .. } => guest_offset,
2909                            ReadState::Open => ItemCount::ZERO,
2910                            _ => bail_bug!("invalid read state"),
2911                        },
2912                        match &transmit.write {
2913                            WriteState::GuestReady { count, .. } => count.as_usize(),
2914                            WriteState::HostReady { .. } => match host_buffer_remaining_before {
2915                                Some(n) => n,
2916                                None => bail_bug!("host_buffer_remaining_before should be set"),
2917                            },
2918                            _ => bail_bug!("invalid write state"),
2919                        },
2920                    ))
2921                })?;
2922
2923                match result {
2924                    StreamResult::Completed => {
2925                        // Progress for a host writer is tracked by its buffer
2926                        // shrinking, whereas for a guest writer it's tracked
2927                        // in `guest_offset`.
2928                        let consumed = match (host_buffer_remaining_before, host_buffer) {
2929                            (Some(before), Some(after)) => before != after.remaining().len(),
2930                            _ => guest_offset > 0,
2931                        };
2932                        if count > 0 && !consumed {
2933                            bail!(
2934                                "StreamConsumer::poll_consume returned StreamResult::Completed \
2935                                 without consuming any items"
2936                            );
2937                        }
2938
2939                        if let TransmitKind::Future = kind {
2940                            tls::get(|store| {
2941                                store.concurrent_state_mut()?.get_mut(id)?.done = true;
2942                                crate::error::Ok(())
2943                            })?;
2944                        }
2945                    }
2946                    StreamResult::Cancelled => {
2947                        if !cancelled {
2948                            bail!(
2949                                "StreamConsumer::poll_consume returned StreamResult::Cancelled \
2950                                 without being given a `finish` parameter value of true"
2951                            );
2952                        }
2953                    }
2954                    StreamResult::Dropped => {}
2955                }
2956
2957                Ok(result)
2958            }
2959        };
2960        let consume = {
2961            let consume = consume_with_buffer.clone();
2962            Box::new(move || {
2963                let consume = consume.clone();
2964                async move { consume(None).await }.boxed()
2965            })
2966        };
2967
2968        match &transmit.write {
2969            WriteState::Open => {
2970                transmit.read = ReadState::HostReady {
2971                    consume,
2972                    guest_offset: ItemCount::ZERO,
2973                    cancel: false,
2974                    cancel_waker: None,
2975                };
2976            }
2977            &WriteState::GuestReady { .. } => {
2978                // If the guest writer has a pending event, then don't start
2979                // another read event here even if there's some partial data
2980                // available. Wait for that to be delivered and wait for the
2981                // writer to do something else before a read is triggered.
2982                let write_handle = transmit.write_handle;
2983                let write_has_event = Waitable::Transmit(write_handle)
2984                    .common(state)?
2985                    .event
2986                    .is_some();
2987                let future = if write_has_event {
2988                    None
2989                } else {
2990                    Some(consume())
2991                };
2992                state.get_mut(id)?.read = ReadState::HostReady {
2993                    consume,
2994                    guest_offset: ItemCount::ZERO,
2995                    cancel: false,
2996                    cancel_waker: None,
2997                };
2998                if let Some(future) = future {
2999                    self.0.pipe_from_guest(kind, id, future);
3000                }
3001            }
3002            WriteState::HostReady { .. } => {
3003                let WriteState::HostReady { produce, .. } = mem::replace(
3004                    &mut transmit.write,
3005                    WriteState::HostReady {
3006                        produce: Box::new(|| {
3007                            Box::pin(async { bail_bug!("unexpected invocation of `produce`") })
3008                        }),
3009                        try_into: Box::new(|_| None),
3010                        guest_offset: ItemCount::ZERO,
3011                        cancel: false,
3012                        cancel_waker: None,
3013                    },
3014                ) else {
3015                    bail_bug!("expected WriteState::HostReady")
3016                };
3017
3018                transmit.read = ReadState::HostToHost {
3019                    accept: Box::new(move |input| {
3020                        let consume = consume_with_buffer.clone();
3021                        async move { consume(Some(input.get_mut::<C::Item>()?)).await }.boxed()
3022                    }),
3023                    buffer: Vec::new(),
3024                    limit: 0,
3025                };
3026
3027                let future = async move {
3028                    loop {
3029                        if tls::get(|store| {
3030                            crate::error::Ok(matches!(
3031                                store.concurrent_state_mut()?.get_mut(id)?.read,
3032                                ReadState::Dropped
3033                            ))
3034                        })? {
3035                            break Ok(());
3036                        }
3037
3038                        match produce().await? {
3039                            StreamResult::Completed | StreamResult::Cancelled => {}
3040                            StreamResult::Dropped => break Ok(()),
3041                        }
3042
3043                        if let TransmitKind::Future = kind {
3044                            break Ok(());
3045                        }
3046                    }
3047                }
3048                .map(move |result| {
3049                    tls::get(|store| store.concurrent_state_mut()?.delete_transmit(id))?;
3050                    result
3051                });
3052
3053                state.push_future(Box::pin(future));
3054            }
3055            WriteState::Dropped => {
3056                let reader = transmit.read_handle;
3057                self.0.host_drop_reader(reader, kind)?;
3058            }
3059        }
3060        Ok(())
3061    }
3062}
3063
3064async fn write<D: 'static, P: Send + 'static, T: func::Lower + 'static, B: WriteBuffer<T>>(
3065    token: StoreToken<D>,
3066    id: TableId<TransmitState>,
3067    pair: Arc<LockedState<(P, B)>>,
3068    kind: TransmitKind,
3069) -> Result<()> {
3070    let (read, guest_offset) = tls::get(|store| {
3071        let transmit = store.concurrent_state_mut()?.get_mut(id)?;
3072
3073        let guest_offset = if let &WriteState::HostReady { guest_offset, .. } = &transmit.write {
3074            Some(guest_offset)
3075        } else {
3076            None
3077        };
3078
3079        crate::error::Ok((
3080            mem::replace(&mut transmit.read, ReadState::Open),
3081            guest_offset,
3082        ))
3083    })?;
3084
3085    match read {
3086        ReadState::GuestReady {
3087            ty,
3088            flat_abi,
3089            options,
3090            address,
3091            count,
3092            handle,
3093            instance,
3094            caller_instance,
3095        } => {
3096            let guest_offset = match guest_offset {
3097                Some(i) => i,
3098                None => bail_bug!("guest_offset should be present if ready"),
3099            };
3100
3101            if let TransmitKind::Future = kind {
3102                tls::get(|store| {
3103                    store.concurrent_state_mut()?.get_mut(id)?.done = true;
3104                    crate::error::Ok(())
3105                })?;
3106            }
3107
3108            let old_remaining = pair.with(|p| p.1.remaining().len())?;
3109            let accept = {
3110                let pair = pair.clone();
3111                move |mut store: StoreContextMut<D>| {
3112                    let mut state = pair.take()?;
3113                    lower::<T, B, D>(
3114                        store.as_context_mut(),
3115                        instance,
3116                        options,
3117                        ty,
3118                        address + (T::SIZE32 * guest_offset.as_usize()),
3119                        count.as_usize() - guest_offset.as_usize(),
3120                        &mut state.1,
3121                    )?;
3122                    crate::error::Ok(())
3123                }
3124            };
3125
3126            if guest_offset < count {
3127                if T::MAY_REQUIRE_REALLOC {
3128                    // For payloads which may require a realloc call, use a
3129                    // oneshot::channel and background task.  This is
3130                    // necessary because calling the guest while there are
3131                    // host embedder frames on the stack is unsound.
3132                    let (tx, rx) = oneshot::channel();
3133                    tls::get(move |store| {
3134                        store
3135                            .concurrent_state_mut()?
3136                            .push_high_priority(WorkItem::WorkerFunction(AlwaysMut::new(
3137                                Box::new(move |store| {
3138                                    _ = tx.send(accept(token.as_context_mut(store))?);
3139                                    Ok(())
3140                                }),
3141                            )));
3142                        crate::error::Ok(())
3143                    })?;
3144                    match rx.await {
3145                        Ok(r) => r,
3146                        // The worker function was dropped without sending a
3147                        // result, which can happen when lowering trapped. This
3148                        // isn't a bug but this specific handling here might be
3149                        // the last hurray of this task in the event queue, so
3150                        // don't kill anything other than our caller. The
3151                        // original trap should make its way back to the best
3152                        // location naturally.
3153                        Err(oneshot::Canceled) => {
3154                            bail!("stream write cancelled before it could complete")
3155                        }
3156                    }
3157                } else {
3158                    // Optimize flat payloads (i.e. those which do not
3159                    // require calling the guest's realloc function) by
3160                    // lowering directly instead of using a oneshot::channel
3161                    // and background task.
3162                    tls::get(|store| accept(token.as_context_mut(store)))?
3163                }
3164            }
3165
3166            tls::get(|store| {
3167                let count = old_remaining - pair.with(|p| p.1.remaining().len())?;
3168
3169                let transmit = store.concurrent_state_mut()?.get_mut(id)?;
3170
3171                let WriteState::HostReady { guest_offset, .. } = &mut transmit.write else {
3172                    bail_bug!("expected WriteState::HostReady")
3173                };
3174
3175                guest_offset.inc(count)?;
3176
3177                transmit.read = ReadState::GuestReady {
3178                    ty,
3179                    flat_abi,
3180                    options,
3181                    address,
3182                    count: ItemCount::new_usize(count)?,
3183                    handle,
3184                    instance,
3185                    caller_instance,
3186                };
3187
3188                crate::error::Ok(())
3189            })?;
3190
3191            Ok(())
3192        }
3193
3194        ReadState::HostToHost {
3195            accept,
3196            mut buffer,
3197            limit,
3198        } => {
3199            let mut state = StreamResult::Completed;
3200            let mut position = 0;
3201
3202            while !matches!(state, StreamResult::Dropped) && position < limit {
3203                let mut slice_buffer = SliceBuffer::new(buffer, position, limit);
3204                state = accept(&mut UntypedWriteBuffer::new(&mut slice_buffer)).await?;
3205                (buffer, position, _) = slice_buffer.into_parts();
3206            }
3207
3208            {
3209                let mut pair = pair.take()?;
3210                let (_, buffer) = &mut *pair;
3211
3212                while !(matches!(state, StreamResult::Dropped) || buffer.remaining().is_empty()) {
3213                    state = accept(&mut UntypedWriteBuffer::new(buffer)).await?;
3214                }
3215            }
3216
3217            tls::get(|store| {
3218                store.concurrent_state_mut()?.get_mut(id)?.read = match state {
3219                    StreamResult::Dropped => ReadState::Dropped,
3220                    StreamResult::Completed | StreamResult::Cancelled => ReadState::HostToHost {
3221                        accept,
3222                        buffer,
3223                        limit: 0,
3224                    },
3225                };
3226
3227                crate::error::Ok(())
3228            })?;
3229            Ok(())
3230        }
3231
3232        _ => bail_bug!("unexpected read state"),
3233    }
3234}
3235
3236/// Polls a `StreamProducer` or `StreamConsumer` future once, inline, from
3237/// within a guest's `{stream,future}.{read,write}` intrinsic.
3238///
3239/// The producer or consumer is embedder code which doesn't belong to the
3240/// reading or writing guest's task, so it's polled as though from the
3241/// top-level host, just as it is when it's later polled from the event loop.
3242fn poll_inline(
3243    store: &mut dyn VMStore,
3244    future: &mut Pin<Box<dyn Future<Output = Result<StreamResult>> + Send + 'static>>,
3245) -> Result<Poll<Result<StreamResult>>> {
3246    let old_thread = store.set_thread(CurrentThread::None)?;
3247    let poll = tls::set(store, || {
3248        future
3249            .as_mut()
3250            .poll(&mut Context::from_waker(&Waker::noop()))
3251    });
3252    store.set_thread(old_thread)?;
3253    Ok(poll)
3254}
3255
3256impl Instance {
3257    /// Handle a guest-initiated write by delivering the item(s) to the
3258    /// `StreamConsumer` for the specified stream or future.
3259    fn consume(
3260        self,
3261        store: &mut dyn VMStore,
3262        kind: TransmitKind,
3263        transmit_id: TableId<TransmitState>,
3264        consume: PollStream,
3265        guest_offset: ItemCount,
3266        cancel: bool,
3267    ) -> Result<ReturnCode> {
3268        let mut future = consume();
3269        store.concurrent_state_mut()?.get_mut(transmit_id)?.read = ReadState::HostReady {
3270            consume,
3271            guest_offset,
3272            cancel,
3273            cancel_waker: None,
3274        };
3275        let poll = poll_inline(store, &mut future)?;
3276
3277        Ok(match poll {
3278            Poll::Ready(state) => {
3279                let transmit = store.concurrent_state_mut()?.get_mut(transmit_id)?;
3280                let code = settle_host_read(transmit, kind, state?)?;
3281                transmit.write = WriteState::Open;
3282                code
3283            }
3284            Poll::Pending => {
3285                store.pipe_from_guest(kind, transmit_id, future);
3286                ReturnCode::Blocked
3287            }
3288        })
3289    }
3290
3291    /// Handle a guest-initiated read by polling the `StreamProducer` for the
3292    /// specified stream or future for items.
3293    fn produce(
3294        self,
3295        store: &mut dyn VMStore,
3296        kind: TransmitKind,
3297        transmit_id: TableId<TransmitState>,
3298        produce: PollStream,
3299        try_into: TryInto,
3300        guest_offset: ItemCount,
3301        cancel: bool,
3302    ) -> Result<ReturnCode> {
3303        let mut future = produce();
3304        store.concurrent_state_mut()?.get_mut(transmit_id)?.write = WriteState::HostReady {
3305            produce,
3306            try_into,
3307            guest_offset,
3308            cancel,
3309            cancel_waker: None,
3310        };
3311        let poll = poll_inline(store, &mut future)?;
3312
3313        Ok(match poll {
3314            Poll::Ready(state) => {
3315                let transmit = store.concurrent_state_mut()?.get_mut(transmit_id)?;
3316                let code = settle_host_write(transmit, kind, state?)?;
3317                transmit.read = ReadState::Open;
3318                code
3319            }
3320            Poll::Pending => {
3321                store.pipe_to_guest(kind, transmit_id, future);
3322                ReturnCode::Blocked
3323            }
3324        })
3325    }
3326
3327    /// Drop the writable end of the specified stream or future from the guest.
3328    pub(super) fn guest_drop_writable(
3329        self,
3330        store: &mut StoreOpaque,
3331        ty: TransmitIndex,
3332        writer: u32,
3333    ) -> Result<()> {
3334        let (instance, counter) = self.id().get_mut_and_handle_counter(store);
3335        let table = instance.table_for_transmit(ty);
3336        let (transmit_rep, is_done) = match ty {
3337            TransmitIndex::Future(ty) => table.future_remove_writable(counter, ty, writer)?,
3338            TransmitIndex::Stream(ty) => {
3339                (table.stream_remove_writable(counter, ty, writer)?, false)
3340            }
3341        };
3342
3343        let id = TableId::<TransmitHandle>::new(transmit_rep);
3344        log::trace!("guest_drop_writable: drop writer {id:?}");
3345        match ty {
3346            TransmitIndex::Stream(_) => store.host_drop_writer(id, None),
3347            TransmitIndex::Future(_) => store.host_drop_writer(
3348                id,
3349                if is_done {
3350                    None
3351                } else {
3352                    Some(|| {
3353                        Err(format_err!(
3354                            "cannot drop future write end without first writing a value"
3355                        ))
3356                    })
3357                },
3358            ),
3359        }
3360    }
3361
3362    /// Copy `count` items from `read_address` to `write_address` for the
3363    /// specified stream or future.
3364    fn copy<T: 'static>(
3365        store: StoreContextMut<T>,
3366        flat_abi: Option<FlatAbi>,
3367        write_runtime_instance: RuntimeInstance,
3368        write_ty: TransmitIndex,
3369        write_options: OptionsIndex,
3370        write_address: usize,
3371        read_runtime_instance: RuntimeInstance,
3372        read_ty: TransmitIndex,
3373        read_options: OptionsIndex,
3374        read_address: usize,
3375        count: ItemCount,
3376        rep: u32,
3377    ) -> Result<()> {
3378        let write_instance = Instance::from_runtime_instance(store.0, write_runtime_instance);
3379        let read_instance = Instance::from_runtime_instance(store.0, read_runtime_instance);
3380        let (write_component, store) = write_instance.component_and_store_mut(store.0);
3381        let (read_component, mut store) = read_instance.component_and_store_mut(store);
3382        let write_types = write_component.types();
3383        let read_types = read_component.types();
3384        let count = count.as_usize();
3385
3386        // Validate `write_ty` w.r.t. `write_address` to ensure it's properly
3387        // aligned and in-bounds.
3388        //
3389        // FIXME(#4311): this and the rest of this function use the 32-bit
3390        // layout of payloads, which needs memory64 handling.
3391        let write_payload_ty = write_ty.payload(write_types);
3392        let write_abi = match write_payload_ty {
3393            Some(ty) => write_types.canonical_abi(ty),
3394            None => &CanonicalAbiInfo::ZERO,
3395        };
3396        let write_length_in_bytes = match flat_abi {
3397            Some(abi) => usize::try_from(abi.size)? * count,
3398            None => usize::try_from(write_abi.size32)? * count,
3399        };
3400        if write_length_in_bytes > 0 {
3401            if write_address % usize::try_from(write_abi.align32)? != 0 {
3402                bail!("write pointer not aligned");
3403            }
3404            write_instance
3405                .options_memory(store, write_options)
3406                .get(write_address..)
3407                .and_then(|b| b.get(..write_length_in_bytes))
3408                .ok_or_else(|| crate::format_err!("write pointer out of bounds"))?;
3409        }
3410
3411        let read_payload_ty = read_ty.payload(read_types);
3412        let read_abi = match read_payload_ty {
3413            Some(ty) => read_types.canonical_abi(ty),
3414            None => &CanonicalAbiInfo::ZERO,
3415        };
3416        let read_length_in_bytes = match flat_abi {
3417            Some(abi) => usize::try_from(abi.size)? * count,
3418            None => usize::try_from(read_abi.size32)? * count,
3419        };
3420        if read_length_in_bytes > 0 {
3421            if read_address % usize::try_from(read_abi.align32)? != 0 {
3422                bail!("read pointer not aligned");
3423            }
3424            read_instance
3425                .options_memory(store, read_options)
3426                .get(read_address..)
3427                .and_then(|b| b.get(..read_length_in_bytes))
3428                .ok_or_else(|| crate::format_err!("read pointer out of bounds"))?;
3429        }
3430
3431        if write_runtime_instance == read_runtime_instance
3432            && !allow_intra_component_read_write(write_payload_ty)
3433        {
3434            bail!(
3435                "cannot read from and write to intra-component future/stream with non-numeric payload"
3436            )
3437        }
3438
3439        match (write_ty, read_ty) {
3440            (TransmitIndex::Future(_), TransmitIndex::Future(_)) => {
3441                if count != 1 {
3442                    bail_bug!("futures can only send 1 item");
3443                }
3444
3445                let val = write_payload_ty
3446                    .map(|ty| {
3447                        let lift = &mut LiftContext::new(store, write_options, write_instance)?;
3448                        let bytes = &lift.memory()[write_address..][..write_length_in_bytes];
3449                        Val::load(lift, *ty, bytes)
3450                    })
3451                    .transpose()?;
3452
3453                if let Some(val) = val {
3454                    let lower =
3455                        &mut LowerContext::new(store.as_context_mut(), read_options, read_instance);
3456                    let ptr = func::validate_inbounds_dynamic(
3457                        read_abi,
3458                        lower.as_slice_mut(),
3459                        &ValRaw::u32(read_address.try_into()?),
3460                    )?;
3461                    let ty = match read_payload_ty {
3462                        Some(ty) => ty,
3463                        None => bail_bug!("expected read payload type to be present"),
3464                    };
3465                    val.store(lower, *ty, ptr)?;
3466                }
3467            }
3468            (TransmitIndex::Stream(_), TransmitIndex::Stream(_)) => {
3469                if write_length_in_bytes == 0 {
3470                    return Ok(());
3471                }
3472                let write_payload_ty = match write_payload_ty {
3473                    Some(ty) => ty,
3474                    None => bail_bug!("expected write payload type to be present"),
3475                };
3476                let read_payload_ty = match read_payload_ty {
3477                    Some(ty) => ty,
3478                    None => bail_bug!("expected read payload type to be present"),
3479                };
3480                if flat_abi.is_some() {
3481                    // Fast path memcpy for "flat" (i.e. no pointers or handles) payloads:
3482                    let store_opaque = store.store_opaque_mut();
3483
3484                    assert_eq!(read_length_in_bytes, write_length_in_bytes);
3485
3486                    if read_instance
3487                        .options_memory(store_opaque, read_options)
3488                        .as_ptr()
3489                        == write_instance
3490                            .options_memory(store_opaque, write_options)
3491                            .as_ptr()
3492                    {
3493                        let memory = read_instance.options_memory_mut(store_opaque, read_options);
3494                        memory.copy_within(
3495                            write_address..write_address + write_length_in_bytes,
3496                            read_address,
3497                        );
3498                    } else {
3499                        let src = write_instance.options_memory(store_opaque, write_options)
3500                            [write_address..][..write_length_in_bytes]
3501                            .as_ptr();
3502                        let dst = read_instance.options_memory_mut(store_opaque, read_options)
3503                            [read_address..][..read_length_in_bytes]
3504                            .as_mut_ptr();
3505
3506                        // SAFETY: Both `src` and `dst` have been validated
3507                        // above to be valid pointers as they're derived from
3508                        // slices that have the desired length with the desired
3509                        // read/write permission. The `unsafe` bit here is that
3510                        // the memories are disjoint (different base pointers)
3511                        // and there's no easy way to borrow both
3512                        // simultaneously from the store. Different memories
3513                        // are guaranteed to be disjoint, however, so the
3514                        // `unsafe` here should be ok.
3515                        unsafe {
3516                            src.copy_to_nonoverlapping(dst, write_length_in_bytes);
3517                        }
3518                    }
3519                } else {
3520                    let store_opaque = store.store_opaque_mut();
3521                    let lift = &mut LiftContext::new(store_opaque, write_options, write_instance)?;
3522                    let bytes = &lift.memory()[write_address..][..write_length_in_bytes];
3523                    lift.consume_fuel_array(count, size_of::<Val>())?;
3524
3525                    let values = (0..count)
3526                        .map(|index| {
3527                            let size = usize::try_from(write_abi.size32)?;
3528                            Val::load(lift, *write_payload_ty, &bytes[(index * size)..][..size])
3529                        })
3530                        .collect::<Result<Vec<_>>>()?;
3531
3532                    let id = TableId::<TransmitHandle>::new(rep);
3533                    log::trace!("copy values {values:?} for {id:?}");
3534
3535                    let lower =
3536                        &mut LowerContext::new(store.as_context_mut(), read_options, read_instance);
3537                    let mut ptr = read_address;
3538                    for value in values {
3539                        value.store(lower, *read_payload_ty, ptr)?;
3540                        ptr += usize::try_from(read_abi.size32)?;
3541                    }
3542                }
3543            }
3544            _ => bail_bug!("mismatched transmit types in copy"),
3545        }
3546
3547        Ok(())
3548    }
3549
3550    fn check_bounds(
3551        self,
3552        store: &StoreOpaque,
3553        options: OptionsIndex,
3554        ty: TransmitIndex,
3555        address: usize,
3556        count: usize,
3557    ) -> Result<()> {
3558        let types = self.id().get(store).component().types();
3559        let payload = match ty {
3560            TransmitIndex::Future(ty) => types[types[ty].ty].payload,
3561            TransmitIndex::Stream(ty) => types[types[ty].ty].payload,
3562        };
3563        let Some(payload) = payload else {
3564            return Ok(());
3565        };
3566        // As-written the spec currently exempts zero-length reads/writes from
3567        // bounds/alignment checks.
3568        if count == 0 {
3569            return Ok(());
3570        }
3571        let abi = types.canonical_abi(&payload);
3572        // FIXME(#4311): needs memory64 handling
3573        let size = usize::try_from(abi.size32)?;
3574        let align = usize::try_from(abi.align32)?;
3575        if address % align != 0 {
3576            bail!("buffer pointer not aligned");
3577        }
3578        self.options_memory(store, options)
3579            .get(address..)
3580            .and_then(|b| b.get(..size.checked_mul(count)?))
3581            .map(drop)
3582            .ok_or_else(|| crate::format_err!("buffer pointer out of bounds of memory"))
3583    }
3584
3585    /// Write to the specified stream or future from the guest.
3586    pub(super) fn guest_write<T: 'static>(
3587        self,
3588        mut store: StoreContextMut<T>,
3589        caller: RuntimeComponentInstanceIndex,
3590        ty: TransmitIndex,
3591        options: OptionsIndex,
3592        flat_abi: Option<FlatAbi>,
3593        handle: u32,
3594        address: u32,
3595        count: u32,
3596    ) -> Result<ReturnCode> {
3597        let count = ItemCount::new(count)?;
3598
3599        let address = usize::try_from(address)?;
3600        self.check_bounds(store.0, options, ty, address, count.as_usize())?;
3601        let async_ = self.options(store.0, options).async_;
3602        let (rep, state) = self.id().get_mut(store.0).get_mut_by_index(ty, handle)?;
3603        let TransmitLocalState::Write { done } = *state else {
3604            bail!(Trap::ConcurrentFutureStreamOp);
3605        };
3606
3607        *state = TransmitLocalState::Busy;
3608        let transmit_handle = TableId::<TransmitHandle>::new(rep);
3609        let concurrent_state = store.0.concurrent_state_mut()?;
3610        let transmit_id = concurrent_state.get_mut(transmit_handle)?.state;
3611        let transmit = concurrent_state.get_mut(transmit_id)?;
3612        log::trace!(
3613            "guest_write {count} to {transmit_handle:?} (handle {handle}; state {transmit_id:?}); {:?}",
3614            transmit.read
3615        );
3616
3617        if done || transmit.done {
3618            bail!(match ty {
3619                TransmitIndex::Future(_) => Trap::WriteToDroppedFuture,
3620                TransmitIndex::Stream(_) => Trap::WriteToDroppedStream,
3621            });
3622        }
3623
3624        if !async_ {
3625            Waitable::Transmit(transmit_handle).trap_if_in_waitable_set(concurrent_state)?;
3626        }
3627        let transmit = concurrent_state.get_mut(transmit_id)?;
3628
3629        let new_state = if let ReadState::Dropped = &transmit.read {
3630            ReadState::Dropped
3631        } else {
3632            ReadState::Open
3633        };
3634
3635        let set_guest_ready = |me: &mut ConcurrentState| {
3636            let transmit = me.get_mut(transmit_id)?;
3637            if !matches!(&transmit.write, WriteState::Open) {
3638                bail_bug!("expected `WriteState::Open`; got `{:?}`", transmit.write);
3639            }
3640            transmit.write = WriteState::GuestReady {
3641                instance: self,
3642                caller,
3643                ty,
3644                flat_abi,
3645                options,
3646                address,
3647                count,
3648                handle,
3649            };
3650            Ok::<_, crate::Error>(())
3651        };
3652
3653        let mut result = match mem::replace(&mut transmit.read, new_state) {
3654            ReadState::GuestReady {
3655                ty: read_ty,
3656                flat_abi: read_flat_abi,
3657                options: read_options,
3658                address: read_address,
3659                count: read_count,
3660                handle: read_handle,
3661                instance: read_instance,
3662                caller_instance: read_caller_instance,
3663            } => {
3664                if flat_abi != read_flat_abi {
3665                    bail_bug!("expected flat ABI calculations to be the same");
3666                }
3667
3668                if let TransmitIndex::Future(_) = ty {
3669                    transmit.done = true;
3670                }
3671
3672                // Note that zero-length reads and writes are handling specially
3673                // by the spec to allow each end to signal readiness to the
3674                // other.  Quoting the spec:
3675                //
3676                // ```
3677                // The meaning of a read or write when the length is 0 is that
3678                // the caller is querying the "readiness" of the other
3679                // side. When a 0-length read/write rendezvous with a
3680                // non-0-length read/write, only the 0-length read/write
3681                // completes; the non-0-length read/write is kept pending (and
3682                // ready for a subsequent rendezvous).
3683                //
3684                // In the corner case where a 0-length read and write
3685                // rendezvous, only the writer is notified of readiness. To
3686                // avoid livelock, the Canonical ABI requires that a writer must
3687                // (eventually) follow a completed 0-length write with a
3688                // non-0-length write that is allowed to block (allowing the
3689                // reader end to run and rendezvous with its own non-0-length
3690                // read).
3691                // ```
3692
3693                let write_count = count;
3694                let write_complete = count == 0 || read_count > 0;
3695                let read_complete = count > 0;
3696                let read_buffer_remaining = count < read_count;
3697
3698                let read_handle_rep = transmit.read_handle.rep();
3699
3700                let count = count.min(read_count);
3701
3702                Instance::copy(
3703                    store.as_context_mut(),
3704                    flat_abi,
3705                    self.runtime_instance(caller),
3706                    ty,
3707                    options,
3708                    address,
3709                    read_instance.runtime_instance(read_caller_instance),
3710                    read_ty,
3711                    read_options,
3712                    read_address,
3713                    count,
3714                    rep,
3715                )?;
3716
3717                let instance = read_instance.id().get(store.0);
3718                let types = instance.component().types();
3719                // FIXME(#4311): needs memory64 handling
3720                let item_size = match read_ty.payload(types) {
3721                    Some(ty) => usize::try_from(types.canonical_abi(ty).size32)?,
3722                    None => 0,
3723                };
3724                let concurrent_state = store.0.concurrent_state_mut()?;
3725                if read_complete {
3726                    let total = if let Some(Event::StreamRead {
3727                        code: ReturnCode::Completed(old_total),
3728                        ..
3729                    }) = concurrent_state.take_event(read_handle_rep)?
3730                    {
3731                        count.add(old_total)?
3732                    } else {
3733                        count
3734                    };
3735
3736                    let code = ReturnCode::completed(ty.kind(), total);
3737
3738                    concurrent_state.send_read_result(read_ty, transmit_id, read_handle, code)?;
3739                }
3740
3741                // If the reader still has buffer remaining, or if this was a
3742                // zero-length rendezvous, then restore the state of the reader
3743                // back to what it was when we found it. Note that for the
3744                // zero-length rendezvous case this specifically won't execute
3745                // the `read_complete` logic above, which is intentional, as the
3746                // reader remains blocked.
3747                if read_buffer_remaining || (write_count == 0 && read_count == 0) {
3748                    let transmit = concurrent_state.get_mut(transmit_id)?;
3749                    transmit.read = ReadState::GuestReady {
3750                        ty: read_ty,
3751                        flat_abi: read_flat_abi,
3752                        options: read_options,
3753                        address: read_address + (count.as_usize() * item_size),
3754                        count: read_count.sub(count)?,
3755                        handle: read_handle,
3756                        instance: read_instance,
3757                        caller_instance: read_caller_instance,
3758                    };
3759                }
3760
3761                if write_complete {
3762                    ReturnCode::completed(ty.kind(), count)
3763                } else {
3764                    set_guest_ready(concurrent_state)?;
3765                    ReturnCode::Blocked
3766                }
3767            }
3768
3769            ReadState::HostReady {
3770                consume,
3771                guest_offset,
3772                cancel,
3773                cancel_waker,
3774            } => {
3775                if cancel_waker.is_some() {
3776                    bail_bug!("expected cancel_waker to be none");
3777                }
3778                if cancel {
3779                    bail_bug!("expected cancel to be false");
3780                }
3781                if guest_offset != 0 {
3782                    bail_bug!("expected guest_offset to be 0");
3783                }
3784
3785                set_guest_ready(concurrent_state)?;
3786                self.consume(
3787                    store.0,
3788                    ty.kind(),
3789                    transmit_id,
3790                    consume,
3791                    ItemCount::ZERO,
3792                    false,
3793                )?
3794            }
3795
3796            ReadState::HostToHost { .. } => bail_bug!("unexpected HostToHost"),
3797
3798            ReadState::Open => {
3799                set_guest_ready(concurrent_state)?;
3800                ReturnCode::Blocked
3801            }
3802
3803            ReadState::Dropped => {
3804                if let TransmitIndex::Future(_) = ty {
3805                    transmit.done = true;
3806                }
3807
3808                match Waitable::Transmit(transmit_handle).take_event(concurrent_state)? {
3809                    Some(
3810                        Event::StreamWrite {
3811                            code: ReturnCode::Dropped(ItemCount::ZERO),
3812                            ..
3813                        }
3814                        | Event::FutureWrite {
3815                            code: ReturnCode::Dropped(ItemCount::ZERO),
3816                            ..
3817                        },
3818                    ) => {}
3819                    event => bail_bug!("expected pending dropped event for writer; got {event:?}"),
3820                }
3821
3822                ReturnCode::Dropped(ItemCount::ZERO)
3823            }
3824        };
3825
3826        if result == ReturnCode::Blocked && !async_ {
3827            result = self.wait_for_write(store.0, caller, transmit_handle)?;
3828        }
3829
3830        if result != ReturnCode::Blocked {
3831            *self.id().get_mut(store.0).get_mut_by_index(ty, handle)?.1 =
3832                TransmitLocalState::Write {
3833                    done: matches!(result, ReturnCode::Dropped(_)),
3834                };
3835        }
3836
3837        log::trace!(
3838            "guest_write result for {transmit_handle:?} (handle {handle}; state {transmit_id:?}): {result:?}",
3839        );
3840
3841        Ok(result)
3842    }
3843
3844    /// Read from the specified stream or future from the guest.
3845    pub(super) fn guest_read<T: 'static>(
3846        self,
3847        mut store: StoreContextMut<T>,
3848        caller_instance: RuntimeComponentInstanceIndex,
3849        ty: TransmitIndex,
3850        options: OptionsIndex,
3851        flat_abi: Option<FlatAbi>,
3852        handle: u32,
3853        address: u32,
3854        count: u32,
3855    ) -> Result<ReturnCode> {
3856        let count = ItemCount::new(count)?;
3857
3858        let address = usize::try_from(address)?;
3859        self.check_bounds(store.0, options, ty, address, count.as_usize())?;
3860        let async_ = self.options(store.0, options).async_;
3861        let (rep, state) = self.id().get_mut(store.0).get_mut_by_index(ty, handle)?;
3862        let TransmitLocalState::Read { done } = *state else {
3863            bail!(Trap::ConcurrentFutureStreamOp);
3864        };
3865
3866        *state = TransmitLocalState::Busy;
3867        let transmit_handle = TableId::<TransmitHandle>::new(rep);
3868        let concurrent_state = store.0.concurrent_state_mut()?;
3869        let transmit_id = concurrent_state.get_mut(transmit_handle)?.state;
3870        let transmit = concurrent_state.get_mut(transmit_id)?;
3871        log::trace!(
3872            "guest_read {count} from {transmit_handle:?} (handle {handle}; state {transmit_id:?}); {:?}",
3873            transmit.write
3874        );
3875
3876        if done || transmit.done {
3877            match ty {
3878                TransmitIndex::Future(_) => {
3879                    bail!("cannot read from future after previous read succeeded")
3880                }
3881                TransmitIndex::Stream(_) => bail!(Trap::ReadFromDroppedStream),
3882            }
3883        }
3884
3885        if !async_ {
3886            Waitable::Transmit(transmit_handle).trap_if_in_waitable_set(concurrent_state)?;
3887        }
3888        let transmit = concurrent_state.get_mut(transmit_id)?;
3889
3890        let new_state = if let WriteState::Dropped = &transmit.write {
3891            WriteState::Dropped
3892        } else {
3893            WriteState::Open
3894        };
3895
3896        let set_guest_ready = |me: &mut ConcurrentState| {
3897            let transmit = me.get_mut(transmit_id)?;
3898            if !matches!(&transmit.read, ReadState::Open) {
3899                bail_bug!("expected `ReadState::Open`; got `{:?}`", transmit.read);
3900            }
3901            transmit.read = ReadState::GuestReady {
3902                ty,
3903                flat_abi,
3904                options,
3905                address,
3906                count,
3907                handle,
3908                instance: self,
3909                caller_instance,
3910            };
3911            Ok::<_, crate::Error>(())
3912        };
3913
3914        let mut result = match mem::replace(&mut transmit.write, new_state) {
3915            WriteState::GuestReady {
3916                instance: write_instance,
3917                ty: write_ty,
3918                flat_abi: write_flat_abi,
3919                options: write_options,
3920                address: write_address,
3921                count: write_count,
3922                handle: write_handle,
3923                caller: write_caller,
3924            } => {
3925                if flat_abi != write_flat_abi {
3926                    bail_bug!("expected flat ABI calculations to be the same");
3927                }
3928
3929                if let TransmitIndex::Future(_) = ty {
3930                    transmit.done = true;
3931                }
3932
3933                let write_handle_rep = transmit.write_handle.rep();
3934
3935                // See the comment in `guest_write` for the
3936                // `ReadState::GuestReady` case concerning zero-length reads and
3937                // writes.
3938
3939                let write_complete = write_count == 0 || count > 0;
3940                let read_complete = write_count > 0;
3941                let write_buffer_remaining = count < write_count;
3942
3943                let count = count.min(write_count);
3944
3945                Instance::copy(
3946                    store.as_context_mut(),
3947                    flat_abi,
3948                    write_instance.runtime_instance(write_caller),
3949                    write_ty,
3950                    write_options,
3951                    write_address,
3952                    self.runtime_instance(caller_instance),
3953                    ty,
3954                    options,
3955                    address,
3956                    count,
3957                    rep,
3958                )?;
3959
3960                let instance = write_instance.id().get(store.0);
3961                let types = instance.component().types();
3962                // FIXME(#4311): needs memory64 handling
3963                let item_size = match write_ty.payload(types) {
3964                    Some(ty) => usize::try_from(types.canonical_abi(ty).size32)?,
3965                    None => 0,
3966                };
3967                let concurrent_state = store.0.concurrent_state_mut()?;
3968
3969                if write_complete {
3970                    let total = if let Some(Event::StreamWrite {
3971                        code: ReturnCode::Completed(old_total),
3972                        ..
3973                    }) = concurrent_state.take_event(write_handle_rep)?
3974                    {
3975                        count.add(old_total)?
3976                    } else {
3977                        count
3978                    };
3979
3980                    let code = ReturnCode::completed(ty.kind(), total);
3981
3982                    concurrent_state.send_write_result(
3983                        write_ty,
3984                        transmit_id,
3985                        write_handle,
3986                        code,
3987                    )?;
3988                }
3989
3990                if write_buffer_remaining {
3991                    let transmit = concurrent_state.get_mut(transmit_id)?;
3992                    transmit.write = WriteState::GuestReady {
3993                        instance: write_instance,
3994                        caller: write_caller,
3995                        ty: write_ty,
3996                        flat_abi: write_flat_abi,
3997                        options: write_options,
3998                        address: write_address + (count.as_usize() * item_size),
3999                        count: write_count.sub(count)?,
4000                        handle: write_handle,
4001                    };
4002                }
4003
4004                if read_complete {
4005                    ReturnCode::completed(ty.kind(), count)
4006                } else {
4007                    set_guest_ready(concurrent_state)?;
4008                    ReturnCode::Blocked
4009                }
4010            }
4011
4012            WriteState::HostReady {
4013                produce,
4014                try_into,
4015                guest_offset,
4016                cancel,
4017                cancel_waker,
4018            } => {
4019                if cancel_waker.is_some() {
4020                    bail_bug!("expected cancel_waker to be none");
4021                }
4022                if cancel {
4023                    bail_bug!("expected cancel to be false");
4024                }
4025                if guest_offset != 0 {
4026                    bail_bug!("expected guest_offset to be 0");
4027                }
4028
4029                set_guest_ready(concurrent_state)?;
4030
4031                let code = self.produce(
4032                    store.0,
4033                    ty.kind(),
4034                    transmit_id,
4035                    produce,
4036                    try_into,
4037                    ItemCount::ZERO,
4038                    false,
4039                )?;
4040
4041                if let (TransmitIndex::Future(_), ReturnCode::Completed(_)) = (ty, code) {
4042                    store.0.concurrent_state_mut()?.get_mut(transmit_id)?.done = true;
4043                }
4044
4045                code
4046            }
4047
4048            WriteState::Open => {
4049                set_guest_ready(concurrent_state)?;
4050                ReturnCode::Blocked
4051            }
4052
4053            WriteState::Dropped => {
4054                if let TransmitIndex::Future(_) = ty {
4055                    bail_bug!(
4056                        "should not be possible to read from a future whose write end was dropped"
4057                    );
4058                }
4059
4060                match Waitable::Transmit(transmit_handle).take_event(concurrent_state)? {
4061                    Some(Event::StreamRead {
4062                        code: ReturnCode::Dropped(ItemCount::ZERO),
4063                        ..
4064                    }) => {}
4065                    event => bail_bug!("expected pending dropped event for reader; got {event:?}"),
4066                }
4067
4068                ReturnCode::Dropped(ItemCount::ZERO)
4069            }
4070        };
4071
4072        if result == ReturnCode::Blocked && !async_ {
4073            result = self.wait_for_read(store.0, caller_instance, transmit_handle)?;
4074        }
4075
4076        if result != ReturnCode::Blocked {
4077            *self.id().get_mut(store.0).get_mut_by_index(ty, handle)?.1 =
4078                TransmitLocalState::Read {
4079                    done: matches!(
4080                        (result, ty),
4081                        (ReturnCode::Dropped(_), TransmitIndex::Stream(_))
4082                    ),
4083                };
4084        }
4085
4086        log::trace!(
4087            "guest_read result for {transmit_handle:?} (handle {handle}; state {transmit_id:?}): {result:?}",
4088        );
4089
4090        Ok(result)
4091    }
4092
4093    fn wait_for_write(
4094        self,
4095        store: &mut StoreOpaque,
4096        caller: RuntimeComponentInstanceIndex,
4097        handle: TableId<TransmitHandle>,
4098    ) -> Result<ReturnCode> {
4099        let waitable = Waitable::Transmit(handle);
4100        store.wait_for_event(self.runtime_instance(caller), waitable)?;
4101        let event = waitable.take_event(store.concurrent_state_mut()?)?;
4102        if let Some(event @ (Event::StreamWrite { code, .. } | Event::FutureWrite { code, .. })) =
4103            event
4104        {
4105            waitable.on_delivery(store, self, event)?;
4106            Ok(code)
4107        } else {
4108            bail_bug!("expected either a stream or future write event")
4109        }
4110    }
4111
4112    /// Cancel a pending stream or future write.
4113    fn cancel_write(
4114        self,
4115        store: &mut StoreOpaque,
4116        caller: RuntimeComponentInstanceIndex,
4117        transmit_id: TableId<TransmitState>,
4118        async_: bool,
4119    ) -> Result<ReturnCode> {
4120        let state = store.concurrent_state_mut()?;
4121        let transmit = state.get_mut(transmit_id)?;
4122        log::trace!(
4123            "host_cancel_write state {transmit_id:?}; write state {:?} read state {:?}",
4124            transmit.read,
4125            transmit.write
4126        );
4127        let waitable = Waitable::Transmit(transmit.write_handle);
4128
4129        waitable.trap_if_has_sync_waiter(state)?;
4130        if !async_ {
4131            waitable.trap_if_in_waitable_set(state)?;
4132        }
4133
4134        let code = if let Some(event) = waitable.take_event(state)? {
4135            let (Event::FutureWrite { code, .. } | Event::StreamWrite { code, .. }) = event else {
4136                bail_bug!("expected either a stream or future write event")
4137            };
4138            waitable.on_delivery(store, self, event)?;
4139            match (code, event) {
4140                (ReturnCode::Completed(count), Event::StreamWrite { .. }) => {
4141                    ReturnCode::Cancelled(count)
4142                }
4143                (ReturnCode::Dropped(_) | ReturnCode::Completed(_), _) => code,
4144                _ => bail_bug!("unexpected code/event combo"),
4145            }
4146        } else if let ReadState::HostReady {
4147            cancel,
4148            cancel_waker,
4149            ..
4150        } = &mut state.get_mut(transmit_id)?.read
4151        {
4152            *cancel = true;
4153            if let Some(waker) = cancel_waker.take() {
4154                waker.wake();
4155            }
4156
4157            if async_ {
4158                ReturnCode::Blocked
4159            } else {
4160                let handle = store
4161                    .concurrent_state_mut()?
4162                    .get_mut(transmit_id)?
4163                    .write_handle;
4164                self.wait_for_write(store, caller, handle)?
4165            }
4166        } else {
4167            ReturnCode::Cancelled(ItemCount::ZERO)
4168        };
4169
4170        if !matches!(code, ReturnCode::Blocked) {
4171            let transmit = store.concurrent_state_mut()?.get_mut(transmit_id)?;
4172
4173            match &transmit.write {
4174                WriteState::GuestReady { .. } => {
4175                    transmit.write = WriteState::Open;
4176                }
4177                WriteState::HostReady { .. } => bail_bug!("support host write cancellation"),
4178                WriteState::Open | WriteState::Dropped => {}
4179            }
4180        }
4181
4182        log::trace!("cancelled write {transmit_id:?}: {code:?}");
4183
4184        Ok(code)
4185    }
4186
4187    fn wait_for_read(
4188        self,
4189        store: &mut StoreOpaque,
4190        caller: RuntimeComponentInstanceIndex,
4191        handle: TableId<TransmitHandle>,
4192    ) -> Result<ReturnCode> {
4193        let waitable = Waitable::Transmit(handle);
4194        store.wait_for_event(self.runtime_instance(caller), waitable)?;
4195        let event = waitable.take_event(store.concurrent_state_mut()?)?;
4196        if let Some(event @ (Event::StreamRead { code, .. } | Event::FutureRead { code, .. })) =
4197            event
4198        {
4199            waitable.on_delivery(store, self, event)?;
4200            Ok(code)
4201        } else {
4202            bail_bug!("expected either a stream or future read event")
4203        }
4204    }
4205
4206    /// Cancel a pending stream or future read.
4207    fn cancel_read(
4208        self,
4209        store: &mut StoreOpaque,
4210        caller: RuntimeComponentInstanceIndex,
4211        transmit_id: TableId<TransmitState>,
4212        async_: bool,
4213    ) -> Result<ReturnCode> {
4214        let state = store.concurrent_state_mut()?;
4215        let transmit = state.get_mut(transmit_id)?;
4216        log::trace!(
4217            "host_cancel_read state {transmit_id:?}; read state {:?} write state {:?}",
4218            transmit.read,
4219            transmit.write
4220        );
4221
4222        let waitable = Waitable::Transmit(transmit.read_handle);
4223
4224        waitable.trap_if_has_sync_waiter(state)?;
4225        if !async_ {
4226            waitable.trap_if_in_waitable_set(state)?;
4227        }
4228
4229        let code = if let Some(event) = waitable.take_event(state)? {
4230            let (Event::FutureRead { code, .. } | Event::StreamRead { code, .. }) = event else {
4231                bail_bug!("expected either a stream or future read event")
4232            };
4233            waitable.on_delivery(store, self, event)?;
4234            match (code, event) {
4235                (ReturnCode::Completed(count), Event::StreamRead { .. }) => {
4236                    ReturnCode::Cancelled(count)
4237                }
4238                (ReturnCode::Dropped(_) | ReturnCode::Completed(_), _) => code,
4239                _ => bail_bug!("unexpected code/event combo"),
4240            }
4241        } else if let WriteState::HostReady {
4242            cancel,
4243            cancel_waker,
4244            ..
4245        } = &mut state.get_mut(transmit_id)?.write
4246        {
4247            *cancel = true;
4248            if let Some(waker) = cancel_waker.take() {
4249                waker.wake();
4250            }
4251
4252            if async_ {
4253                ReturnCode::Blocked
4254            } else {
4255                let handle = store
4256                    .concurrent_state_mut()?
4257                    .get_mut(transmit_id)?
4258                    .read_handle;
4259                self.wait_for_read(store, caller, handle)?
4260            }
4261        } else {
4262            ReturnCode::Cancelled(ItemCount::ZERO)
4263        };
4264
4265        if !matches!(code, ReturnCode::Blocked) {
4266            let transmit = store.concurrent_state_mut()?.get_mut(transmit_id)?;
4267
4268            match &transmit.read {
4269                ReadState::GuestReady { .. } => {
4270                    transmit.read = ReadState::Open;
4271                }
4272                ReadState::HostReady { .. } | ReadState::HostToHost { .. } => {
4273                    bail_bug!("support host read cancellation")
4274                }
4275                ReadState::Open | ReadState::Dropped => {}
4276            }
4277        }
4278
4279        log::trace!("cancelled read {transmit_id:?}: {code:?}");
4280
4281        Ok(code)
4282    }
4283
4284    /// Cancel a pending write for the specified stream or future from the guest.
4285    fn guest_cancel_write(
4286        self,
4287        store: &mut StoreOpaque,
4288        caller: RuntimeComponentInstanceIndex,
4289        ty: TransmitIndex,
4290        async_: bool,
4291        writer: u32,
4292    ) -> Result<ReturnCode> {
4293        let (rep, state) =
4294            get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, writer)?;
4295        let id = TableId::<TransmitHandle>::new(rep);
4296        log::trace!("guest cancel write {id:?} (handle {writer})");
4297        match state {
4298            TransmitLocalState::Write { .. } => {
4299                bail!("stream or future write cancelled when no write is pending")
4300            }
4301            TransmitLocalState::Read { .. } => {
4302                bail!("passed read end to `{{stream|future}}.cancel-write`")
4303            }
4304            TransmitLocalState::Cancelling => {
4305                bail!("stream or future write cancelled when a cancellation is already pending")
4306            }
4307            TransmitLocalState::Busy => {}
4308        }
4309        let transmit_id = store.concurrent_state_mut()?.get_mut(id)?.state;
4310        let code = self.cancel_write(store, caller, transmit_id, async_)?;
4311        let state =
4312            get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, writer)?.1;
4313        if let TransmitLocalState::Busy = state {
4314            *state = if matches!(code, ReturnCode::Blocked) {
4315                TransmitLocalState::Cancelling
4316            } else {
4317                TransmitLocalState::Write { done: false }
4318            };
4319        }
4320        Ok(code)
4321    }
4322
4323    /// Cancel a pending read for the specified stream or future from the guest.
4324    fn guest_cancel_read(
4325        self,
4326        store: &mut StoreOpaque,
4327        caller: RuntimeComponentInstanceIndex,
4328        ty: TransmitIndex,
4329        async_: bool,
4330        reader: u32,
4331    ) -> Result<ReturnCode> {
4332        let (rep, state) =
4333            get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, reader)?;
4334        let id = TableId::<TransmitHandle>::new(rep);
4335        log::trace!("guest cancel read {id:?} (handle {reader})");
4336        match state {
4337            TransmitLocalState::Read { .. } => {
4338                bail!("stream or future read cancelled when no read is pending")
4339            }
4340            TransmitLocalState::Write { .. } => {
4341                bail!("passed write end to `{{stream|future}}.cancel-read`")
4342            }
4343            TransmitLocalState::Cancelling => {
4344                bail!("stream or future read cancelled when a cancellation is already pending")
4345            }
4346            TransmitLocalState::Busy => {}
4347        }
4348        let transmit_id = store.concurrent_state_mut()?.get_mut(id)?.state;
4349        let code = self.cancel_read(store, caller, transmit_id, async_)?;
4350        let state =
4351            get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, reader)?.1;
4352        if let TransmitLocalState::Busy = state {
4353            *state = if matches!(code, ReturnCode::Blocked) {
4354                TransmitLocalState::Cancelling
4355            } else {
4356                TransmitLocalState::Read { done: false }
4357            };
4358        }
4359        Ok(code)
4360    }
4361
4362    /// Drop the readable end of the specified stream or future from the guest.
4363    fn guest_drop_readable(
4364        self,
4365        store: &mut StoreOpaque,
4366        ty: TransmitIndex,
4367        reader: u32,
4368    ) -> Result<()> {
4369        let (instance, counter) = self.id().get_mut_and_handle_counter(store);
4370        let table = instance.table_for_transmit(ty);
4371        let (rep, _is_done) = match ty {
4372            TransmitIndex::Stream(ty) => table.stream_remove_readable(counter, ty, reader)?,
4373            TransmitIndex::Future(ty) => table.future_remove_readable(counter, ty, reader)?,
4374        };
4375        let kind = match ty {
4376            TransmitIndex::Stream(_) => TransmitKind::Stream,
4377            TransmitIndex::Future(_) => TransmitKind::Future,
4378        };
4379        let id = TableId::<TransmitHandle>::new(rep);
4380        log::trace!("guest_drop_readable: drop reader {id:?}");
4381        store.host_drop_reader(id, kind)
4382    }
4383
4384    /// Create a new error context for the given component.
4385    pub(crate) fn error_context_new(
4386        self,
4387        store: &mut StoreOpaque,
4388        ty: TypeComponentLocalErrorContextTableIndex,
4389        options: OptionsIndex,
4390        debug_msg_address: u32,
4391        debug_msg_len: u32,
4392    ) -> Result<u32> {
4393        let lift_ctx = &mut LiftContext::new(store, options, self)?;
4394        let debug_msg = String::linear_lift_from_flat(
4395            lift_ctx,
4396            InterfaceType::String,
4397            &[ValRaw::u32(debug_msg_address), ValRaw::u32(debug_msg_len)],
4398        )?;
4399
4400        // Create a new ErrorContext that is tracked along with other concurrent state
4401        let err_ctx = ErrorContextState { debug_msg };
4402        let state = store.concurrent_state_mut()?;
4403        let table_id = state.push(err_ctx)?;
4404        let global_ref_count_idx =
4405            TypeComponentGlobalErrorContextTableIndex::from_u32(table_id.rep());
4406
4407        // Add to the global error context ref counts
4408        let _ = state
4409            .global_error_context_ref_counts
4410            .insert(global_ref_count_idx, GlobalErrorContextRefCount(1));
4411
4412        // Error context are tracked both locally (to a single component instance) and globally
4413        // the counts for both must stay in sync.
4414        //
4415        // Here we reflect the newly created global concurrent error context state into the
4416        // component instance's locally tracked count, along with the appropriate key into the global
4417        // ref tracking data structures to enable later lookup
4418        let (instance, counter) = self.id().get_mut_and_handle_counter(store);
4419        let local_idx = instance
4420            .table_for_error_context(ty)
4421            .error_context_insert(counter, table_id.rep())?;
4422
4423        Ok(local_idx)
4424    }
4425
4426    /// Retrieve the debug message from the specified error context.
4427    pub(super) fn error_context_debug_message<T>(
4428        self,
4429        store: StoreContextMut<T>,
4430        ty: TypeComponentLocalErrorContextTableIndex,
4431        options: OptionsIndex,
4432        err_ctx_handle: u32,
4433        debug_msg_address: u32,
4434    ) -> Result<()> {
4435        // Retrieve the error context and internal debug message
4436        let handle_table_id_rep = self
4437            .id()
4438            .get_mut(store.0)
4439            .table_for_error_context(ty)
4440            .error_context_rep(err_ctx_handle)?;
4441
4442        let state = store.0.concurrent_state_mut()?;
4443        // Get the state associated with the error context
4444        let ErrorContextState { debug_msg } =
4445            state.get_mut(TableId::<ErrorContextState>::new(handle_table_id_rep))?;
4446        let debug_msg = debug_msg.clone();
4447
4448        let lower_cx = &mut LowerContext::new(store, options, self);
4449        let debug_msg_address = usize::try_from(debug_msg_address)?;
4450        if debug_msg_address % usize::try_from(<str as ComponentType>::ALIGN32)? != 0 {
4451            bail!("invalid debug message pointer: not aligned");
4452        }
4453        // Lower the string into the component's memory.
4454        //
4455        // Note that the "8" here is the size of a WIT `string` in linear
4456        // memory, the ptr+length. This'll need to be updated when `memory64`
4457        // comes along. (FIXME(#4311))
4458        let offset = lower_cx
4459            .as_slice_mut()
4460            .get(debug_msg_address..)
4461            .and_then(|b| b.get(..8))
4462            .map(|_| debug_msg_address)
4463            .ok_or_else(|| crate::format_err!("invalid debug message pointer: out of bounds"))?;
4464        debug_msg
4465            .as_str()
4466            .linear_lower_to_memory(lower_cx, InterfaceType::String, offset)?;
4467
4468        Ok(())
4469    }
4470
4471    /// Implements the `future.cancel-read` intrinsic.
4472    pub(crate) fn future_cancel_read(
4473        self,
4474        store: &mut StoreOpaque,
4475        caller: RuntimeComponentInstanceIndex,
4476        ty: TypeFutureTableIndex,
4477        async_: bool,
4478        reader: u32,
4479    ) -> Result<u32> {
4480        self.guest_cancel_read(store, caller, TransmitIndex::Future(ty), async_, reader)
4481            .map(|v| v.encode())
4482    }
4483
4484    /// Implements the `future.cancel-write` intrinsic.
4485    pub(crate) fn future_cancel_write(
4486        self,
4487        store: &mut StoreOpaque,
4488        caller: RuntimeComponentInstanceIndex,
4489        ty: TypeFutureTableIndex,
4490        async_: bool,
4491        writer: u32,
4492    ) -> Result<u32> {
4493        self.guest_cancel_write(store, caller, TransmitIndex::Future(ty), async_, writer)
4494            .map(|v| v.encode())
4495    }
4496
4497    /// Implements the `stream.cancel-read` intrinsic.
4498    pub(crate) fn stream_cancel_read(
4499        self,
4500        store: &mut StoreOpaque,
4501        caller: RuntimeComponentInstanceIndex,
4502        ty: TypeStreamTableIndex,
4503        async_: bool,
4504        reader: u32,
4505    ) -> Result<u32> {
4506        self.guest_cancel_read(store, caller, TransmitIndex::Stream(ty), async_, reader)
4507            .map(|v| v.encode())
4508    }
4509
4510    /// Implements the `stream.cancel-write` intrinsic.
4511    pub(crate) fn stream_cancel_write(
4512        self,
4513        store: &mut StoreOpaque,
4514        caller: RuntimeComponentInstanceIndex,
4515        ty: TypeStreamTableIndex,
4516        async_: bool,
4517        writer: u32,
4518    ) -> Result<u32> {
4519        self.guest_cancel_write(store, caller, TransmitIndex::Stream(ty), async_, writer)
4520            .map(|v| v.encode())
4521    }
4522
4523    /// Implements the `future.drop-readable` intrinsic.
4524    pub(crate) fn future_drop_readable(
4525        self,
4526        store: &mut StoreOpaque,
4527        ty: TypeFutureTableIndex,
4528        reader: u32,
4529    ) -> Result<()> {
4530        self.guest_drop_readable(store, TransmitIndex::Future(ty), reader)
4531    }
4532
4533    /// Implements the `stream.drop-readable` intrinsic.
4534    pub(crate) fn stream_drop_readable(
4535        self,
4536        store: &mut StoreOpaque,
4537        ty: TypeStreamTableIndex,
4538        reader: u32,
4539    ) -> Result<()> {
4540        self.guest_drop_readable(store, TransmitIndex::Stream(ty), reader)
4541    }
4542
4543    /// Allocate a new future or stream and grant ownership of both the read and
4544    /// write ends to the (sub-)component instance to which the specified
4545    /// `TransmitIndex` belongs.
4546    fn guest_new(self, store: &mut StoreOpaque, ty: TransmitIndex) -> Result<ResourcePair> {
4547        let (write, read) = store
4548            .concurrent_state_mut()?
4549            .new_transmit(TransmitOrigin::guest(self.id().instance(), ty))?;
4550
4551        let (instance, counter) = self.id().get_mut_and_handle_counter(store);
4552        let table = instance.table_for_transmit(ty);
4553        let (read_handle, write_handle) = match ty {
4554            TransmitIndex::Future(ty) => (
4555                table.future_insert_read(counter, ty, read.rep())?,
4556                table.future_insert_write(counter, ty, write.rep())?,
4557            ),
4558            TransmitIndex::Stream(ty) => (
4559                table.stream_insert_read(counter, ty, read.rep())?,
4560                table.stream_insert_write(counter, ty, write.rep())?,
4561            ),
4562        };
4563
4564        let state = store.concurrent_state_mut()?;
4565        state.get_mut(read)?.common.handle = Some(read_handle);
4566        state.get_mut(write)?.common.handle = Some(write_handle);
4567
4568        Ok(ResourcePair {
4569            write: write_handle,
4570            read: read_handle,
4571        })
4572    }
4573
4574    /// Drop the specified error context.
4575    pub(crate) fn error_context_drop(
4576        self,
4577        store: &mut StoreOpaque,
4578        ty: TypeComponentLocalErrorContextTableIndex,
4579        error_context: u32,
4580    ) -> Result<()> {
4581        let (instance, counter) = self.id().get_mut_and_handle_counter(store);
4582
4583        let local_handle_table = instance.table_for_error_context(ty);
4584
4585        let rep = local_handle_table.error_context_drop(counter, error_context)?;
4586
4587        let global_ref_count_idx = TypeComponentGlobalErrorContextTableIndex::from_u32(rep);
4588
4589        let state = store.concurrent_state_mut()?;
4590        let Some(GlobalErrorContextRefCount(global_ref_count)) = state
4591            .global_error_context_ref_counts
4592            .get_mut(&global_ref_count_idx)
4593        else {
4594            bail_bug!("retrieve concurrent state for error context during drop")
4595        };
4596
4597        // Reduce the component-global ref count, removing tracking if necessary
4598        if *global_ref_count < 1 {
4599            bail_bug!("ref count unexpectedly zero");
4600        }
4601        *global_ref_count -= 1;
4602        if *global_ref_count == 0 {
4603            state
4604                .global_error_context_ref_counts
4605                .remove(&global_ref_count_idx);
4606
4607            state
4608                .delete(TableId::<ErrorContextState>::new(rep))
4609                .context("deleting component-global error context data")?;
4610        }
4611
4612        Ok(())
4613    }
4614
4615    /// Transfer ownership of the specified stream or future read end from one
4616    /// guest to another.
4617    fn guest_transfer(
4618        self,
4619        store: &mut StoreOpaque,
4620        src_idx: u32,
4621        src: TransmitIndex,
4622        dst: TransmitIndex,
4623    ) -> Result<u32> {
4624        let id = self.lift_index_to_transmit(store, src, src_idx)?;
4625        self.lower_transmit_to_index(store, dst, id)
4626    }
4627
4628    fn lift_index_to_transmit(
4629        self,
4630        store: &mut StoreOpaque,
4631        ty: TransmitIndex,
4632        src_idx: u32,
4633    ) -> Result<TableId<TransmitHandle>> {
4634        let (state, _, counter, _, instance) = store.lift_context_parts(self);
4635        lift_index_to_transmit(instance, counter, state.concurrent_state_mut(), ty, src_idx)
4636    }
4637
4638    fn lower_transmit_to_index(
4639        self,
4640        store: &mut StoreOpaque,
4641        ty: TransmitIndex,
4642        id: TableId<TransmitHandle>,
4643    ) -> Result<u32> {
4644        let (state, _, counter, _, instance) = store.lift_context_parts(self);
4645        lower_transmit_to_index(instance, counter, state.concurrent_state_mut(), ty, id)
4646    }
4647
4648    /// Implements the `future.new` intrinsic.
4649    pub(crate) fn future_new(
4650        self,
4651        store: &mut StoreOpaque,
4652        ty: TypeFutureTableIndex,
4653    ) -> Result<ResourcePair> {
4654        self.guest_new(store, TransmitIndex::Future(ty))
4655    }
4656
4657    /// Implements the `stream.new` intrinsic.
4658    pub(crate) fn stream_new(
4659        self,
4660        store: &mut StoreOpaque,
4661        ty: TypeStreamTableIndex,
4662    ) -> Result<ResourcePair> {
4663        self.guest_new(store, TransmitIndex::Stream(ty))
4664    }
4665
4666    /// Transfer ownership of the specified future read end from one guest to
4667    /// another.
4668    pub(crate) fn future_transfer(
4669        self,
4670        store: &mut StoreOpaque,
4671        src_idx: u32,
4672        src: TypeFutureTableIndex,
4673        dst: TypeFutureTableIndex,
4674    ) -> Result<u32> {
4675        self.guest_transfer(
4676            store,
4677            src_idx,
4678            TransmitIndex::Future(src),
4679            TransmitIndex::Future(dst),
4680        )
4681    }
4682
4683    /// Transfer ownership of the specified stream read end from one guest to
4684    /// another.
4685    pub(crate) fn stream_transfer(
4686        self,
4687        store: &mut StoreOpaque,
4688        src_idx: u32,
4689        src: TypeStreamTableIndex,
4690        dst: TypeStreamTableIndex,
4691    ) -> Result<u32> {
4692        self.guest_transfer(
4693            store,
4694            src_idx,
4695            TransmitIndex::Stream(src),
4696            TransmitIndex::Stream(dst),
4697        )
4698    }
4699
4700    /// Copy the specified error context from one component to another.
4701    pub(crate) fn error_context_transfer(
4702        self,
4703        store: &mut StoreOpaque,
4704        src_idx: u32,
4705        src: TypeComponentLocalErrorContextTableIndex,
4706        dst: TypeComponentLocalErrorContextTableIndex,
4707    ) -> Result<u32> {
4708        let (mut instance, counter) = self.id().get_mut_and_handle_counter(store);
4709        let rep = instance
4710            .as_mut()
4711            .table_for_error_context(src)
4712            .error_context_rep(src_idx)?;
4713        let dst_idx = instance
4714            .table_for_error_context(dst)
4715            .error_context_insert(counter, rep)?;
4716
4717        // Update the global (cross-subcomponent) count for error contexts
4718        // as the new component has essentially created a new reference that will
4719        // be dropped/handled independently
4720        let global_ref_count = store
4721            .concurrent_state_mut()?
4722            .global_error_context_ref_counts
4723            .get_mut(&TypeComponentGlobalErrorContextTableIndex::from_u32(rep))
4724            .context("global ref count present for existing (sub)component error context")?;
4725
4726        global_ref_count.0 = global_ref_count
4727            .0
4728            .checked_add(1)
4729            .ok_or_else(|| format_err!(Trap::ReferenceCountOverflow))?;
4730
4731        Ok(dst_idx)
4732    }
4733}
4734
4735/// Performs the opertion of lifting a future or stream from and instance into
4736/// the `TransmitHandle` for it.
4737///
4738/// The `src_idx` is the guest-specified index within `instance` and `ty` is the
4739/// expected type of future/stream.
4740fn lift_index_to_transmit(
4741    instance: Pin<&mut ComponentInstance>,
4742    counter: &mut HandleCounter,
4743    concurrent_state: &mut ConcurrentState,
4744    ty: TransmitIndex,
4745    src_idx: u32,
4746) -> Result<TableId<TransmitHandle>> {
4747    let handle_table = instance.table_for_transmit(ty);
4748    let (rep, is_done) = match ty {
4749        TransmitIndex::Future(idx) => handle_table.future_remove_readable(counter, idx, src_idx)?,
4750        TransmitIndex::Stream(idx) => handle_table.stream_remove_readable(counter, idx, src_idx)?,
4751    };
4752    let desc = match ty {
4753        TransmitIndex::Future(_) => "future",
4754        TransmitIndex::Stream(_) => "stream",
4755    };
4756    if is_done {
4757        bail!("cannot lift {desc} after being notified that the writable end dropped");
4758    }
4759    let id = TableId::<TransmitHandle>::new(rep);
4760    let future = concurrent_state.get_mut(id)?;
4761    if future.common.set.is_some() {
4762        bail!("cannot lift {desc} while it's in a waitable set");
4763    }
4764    future.common.handle = None;
4765
4766    let state = future.state;
4767    if concurrent_state.get_mut(state)?.done {
4768        match ty {
4769            TransmitIndex::Future(_) => bail!("cannot lift {desc} after previous read succeeded"),
4770            TransmitIndex::Stream(_) => bail!(Trap::LiftDroppedStream),
4771        };
4772    }
4773
4774    Ok(id)
4775}
4776
4777/// Performs the opertion of lowering a future or stream `TransmitHandle` into
4778/// an instance.
4779fn lower_transmit_to_index(
4780    instance: Pin<&mut ComponentInstance>,
4781    counter: &mut HandleCounter,
4782    concurrent_state: &mut ConcurrentState,
4783    ty: TransmitIndex,
4784    id: TableId<TransmitHandle>,
4785) -> Result<u32> {
4786    let state = concurrent_state.get_mut(id)?.state;
4787    debug_assert_eq!(concurrent_state.get_mut(state)?.read_handle, id);
4788    let handle_table = instance.table_for_transmit(ty);
4789    let handle = match ty {
4790        TransmitIndex::Future(idx) => handle_table.future_insert_read(counter, idx, id.rep()),
4791        TransmitIndex::Stream(idx) => handle_table.stream_insert_read(counter, idx, id.rep()),
4792    }?;
4793    concurrent_state.get_mut(id)?.common.handle = Some(handle);
4794    Ok(handle)
4795}
4796
4797impl ComponentInstance {
4798    fn table_for_transmit(self: Pin<&mut Self>, ty: TransmitIndex) -> &mut HandleTable {
4799        let (states, types) = self.instance_states();
4800        let runtime_instance = match ty {
4801            TransmitIndex::Stream(ty) => types[ty].instance,
4802            TransmitIndex::Future(ty) => types[ty].instance,
4803        };
4804        states[runtime_instance].handle_table()
4805    }
4806
4807    fn table_for_error_context(
4808        self: Pin<&mut Self>,
4809        ty: TypeComponentLocalErrorContextTableIndex,
4810    ) -> &mut HandleTable {
4811        let (states, types) = self.instance_states();
4812        let runtime_instance = types[ty].instance;
4813        states[runtime_instance].handle_table()
4814    }
4815
4816    fn get_mut_by_index(
4817        self: Pin<&mut Self>,
4818        ty: TransmitIndex,
4819        index: u32,
4820    ) -> Result<(u32, &mut TransmitLocalState)> {
4821        get_mut_by_index_from(self.table_for_transmit(ty), ty, index)
4822    }
4823}
4824
4825impl ConcurrentState {
4826    fn send_write_result(
4827        &mut self,
4828        ty: TransmitIndex,
4829        id: TableId<TransmitState>,
4830        handle: u32,
4831        code: ReturnCode,
4832    ) -> Result<()> {
4833        let write_handle = self.get_mut(id)?.write_handle.rep();
4834        self.set_event(
4835            write_handle,
4836            match ty {
4837                TransmitIndex::Future(ty) => Event::FutureWrite {
4838                    code,
4839                    pending: Some((ty, handle)),
4840                },
4841                TransmitIndex::Stream(ty) => Event::StreamWrite {
4842                    code,
4843                    pending: Some((ty, handle)),
4844                },
4845            },
4846        )
4847    }
4848
4849    fn send_read_result(
4850        &mut self,
4851        ty: TransmitIndex,
4852        id: TableId<TransmitState>,
4853        handle: u32,
4854        code: ReturnCode,
4855    ) -> Result<()> {
4856        let read_handle = self.get_mut(id)?.read_handle.rep();
4857        self.set_event(
4858            read_handle,
4859            match ty {
4860                TransmitIndex::Future(ty) => Event::FutureRead {
4861                    code,
4862                    pending: Some((ty, handle)),
4863                },
4864                TransmitIndex::Stream(ty) => Event::StreamRead {
4865                    code,
4866                    pending: Some((ty, handle)),
4867                },
4868            },
4869        )
4870    }
4871
4872    fn take_event(&mut self, waitable: u32) -> Result<Option<Event>> {
4873        Waitable::Transmit(TableId::<TransmitHandle>::new(waitable)).take_event(self)
4874    }
4875
4876    fn set_event(&mut self, waitable: u32, event: Event) -> Result<()> {
4877        Waitable::Transmit(TableId::<TransmitHandle>::new(waitable)).set_event(self, Some(event))
4878    }
4879
4880    /// Set or update the event for the specified waitable.
4881    ///
4882    /// If there is already an event set for this waitable, we assert that it is
4883    /// of the same variant as the new one and reuse the `ReturnCode` count and
4884    /// the `pending` field if applicable.
4885    // TODO: This is a bit awkward due to how
4886    // `Event::{Stream,Future}{Write,Read}` and
4887    // `ReturnCode::{Completed,Dropped,Cancelled}` are currently represented.
4888    // Consider updating those representations in a way that allows this
4889    // function to be simplified.
4890    fn update_event(&mut self, waitable: u32, event: Event) -> Result<()> {
4891        let waitable = Waitable::Transmit(TableId::<TransmitHandle>::new(waitable));
4892
4893        fn update_code(old: ReturnCode, new: ReturnCode) -> Result<ReturnCode> {
4894            let (ReturnCode::Completed(count)
4895            | ReturnCode::Dropped(count)
4896            | ReturnCode::Cancelled(count)) = old
4897            else {
4898                bail_bug!("unexpected old return code")
4899            };
4900
4901            Ok(match new {
4902                ReturnCode::Dropped(ItemCount::ZERO) => ReturnCode::Dropped(count),
4903                ReturnCode::Cancelled(ItemCount::ZERO) => ReturnCode::Cancelled(count),
4904                _ => bail_bug!("unexpected new return code"),
4905            })
4906        }
4907
4908        let event = match (waitable.take_event(self)?, event) {
4909            (None, _) => event,
4910            (Some(old @ Event::FutureWrite { .. }), Event::FutureWrite { .. }) => old,
4911            (Some(old @ Event::FutureRead { .. }), Event::FutureRead { .. }) => old,
4912            (
4913                Some(Event::StreamWrite {
4914                    code: old_code,
4915                    pending: old_pending,
4916                }),
4917                Event::StreamWrite { code, pending },
4918            ) => Event::StreamWrite {
4919                code: update_code(old_code, code)?,
4920                pending: old_pending.or(pending),
4921            },
4922            (
4923                Some(Event::StreamRead {
4924                    code: old_code,
4925                    pending: old_pending,
4926                }),
4927                Event::StreamRead { code, pending },
4928            ) => Event::StreamRead {
4929                code: update_code(old_code, code)?,
4930                pending: old_pending.or(pending),
4931            },
4932            _ => bail_bug!("unexpected event combination"),
4933        };
4934
4935        waitable.set_event(self, Some(event))
4936    }
4937
4938    /// Allocate a new future or stream, including the `TransmitState` and the
4939    /// `TransmitHandle`s corresponding to the read and write ends.
4940    fn new_transmit(
4941        &mut self,
4942        origin: TransmitOrigin,
4943    ) -> Result<(TableId<TransmitHandle>, TableId<TransmitHandle>)> {
4944        // All three entries must fit before inserting any of them so a full
4945        // table does not leave an inaccessible state or handle behind.
4946        self.table().ensure_capacity(3)?;
4947        let state_id = self.push(TransmitState::new(origin))?;
4948
4949        let write = self.push(TransmitHandle::new(state_id))?;
4950        let read = self.push(TransmitHandle::new(state_id))?;
4951
4952        let state = self.get_mut(state_id)?;
4953        state.write_handle = write;
4954        state.read_handle = read;
4955
4956        log::trace!("new transmit: state {state_id:?}; write {write:?}; read {read:?}",);
4957
4958        Ok((write, read))
4959    }
4960
4961    /// Delete the specified future or stream, including the read and write ends.
4962    fn delete_transmit(&mut self, state_id: TableId<TransmitState>) -> Result<()> {
4963        let state = self.delete(state_id)?;
4964        self.delete(state.write_handle)?;
4965        self.delete(state.read_handle)?;
4966
4967        log::trace!(
4968            "delete transmit: state {state_id:?}; write {:?}; read {:?}",
4969            state.write_handle,
4970            state.read_handle,
4971        );
4972
4973        Ok(())
4974    }
4975}
4976
4977pub(crate) struct ResourcePair {
4978    pub(crate) write: u32,
4979    pub(crate) read: u32,
4980}
4981
4982impl Waitable {
4983    /// Handle the imminent delivery of the specified event, e.g. by updating
4984    /// the state of the stream or future.
4985    pub(super) fn on_delivery(
4986        &self,
4987        store: &mut StoreOpaque,
4988        instance: Instance,
4989        event: Event,
4990    ) -> Result<()> {
4991        if let Event::FutureRead {
4992            code: ReturnCode::Dropped(_),
4993            ..
4994        }
4995        | Event::FutureWrite {
4996            code: ReturnCode::Dropped(_),
4997            ..
4998        }
4999        | Event::StreamRead {
5000            code: ReturnCode::Dropped(_),
5001            ..
5002        }
5003        | Event::StreamWrite {
5004            code: ReturnCode::Dropped(_),
5005            ..
5006        } = event
5007        {
5008            let Waitable::Transmit(transmit_handle) = self else {
5009                bail_bug!("unexpected `{event:?}` for `{self:?}`");
5010            };
5011            let state = store.concurrent_state_mut()?;
5012            let transmit_id = state.get_mut(*transmit_handle)?.state;
5013            state.get_mut(transmit_id)?.done = true;
5014        }
5015
5016        let instance = instance.id().get_mut(store);
5017        let (rep, state, code) = match event {
5018            Event::FutureRead {
5019                pending: Some((ty, handle)),
5020                code,
5021            }
5022            | Event::FutureWrite {
5023                pending: Some((ty, handle)),
5024                code,
5025            } => {
5026                let runtime_instance = instance.component().types()[ty].instance;
5027                let (rep, state) = instance.instance_states().0[runtime_instance]
5028                    .handle_table()
5029                    .future_rep(ty, handle)?;
5030                (rep, state, code)
5031            }
5032            Event::StreamRead {
5033                pending: Some((ty, handle)),
5034                code,
5035            }
5036            | Event::StreamWrite {
5037                pending: Some((ty, handle)),
5038                code,
5039            } => {
5040                let runtime_instance = instance.component().types()[ty].instance;
5041                let (rep, state) = instance.instance_states().0[runtime_instance]
5042                    .handle_table()
5043                    .stream_rep(ty, handle)?;
5044                (rep, state, code)
5045            }
5046            _ => return Ok(()),
5047        };
5048        if rep != self.rep() {
5049            bail_bug!("unexpected rep mismatch");
5050        }
5051        if !matches!(
5052            *state,
5053            TransmitLocalState::Busy | TransmitLocalState::Cancelling
5054        ) {
5055            bail_bug!("expected state to be busy or cancelling");
5056        }
5057        let done = matches!(code, ReturnCode::Dropped(_));
5058        *state = match event {
5059            Event::FutureRead { .. } | Event::StreamRead { .. } => {
5060                TransmitLocalState::Read { done }
5061            }
5062            Event::FutureWrite { .. } | Event::StreamWrite { .. } => {
5063                TransmitLocalState::Write { done }
5064            }
5065            _ => bail_bug!("unexpected event for stream"),
5066        };
5067
5068        let transmit_handle = TableId::<TransmitHandle>::new(rep);
5069        let state = store.concurrent_state_mut()?;
5070        let transmit_id = state.get_mut(transmit_handle)?.state;
5071        let transmit = state.get_mut(transmit_id)?;
5072
5073        match event {
5074            Event::StreamRead { .. } => {
5075                transmit.read = ReadState::Open;
5076            }
5077            Event::StreamWrite { .. } => transmit.write = WriteState::Open,
5078            _ => {}
5079        }
5080        Ok(())
5081    }
5082}
5083
5084/// Determine whether an intra-component read/write is allowed for the specified
5085/// `stream` or `future` payload type according to the component model
5086/// specification.
5087fn allow_intra_component_read_write(ty: Option<&InterfaceType>) -> bool {
5088    matches!(
5089        ty,
5090        None | Some(
5091            InterfaceType::S8
5092                | InterfaceType::U8
5093                | InterfaceType::S16
5094                | InterfaceType::U16
5095                | InterfaceType::S32
5096                | InterfaceType::U32
5097                | InterfaceType::S64
5098                | InterfaceType::U64
5099                | InterfaceType::Float32
5100                | InterfaceType::Float64
5101        )
5102    )
5103}
5104
5105/// Helper structure to manage moving a `T` in/out of an interior `Mutex` which
5106/// contains an
5107/// `Option<T>`
5108struct LockedState<T> {
5109    inner: TryMutex<Option<T>>,
5110}
5111
5112impl<T> LockedState<T> {
5113    /// Creates a new initial state with `value` stored.
5114    fn new(value: T) -> Self {
5115        Self {
5116            inner: TryMutex::new(Some(value)),
5117        }
5118    }
5119
5120    /// Attempts to lock the inner mutex and return its guard.
5121    ///
5122    /// # Errors
5123    ///
5124    /// Fails if this lock is either poisoned or if it's currently locked.
5125    /// As-used in this file there should never actually be contention on this
5126    /// lock nor recursive access so failing to acquire the lock is a fatal
5127    /// error that gets propagated upwards.
5128    fn try_lock(&self) -> Result<TryMutexGuard<'_, Option<T>>> {
5129        match self.inner.try_lock() {
5130            Some(lock) => Ok(lock),
5131            None => bail_bug!("should not have contention on state lock"),
5132        }
5133    }
5134
5135    /// Takes the inner `T` out of this state, returning it as a guard which
5136    /// will put it back when finished.
5137    ///
5138    /// # Errors
5139    ///
5140    /// Returns an error if the state `T` isn't present.
5141    fn take(&self) -> Result<LockedStateGuard<'_, T>> {
5142        let result = self.try_lock()?.take();
5143        match result {
5144            Some(result) => Ok(LockedStateGuard {
5145                value: ManuallyDrop::new(result),
5146                state: self,
5147            }),
5148            None => bail_bug!("lock value unexpectedly missing"),
5149        }
5150    }
5151
5152    /// Performs the operation `f` on the inner state `&mut T`.
5153    ///
5154    /// This will acquire the internal lock and invoke `f`, so `f` should not
5155    /// expect to be able to recursively acquire this lock.
5156    ///
5157    /// # Errors
5158    ///
5159    /// Returns an error if the state `T` isn't present.
5160    fn with<R>(&self, f: impl FnOnce(&mut T) -> R) -> Result<R> {
5161        let mut inner = self.try_lock()?;
5162        match &mut *inner {
5163            Some(state) => Ok(f(state)),
5164            None => bail_bug!("lock value unexpectedly missing"),
5165        }
5166    }
5167}
5168
5169/// Helper structure returned from [`LockedState::take`] which will put the
5170/// state specified by `value` back into the original lock once this is dropped.
5171struct LockedStateGuard<'a, T> {
5172    value: ManuallyDrop<T>,
5173    state: &'a LockedState<T>,
5174}
5175
5176impl<T> Deref for LockedStateGuard<'_, T> {
5177    type Target = T;
5178
5179    fn deref(&self) -> &T {
5180        &self.value
5181    }
5182}
5183
5184impl<T> DerefMut for LockedStateGuard<'_, T> {
5185    fn deref_mut(&mut self) -> &mut T {
5186        &mut self.value
5187    }
5188}
5189
5190impl<T> Drop for LockedStateGuard<'_, T> {
5191    fn drop(&mut self) {
5192        // SAFETY: `ManuallyDrop::take` requires that after invoked the
5193        // original value is not read. This is the `Drop` for this type which
5194        // means we have exclusive ownership and it is not read further in the
5195        // destructor, satisfying this requirement.
5196        let value = unsafe { ManuallyDrop::take(&mut self.value) };
5197
5198        // If this fails due to contention that's a bug, but we're not in a
5199        // position to panic due to this being a destructor nor return an error,
5200        // so defer the bug to showing up later.
5201        if let Ok(mut lock) = self.state.try_lock() {
5202            *lock = Some(value);
5203        }
5204    }
5205}
5206
5207#[cfg(test)]
5208mod tests {
5209    use super::*;
5210    use crate::{Engine, Store};
5211    use core::future::pending;
5212    use core::pin::pin;
5213    use std::sync::LazyLock;
5214
5215    static ENGINE: LazyLock<Engine> = LazyLock::new(Engine::default);
5216
5217    fn poll_future_producer<T>(rx: Pin<&mut T>, finish: bool) -> Poll<Result<Option<T::Item>>>
5218    where
5219        T: FutureProducer<()>,
5220    {
5221        rx.poll_produce(
5222            &mut Context::from_waker(Waker::noop()),
5223            Store::new(&ENGINE, ()).as_context_mut(),
5224            finish,
5225        )
5226    }
5227
5228    #[test]
5229    fn future_producer() {
5230        let mut fut = pin!(async { crate::error::Ok(()) });
5231        assert!(matches!(
5232            poll_future_producer(fut.as_mut(), false),
5233            Poll::Ready(Ok(Some(()))),
5234        ));
5235
5236        let mut fut = pin!(async { crate::error::Ok(()) });
5237        assert!(matches!(
5238            poll_future_producer(fut.as_mut(), true),
5239            Poll::Ready(Ok(Some(()))),
5240        ));
5241
5242        let mut fut = pin!(pending::<Result<()>>());
5243        assert!(matches!(
5244            poll_future_producer(fut.as_mut(), false),
5245            Poll::Pending,
5246        ));
5247        assert!(matches!(
5248            poll_future_producer(fut.as_mut(), true),
5249            Poll::Ready(Ok(None)),
5250        ));
5251
5252        let (tx, rx) = oneshot::channel();
5253        let mut rx = pin!(rx);
5254        assert!(matches!(
5255            poll_future_producer(rx.as_mut(), false),
5256            Poll::Pending,
5257        ));
5258        assert!(matches!(
5259            poll_future_producer(rx.as_mut(), true),
5260            Poll::Ready(Ok(None)),
5261        ));
5262        tx.send(()).unwrap();
5263        assert!(matches!(
5264            poll_future_producer(rx.as_mut(), true),
5265            Poll::Ready(Ok(Some(()))),
5266        ));
5267
5268        let (tx, rx) = oneshot::channel();
5269        let mut rx = pin!(rx);
5270        tx.send(()).unwrap();
5271        assert!(matches!(
5272            poll_future_producer(rx.as_mut(), false),
5273            Poll::Ready(Ok(Some(()))),
5274        ));
5275
5276        let (tx, rx) = oneshot::channel::<()>();
5277        let mut rx = pin!(rx);
5278        drop(tx);
5279        assert!(matches!(
5280            poll_future_producer(rx.as_mut(), false),
5281            Poll::Ready(Err(..)),
5282        ));
5283
5284        let (tx, rx) = oneshot::channel::<()>();
5285        let mut rx = pin!(rx);
5286        drop(tx);
5287        assert!(matches!(
5288            poll_future_producer(rx.as_mut(), true),
5289            Poll::Ready(Err(..)),
5290        ));
5291    }
5292}