wit_bindgen/rt/async_support/
abi_buffer.rs

1use crate::rt::Cleanup;
2use crate::rt::async_support::StreamOps;
3use alloc::alloc::Layout;
4use alloc::vec::Vec;
5use core::mem::{self, MaybeUninit};
6use core::ptr;
7
8/// A helper structure used with a stream to handle the canonical ABI
9/// representation of lists and track partial writes.
10///
11/// This structure is returned whenever a write to a stream completes. This
12/// keeps track of the original buffer used to perform a write (`Vec<T>`) and
13/// additionally tracks any partial writes. Writes can then be resumed with
14/// this buffer again or the partial write can be converted back to `Vec<T>` to
15/// get access to the remaining values.
16///
17/// This value is created through the [`StreamWrite`](super::StreamWrite)
18/// future's return value.
19pub struct AbiBuffer<O: StreamOps> {
20    rust_storage: Vec<MaybeUninit<O::Payload>>,
21    ops: O,
22    alloc: Option<Cleanup>,
23    cursor: usize,
24}
25
26impl<O: StreamOps> AbiBuffer<O> {
27    pub(crate) fn new(mut vec: Vec<O::Payload>, mut ops: O) -> AbiBuffer<O> {
28        // SAFETY: We're converting `Vec<T>` to `Vec<MaybeUninit<T>>`, which
29        // should be safe.
30        let rust_storage = unsafe {
31            let ptr = vec.as_mut_ptr();
32            let len = vec.len();
33            let cap = vec.capacity();
34            mem::forget(vec);
35            Vec::<MaybeUninit<O::Payload>>::from_raw_parts(ptr.cast(), len, cap)
36        };
37
38        // If `lower` is provided then the canonical ABI format is different
39        // from the native format, so all items are converted at this time.
40        //
41        // Note that this is probably pretty inefficient for "big" use cases
42        // but it's hoped that "big" use cases are using `u8` and therefore
43        // skip this entirely.
44        let alloc = if ops.native_abi_matches_canonical_abi() {
45            None
46        } else {
47            let elem_layout = ops.elem_layout();
48            let layout = Layout::from_size_align(
49                elem_layout.size() * rust_storage.len(),
50                elem_layout.align(),
51            )
52            .unwrap();
53            let (mut ptr, cleanup) = Cleanup::new(layout);
54            // SAFETY: All items in `rust_storage` are already initialized so
55            // it should be safe to read them and move ownership into the
56            // canonical ABI format.
57            unsafe {
58                for item in rust_storage.iter() {
59                    let item = item.assume_init_read();
60                    ops.lower(item, ptr);
61                    ptr = ptr.add(elem_layout.size());
62                }
63            }
64            cleanup
65        };
66        AbiBuffer {
67            rust_storage,
68            alloc,
69            ops,
70            cursor: 0,
71        }
72    }
73
74    /// Returns the canonical ABI pointer/length to pass off to a write
75    /// operation.
76    pub(crate) fn abi_ptr_and_len(&self) -> (*const u8, usize) {
77        // If there's no `lower` operation then it means that `T`'s layout is
78        // the same in the canonical ABI so it can be used as-is. In this
79        // situation the list would have been un-tampered with above.
80        if self.ops.native_abi_matches_canonical_abi() {
81            // SAFETY: this should be in-bounds, so it should be safe.
82            let ptr = unsafe { self.rust_storage.as_ptr().add(self.cursor).cast() };
83            let len = self.rust_storage.len() - self.cursor;
84            return (ptr, len.try_into().unwrap());
85        }
86
87        // Othereise when `lower` is present that means that `self.alloc` has
88        // the ABI pointer we should pass along.
89        let ptr = self
90            .alloc
91            .as_ref()
92            .map(|c| c.ptr.as_ptr())
93            .unwrap_or(ptr::null_mut());
94        (
95            // SAFETY: this should be in-bounds, so it should be safe.
96            unsafe { ptr.add(self.cursor * self.ops.elem_layout().size()) },
97            self.rust_storage.len() - self.cursor,
98        )
99    }
100
101    /// Converts this `AbiBuffer<T>` back into a `Vec<T>`
102    ///
103    /// This commit consumes this buffer and yields back unwritten values as a
104    /// `Vec<T>`. The remaining items in `Vec<T>` have not yet been written and
105    /// all written items have been removed from the front of the list.
106    ///
107    /// Note that the backing storage of the returned `Vec<T>` has not changed
108    /// from whe this buffer was created.
109    ///
110    /// Also note that this can be an expensive operation if a partial write
111    /// occurred as this will involve shifting items from the end of the vector
112    /// to the start of the vector.
113    pub fn into_vec(mut self) -> Vec<O::Payload> {
114        self.take_vec()
115    }
116
117    /// Returns the number of items remaining in this buffer.
118    pub fn remaining(&self) -> usize {
119        self.rust_storage.len() - self.cursor
120    }
121
122    /// Advances this buffer by `amt` items.
123    ///
124    /// This signals that `amt` items are no longer going to be yielded from
125    /// `abi_ptr_and_len`. Additionally this will perform any deallocation
126    /// necessary for the starting `amt` items in this list.
127    pub(crate) fn advance(&mut self, amt: usize) {
128        assert!(amt + self.cursor <= self.rust_storage.len());
129        if !self.ops.contains_lists() {
130            self.cursor += amt;
131            return;
132        }
133        let (mut ptr, len) = self.abi_ptr_and_len();
134        assert!(amt <= len);
135        for _ in 0..amt {
136            // Update self.cursor incrementally for exception safety.
137            // When `self.ops.dealloc_lists` panics (which is a user-provided
138            // callback), we can make sure any item before (including the
139            // panic one) will not be dealloced again, and the remaining items
140            // can still get advanced properly.
141            self.cursor += 1;
142            // SAFETY: we're managing the pointer passed to `dealloc_lists` and
143            // it was initialized with a `lower`, and then the pointer
144            // arithmetic should all be in-bounds.
145            unsafe {
146                self.ops.dealloc_lists(ptr.cast_mut());
147                ptr = ptr.add(self.ops.elem_layout().size());
148            }
149        }
150    }
151
152    fn take_vec(&mut self) -> Vec<O::Payload> {
153        // First, if necessary, convert remaining values within `self.alloc`
154        // back into `self.rust_storage`. This is necessary when a lift
155        // operation is available meaning that the representation of `T` is
156        // different in the canonical ABI.
157        //
158        // Note that when `lift` is provided then when this original
159        // `AbiBuffer` was created it moved ownership of all values from the
160        // original vector into the `alloc` value. This is the reverse
161        // operation, moving all the values back into the vector.
162        if !self.ops.native_abi_matches_canonical_abi() {
163            let (mut ptr, mut len) = self.abi_ptr_and_len();
164            // SAFETY: this should be safe as `lift` is operating on values that
165            // were initialized with a previous `lower`, and the pointer
166            // arithmetic here should all be in-bounds.
167            unsafe {
168                for dst in self.rust_storage[self.cursor..].iter_mut() {
169                    dst.write(self.ops.lift(ptr.cast_mut()));
170                    ptr = ptr.add(self.ops.elem_layout().size());
171                    len -= 1;
172                }
173                assert_eq!(len, 0);
174            }
175        }
176
177        // Next extract the rust storage and zero out this struct's fields.
178        // This is also the location where a "shift" happens to remove items
179        // from the beginning of the returned vector as those have already been
180        // transferred somewhere else.
181        let mut storage = mem::take(&mut self.rust_storage);
182        storage.drain(..self.cursor);
183        self.cursor = 0;
184        self.alloc = None;
185
186        // SAFETY: we're casting `Vec<MaybeUninit<T>>` here to `Vec<T>`. The
187        // elements were either always initialized (`lift` is `None`) or we just
188        // re-initialized them above from `self.alloc`.
189        unsafe {
190            let ptr = storage.as_mut_ptr();
191            let len = storage.len();
192            let cap = storage.capacity();
193            mem::forget(storage);
194            Vec::<O::Payload>::from_raw_parts(ptr.cast(), len, cap)
195        }
196    }
197}
198
199impl<O> Drop for AbiBuffer<O>
200where
201    O: StreamOps,
202{
203    fn drop(&mut self) {
204        let _ = self.take_vec();
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use crate::rt::async_support::StreamVtable;
212    use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
213    use std::vec;
214
215    extern "C" fn cancel(_: u32) -> u32 {
216        todo!()
217    }
218    extern "C" fn drop(_: u32) {
219        todo!()
220    }
221    extern "C" fn new() -> u64 {
222        todo!()
223    }
224    extern "C" fn start_read(_: u32, _: *mut u8, _: usize) -> u32 {
225        todo!()
226    }
227    extern "C" fn start_write(_: u32, _: *const u8, _: usize) -> u32 {
228        todo!()
229    }
230
231    static BLANK: StreamVtable<u8> = StreamVtable {
232        cancel_read: cancel,
233        cancel_write: cancel,
234        drop_readable: drop,
235        drop_writable: drop,
236        dealloc_lists: None,
237        lift: None,
238        lower: None,
239        layout: unsafe { Layout::from_size_align_unchecked(1, 1) },
240        new,
241        start_read,
242        start_write,
243    };
244
245    #[test]
246    fn blank_advance_to_end() {
247        let mut buffer = AbiBuffer::new(vec![1, 2, 3, 4], &BLANK);
248        assert_eq!(buffer.remaining(), 4);
249        buffer.advance(1);
250        assert_eq!(buffer.remaining(), 3);
251        buffer.advance(2);
252        assert_eq!(buffer.remaining(), 1);
253        buffer.advance(1);
254        assert_eq!(buffer.remaining(), 0);
255        assert_eq!(buffer.into_vec(), []);
256    }
257
258    #[test]
259    fn blank_advance_partial() {
260        let buffer = AbiBuffer::new(vec![1, 2, 3, 4], &BLANK);
261        assert_eq!(buffer.into_vec(), [1, 2, 3, 4]);
262        let mut buffer = AbiBuffer::new(vec![1, 2, 3, 4], &BLANK);
263        buffer.advance(1);
264        assert_eq!(buffer.into_vec(), [2, 3, 4]);
265        let mut buffer = AbiBuffer::new(vec![1, 2, 3, 4], &BLANK);
266        buffer.advance(1);
267        buffer.advance(2);
268        assert_eq!(buffer.into_vec(), [4]);
269    }
270
271    #[test]
272    fn blank_ptr_eq() {
273        let mut buf = vec![1, 2, 3, 4];
274        let ptr = buf.as_mut_ptr();
275        let mut buffer = AbiBuffer::new(buf, &BLANK);
276        let (a, b) = buffer.abi_ptr_and_len();
277        assert_eq!(a, ptr);
278        assert_eq!(b, 4);
279        unsafe {
280            assert_eq!(std::slice::from_raw_parts(a, b), [1, 2, 3, 4]);
281        }
282
283        buffer.advance(1);
284        let (a, b) = buffer.abi_ptr_and_len();
285        assert_eq!(a, ptr.wrapping_add(1));
286        assert_eq!(b, 3);
287        unsafe {
288            assert_eq!(std::slice::from_raw_parts(a, b), [2, 3, 4]);
289        }
290
291        buffer.advance(2);
292        let (a, b) = buffer.abi_ptr_and_len();
293        assert_eq!(a, ptr.wrapping_add(3));
294        assert_eq!(b, 1);
295        unsafe {
296            assert_eq!(std::slice::from_raw_parts(a, b), [4]);
297        }
298
299        let ret = buffer.into_vec();
300        assert_eq!(ret, [4]);
301        assert_eq!(ret.as_ptr(), ptr);
302    }
303
304    #[derive(PartialEq, Eq, Debug)]
305    struct B(u8);
306
307    static OP: StreamVtable<B> = StreamVtable {
308        cancel_read: cancel,
309        cancel_write: cancel,
310        drop_readable: drop,
311        drop_writable: drop,
312        dealloc_lists: Some(|_ptr| {}),
313        lift: Some(|ptr| unsafe { B(*ptr - 1) }),
314        lower: Some(|b, ptr| unsafe {
315            *ptr = b.0 + 1;
316        }),
317        layout: unsafe { Layout::from_size_align_unchecked(1, 1) },
318        new,
319        start_read,
320        start_write,
321    };
322
323    #[test]
324    fn op_advance_to_end() {
325        let mut buffer = AbiBuffer::new(vec![B(1), B(2), B(3), B(4)], &OP);
326        assert_eq!(buffer.remaining(), 4);
327        buffer.advance(1);
328        assert_eq!(buffer.remaining(), 3);
329        buffer.advance(2);
330        assert_eq!(buffer.remaining(), 1);
331        buffer.advance(1);
332        assert_eq!(buffer.remaining(), 0);
333        assert_eq!(buffer.into_vec(), []);
334    }
335
336    #[test]
337    fn op_advance_partial() {
338        let buffer = AbiBuffer::new(vec![B(1), B(2), B(3), B(4)], &OP);
339        assert_eq!(buffer.into_vec(), [B(1), B(2), B(3), B(4)]);
340        let mut buffer = AbiBuffer::new(vec![B(1), B(2), B(3), B(4)], &OP);
341        buffer.advance(1);
342        assert_eq!(buffer.into_vec(), [B(2), B(3), B(4)]);
343        let mut buffer = AbiBuffer::new(vec![B(1), B(2), B(3), B(4)], &OP);
344        buffer.advance(1);
345        buffer.advance(2);
346        assert_eq!(buffer.into_vec(), [B(4)]);
347    }
348
349    #[test]
350    fn op_ptrs() {
351        let mut buf = vec![B(1), B(2), B(3), B(4)];
352        let ptr = buf.as_mut_ptr().cast::<u8>();
353        let mut buffer = AbiBuffer::new(buf, &OP);
354        let (a, b) = buffer.abi_ptr_and_len();
355        let base = a;
356        assert_ne!(a, ptr);
357        assert_eq!(b, 4);
358        unsafe {
359            assert_eq!(std::slice::from_raw_parts(a, b), [2, 3, 4, 5]);
360        }
361
362        buffer.advance(1);
363        let (a, b) = buffer.abi_ptr_and_len();
364        assert_ne!(a, ptr.wrapping_add(1));
365        assert_eq!(a, base.wrapping_add(1));
366        assert_eq!(b, 3);
367        unsafe {
368            assert_eq!(std::slice::from_raw_parts(a, b), [3, 4, 5]);
369        }
370
371        buffer.advance(2);
372        let (a, b) = buffer.abi_ptr_and_len();
373        assert_ne!(a, ptr.wrapping_add(3));
374        assert_eq!(a, base.wrapping_add(3));
375        assert_eq!(b, 1);
376        unsafe {
377            assert_eq!(std::slice::from_raw_parts(a, b), [5]);
378        }
379
380        let ret = buffer.into_vec();
381        assert_eq!(ret, [B(4)]);
382        assert_eq!(ret.as_ptr(), ptr.cast());
383    }
384
385    #[test]
386    fn dealloc_lists() {
387        static DEALLOCS: AtomicUsize = AtomicUsize::new(0);
388        static OP: StreamVtable<B> = StreamVtable {
389            cancel_read: cancel,
390            cancel_write: cancel,
391            drop_readable: drop,
392            drop_writable: drop,
393            dealloc_lists: Some(|ptr| {
394                let prev = DEALLOCS.fetch_add(1, Relaxed);
395                assert_eq!(unsafe { usize::from(*ptr) }, prev + 1);
396            }),
397            lift: Some(|ptr| unsafe { B(*ptr) }),
398            lower: Some(|b, ptr| unsafe {
399                *ptr = b.0;
400            }),
401            layout: unsafe { Layout::from_size_align_unchecked(1, 1) },
402            new,
403            start_read,
404            start_write,
405        };
406
407        assert_eq!(DEALLOCS.load(Relaxed), 0);
408        let buf = vec![B(1), B(2), B(3), B(4)];
409        let mut buffer = AbiBuffer::new(buf, &OP);
410        assert_eq!(DEALLOCS.load(Relaxed), 0);
411        buffer.abi_ptr_and_len();
412        assert_eq!(DEALLOCS.load(Relaxed), 0);
413
414        buffer.advance(1);
415        assert_eq!(DEALLOCS.load(Relaxed), 1);
416        buffer.abi_ptr_and_len();
417        assert_eq!(DEALLOCS.load(Relaxed), 1);
418        buffer.advance(2);
419        assert_eq!(DEALLOCS.load(Relaxed), 3);
420        buffer.abi_ptr_and_len();
421        assert_eq!(DEALLOCS.load(Relaxed), 3);
422        buffer.into_vec();
423        assert_eq!(DEALLOCS.load(Relaxed), 3);
424    }
425}
wit_bindgen/rt/async_support/abi_buffer.rs

wit_bindgen/rt/async_support/
abi_buffer.rs
