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