wasmtime/runtime/component/concurrent/futures_and_streams/
buffers.rs

1#[cfg(feature = "component-model-async-bytes")]
2use bytes::{Bytes, BytesMut};
3#[cfg(feature = "component-model-async-bytes")]
4use std::io::Cursor;
5use std::mem::{self, MaybeUninit};
6use std::slice;
7use std::vec::Vec;
8
9// Inner module here to restrict possible readers of the fields of
10// `UntypedWriteBuffer`.
11pub use untyped::*;
12mod untyped {
13    use super::WriteBuffer;
14    use std::any::TypeId;
15    use std::marker;
16    use std::mem;
17
18    /// Helper structure to type-erase the `T` in `WriteBuffer<T>`.
19    ///
20    /// This is constructed with a `&mut dyn WriteBuffer<T>` and then can only
21    /// be viewed as `&mut dyn WriteBuffer<T>` as well. The `T`, however, is
22    /// carried through methods rather than the struct itself.
23    ///
24    /// Note that this structure has a lifetime `'a` which forces an active
25    /// borrow on the original buffer passed in.
26    pub struct UntypedWriteBuffer<'a> {
27        element_type_id: TypeId,
28        buf: *mut dyn WriteBuffer<()>,
29        _marker: marker::PhantomData<&'a mut dyn WriteBuffer<()>>,
30    }
31
32    /// Helper structure to transmute between `WriteBuffer<T>` and
33    /// `WriteBuffer<()>`.
34    union ReinterpretWriteBuffer<T> {
35        typed: *mut dyn WriteBuffer<T>,
36        untyped: *mut dyn WriteBuffer<()>,
37    }
38
39    impl<'a> UntypedWriteBuffer<'a> {
40        /// Creates a new `UntypedWriteBuffer` from the `buf` provided.
41        ///
42        /// The returned value can be used with the `get_mut` method to get the
43        /// original write buffer back.
44        pub fn new<T: 'static>(buf: &'a mut dyn WriteBuffer<T>) -> UntypedWriteBuffer<'a> {
45            UntypedWriteBuffer {
46                element_type_id: TypeId::of::<T>(),
47                // SAFETY: this is `unsafe` due to reading union fields. That
48                // is safe here because `typed` and `untyped` have the same size
49                // and we're otherwise reinterpreting a raw pointer with a type
50                // parameter to one without one.
51                buf: unsafe {
52                    let r = ReinterpretWriteBuffer { typed: buf };
53                    assert_eq!(mem::size_of_val(&r.typed), mem::size_of_val(&r.untyped));
54                    r.untyped
55                },
56                _marker: marker::PhantomData,
57            }
58        }
59
60        /// Acquires the underyling `WriteBuffer<T>` this was created with.
61        ///
62        /// # Panics
63        ///
64        /// Panics if `T` does not match the type that this was created with.
65        pub fn get_mut<T: 'static>(&mut self) -> &mut dyn WriteBuffer<T> {
66            assert_eq!(self.element_type_id, TypeId::of::<T>());
67            // SAFETY: the `T` has been checked with `TypeId` and this
68            // structure also is proof of valid existence of the original
69            // `&mut WriteBuffer<T>`, so taking the raw pointer back to a safe
70            // reference is valid.
71            unsafe { &mut *ReinterpretWriteBuffer { untyped: self.buf }.typed }
72        }
73    }
74}
75
76/// Trait representing a buffer which may be written to a `StreamWriter`.
77///
78/// See also [`crate::component::Instance::stream`].
79///
80/// # Unsafety
81///
82/// This trait is unsafe due to the contract of the `take` function. This trait
83/// is only safe to implement if the `take` function is implemented correctly,
84/// namely that all the items passed to the closure are fully initialized for
85/// `T`.
86pub unsafe trait WriteBuffer<T>: Send + Sync + 'static {
87    /// Slice of items remaining to be read.
88    fn remaining(&self) -> &[T];
89
90    /// Skip and drop the specified number of items.
91    fn skip(&mut self, count: usize);
92
93    /// Take ownership of the specified number of items.
94    ///
95    /// This function will take `count` items from `self` and pass them as a
96    /// contiguous slice to the closure `fun` provided. The `fun` closure may
97    /// assume that the items are all fully initialized and available to read.
98    /// It is expected that `fun` will read all the items provided. Any items
99    /// that aren't read by `fun` will be leaked.
100    ///
101    /// # Panics
102    ///
103    /// Panics if `count` is larger than `self.remaining()`. If `fun` panics
104    /// then items may be leaked.
105    fn take(&mut self, count: usize, fun: &mut dyn FnMut(&[MaybeUninit<T>]));
106}
107
108/// Trait representing a buffer which may be used to read from a `StreamReader`.
109///
110/// See also [`crate::component::Instance::stream`].
111pub trait ReadBuffer<T>: Send + Sync + 'static {
112    /// Move the specified items into this buffer.
113    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I);
114
115    /// Number of items which may be read before this buffer is full.
116    fn remaining_capacity(&self) -> usize;
117
118    /// Move (i.e. take ownership of) the specified items into this buffer.
119    ///
120    /// This method will drain `count` items from the `input` provided and move
121    /// ownership into this buffer.
122    ///
123    /// # Panics
124    ///
125    /// This method will panic if `count` is larger than
126    /// `self.remaining_capacity()` or if it's larger than `input.remaining()`.
127    fn move_from(&mut self, input: &mut dyn WriteBuffer<T>, count: usize);
128}
129
130pub(super) struct Extender<'a, B>(pub(super) &'a mut B);
131
132impl<T, B: ReadBuffer<T>> Extend<T> for Extender<'_, B> {
133    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
134        self.0.extend(iter)
135    }
136}
137
138// SAFETY: the `take` implementation below guarantees that the `fun` closure is
139// provided with fully initialized items.
140unsafe impl<T: Send + Sync + 'static> WriteBuffer<T> for Option<T> {
141    fn remaining(&self) -> &[T] {
142        if let Some(me) = self {
143            slice::from_ref(me)
144        } else {
145            &[]
146        }
147    }
148
149    fn skip(&mut self, count: usize) {
150        match count {
151            0 => {}
152            1 => {
153                assert!(self.is_some());
154                *self = None;
155            }
156            _ => panic!("cannot skip more than {} item(s)", self.remaining().len()),
157        }
158    }
159
160    fn take(&mut self, count: usize, fun: &mut dyn FnMut(&[MaybeUninit<T>])) {
161        match count {
162            0 => fun(&mut []),
163            1 => {
164                let mut item = MaybeUninit::new(self.take().unwrap());
165                fun(slice::from_mut(&mut item));
166            }
167            _ => panic!("cannot forget more than {} item(s)", self.remaining().len()),
168        }
169    }
170}
171
172impl<T: Send + Sync + 'static> ReadBuffer<T> for Option<T> {
173    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
174        let mut iter = iter.into_iter();
175        if self.is_none() {
176            *self = iter.next();
177        }
178        assert!(iter.next().is_none());
179    }
180
181    fn remaining_capacity(&self) -> usize {
182        if self.is_some() { 0 } else { 1 }
183    }
184
185    fn move_from(&mut self, input: &mut dyn WriteBuffer<T>, count: usize) {
186        match count {
187            0 => {}
188            1 => {
189                assert!(self.is_none());
190                input.take(1, &mut |slice| {
191                    // SAFETY: Per the `WriteBuffer` trait contract this block
192                    // has ownership of the items in `slice` and they're all
193                    // valid to take.
194                    unsafe {
195                        *self = Some(slice[0].assume_init_read());
196                    }
197                });
198            }
199            _ => panic!(
200                "cannot take more than {} item(s)",
201                self.remaining_capacity()
202            ),
203        }
204    }
205}
206
207/// A `WriteBuffer` implementation, backed by a `Vec`.
208pub struct VecBuffer<T: Send + Sync + 'static> {
209    buffer: Vec<MaybeUninit<T>>,
210    offset: usize,
211}
212
213impl<T: Send + Sync + 'static> VecBuffer<T> {
214    /// Create a new instance with the specified capacity.
215    pub fn with_capacity(capacity: usize) -> Self {
216        Self {
217            buffer: Vec::with_capacity(capacity),
218            offset: 0,
219        }
220    }
221
222    /// Reset the state of this buffer, removing all items and preserving its
223    /// capacity.
224    pub fn reset(&mut self) {
225        self.skip(self.remaining().len());
226        self.buffer.clear();
227        self.offset = 0;
228    }
229}
230
231// SAFETY: the `take` implementation below guarantees that the `fun` closure is
232// provided with fully initialized items due to `self.offset`-and-onwards being
233// always initialized.
234unsafe impl<T: Send + Sync + 'static> WriteBuffer<T> for VecBuffer<T> {
235    fn remaining(&self) -> &[T] {
236        // SAFETY: This relies on the invariant (upheld in the other methods of
237        // this type) that all the elements from `self.offset` onward are
238        // initialized and valid for `self.buffer`.
239        unsafe { mem::transmute::<&[MaybeUninit<T>], &[T]>(&self.buffer[self.offset..]) }
240    }
241
242    fn skip(&mut self, count: usize) {
243        assert!(count <= self.remaining().len());
244        for item in &mut self.buffer[self.offset..][..count] {
245            // Note that the offset is incremented first here to ensure that if
246            // any destructors panic we don't attempt to re-drop the item.
247            self.offset += 1;
248            // SAFETY: See comment in `Self::remaining`
249            unsafe {
250                item.assume_init_drop();
251            }
252        }
253    }
254
255    fn take(&mut self, count: usize, fun: &mut dyn FnMut(&[MaybeUninit<T>])) {
256        assert!(count <= self.remaining().len());
257        // Note that the offset here is incremented before `fun` is called to
258        // ensure that if `fun` panics that the items are still considered
259        // transferred.
260        self.offset += count;
261        fun(&mut self.buffer[self.offset - count..]);
262    }
263}
264
265impl<T: Send + Sync + 'static> From<Vec<T>> for VecBuffer<T> {
266    fn from(buffer: Vec<T>) -> Self {
267        Self {
268            // SAFETY: Transmuting from `Vec<T>` to `Vec<MaybeUninit<T>>` should
269            // be sound for any `T`.
270            buffer: unsafe { mem::transmute::<Vec<T>, Vec<MaybeUninit<T>>>(buffer) },
271            offset: 0,
272        }
273    }
274}
275
276impl<T: Send + Sync + 'static> Drop for VecBuffer<T> {
277    fn drop(&mut self) {
278        self.reset();
279    }
280}
281
282impl<T: Send + Sync + 'static> ReadBuffer<T> for Vec<T> {
283    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
284        Extend::extend(self, iter)
285    }
286
287    fn remaining_capacity(&self) -> usize {
288        self.capacity().checked_sub(self.len()).unwrap()
289    }
290
291    fn move_from(&mut self, input: &mut dyn WriteBuffer<T>, count: usize) {
292        assert!(count <= self.remaining_capacity());
293        input.take(count, &mut |slice| {
294            for item in slice {
295                // SAFETY: Per the `WriteBuffer` implementation contract this
296                // function has exclusive ownership of all items in `slice` so
297                // this is safe to take and transfer them here.
298                self.push(unsafe { item.assume_init_read() });
299            }
300        });
301    }
302}
303
304// SAFETY: the `take` implementation below guarantees that the `fun` closure is
305// provided with fully initialized items.
306#[cfg(feature = "component-model-async-bytes")]
307unsafe impl WriteBuffer<u8> for Cursor<Bytes> {
308    fn remaining(&self) -> &[u8] {
309        &self.get_ref()[usize::try_from(self.position()).unwrap()..]
310    }
311
312    fn skip(&mut self, count: usize) {
313        assert!(
314            count <= self.remaining().len(),
315            "tried to skip {count} with {} remaining",
316            self.remaining().len()
317        );
318        self.set_position(
319            self.position()
320                .checked_add(u64::try_from(count).unwrap())
321                .unwrap(),
322        );
323    }
324
325    fn take(&mut self, count: usize, fun: &mut dyn FnMut(&[MaybeUninit<u8>])) {
326        assert!(count <= self.remaining().len());
327        fun(unsafe_byte_slice(self.remaining()));
328        self.skip(count);
329    }
330}
331
332// SAFETY: the `take` implementation below guarantees that the `fun` closure is
333// provided with fully initialized items.
334#[cfg(feature = "component-model-async-bytes")]
335unsafe impl WriteBuffer<u8> for Cursor<BytesMut> {
336    fn remaining(&self) -> &[u8] {
337        &self.get_ref()[usize::try_from(self.position()).unwrap()..]
338    }
339
340    fn skip(&mut self, count: usize) {
341        assert!(count <= self.remaining().len());
342        self.set_position(
343            self.position()
344                .checked_add(u64::try_from(count).unwrap())
345                .unwrap(),
346        );
347    }
348
349    fn take(&mut self, count: usize, fun: &mut dyn FnMut(&[MaybeUninit<u8>])) {
350        assert!(count <= self.remaining().len());
351        fun(unsafe_byte_slice(self.remaining()));
352        self.skip(count);
353    }
354}
355
356#[cfg(feature = "component-model-async-bytes")]
357impl ReadBuffer<u8> for BytesMut {
358    fn extend<I: IntoIterator<Item = u8>>(&mut self, iter: I) {
359        Extend::extend(self, iter)
360    }
361
362    fn remaining_capacity(&self) -> usize {
363        self.capacity().checked_sub(self.len()).unwrap()
364    }
365
366    fn move_from(&mut self, input: &mut dyn WriteBuffer<u8>, count: usize) {
367        assert!(count <= self.remaining_capacity());
368        input.take(count, &mut |slice| {
369            // SAFETY: per the contract of `WriteBuffer` all the elements of
370            // the input `slice` are fully initialized so this is safe
371            // to reinterpret the slice.
372            let slice = unsafe { mem::transmute::<&[MaybeUninit<u8>], &[u8]>(slice) };
373            self.extend_from_slice(slice);
374        });
375    }
376}
377
378#[cfg(feature = "component-model-async-bytes")]
379fn unsafe_byte_slice(slice: &[u8]) -> &[MaybeUninit<u8>] {
380    // SAFETY: it's always safe to interpret a slice of items as a
381    // possibly-initialized slice of items.
382    unsafe { mem::transmute::<&[u8], &[MaybeUninit<u8>]>(slice) }
383}
wasmtime/runtime/component/concurrent/futures_and_streams/buffers.rs

wasmtime/runtime/component/concurrent/futures_and_streams/
buffers.rs
