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