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