Skip to main content

proc_macro/bridge/
buffer.rs

1//! Buffer management for same-process client<->server communication.
2
3use std::mem::{self, ManuallyDrop};
4use std::ops::Deref;
5use std::slice;
6
7#[repr(C)]
8pub struct Buffer {
9    data: *mut u8,
10    len: usize,
11    capacity: usize,
12    reserve: extern "C" fn(Buffer, usize) -> Buffer,
13    drop: extern "C" fn(Buffer),
14}
15
16unsafe impl Sync for Buffer {}
17unsafe impl Send for Buffer {}
18
19impl Default for Buffer {
20    #[inline]
21    fn default() -> Self {
22        Self::from(vec![])
23    }
24}
25
26impl Deref for Buffer {
27    type Target = [u8];
28    #[inline]
29    fn deref(&self) -> &[u8] {
30        unsafe { slice::from_raw_parts(self.data as *const u8, self.len) }
31    }
32}
33
34impl Buffer {
35    #[inline]
36    pub(super) fn new() -> Self {
37        Self::default()
38    }
39
40    #[inline]
41    pub(super) fn clear(&mut self) {
42        self.len = 0;
43    }
44
45    #[inline]
46    pub(super) fn take(&mut self) -> Self {
47        mem::take(self)
48    }
49
50    // We have the array method separate from extending from a slice. This is
51    // because in the case of small arrays, codegen can be more efficient
52    // (avoiding a memmove call). With extend_from_slice, LLVM at least
53    // currently is not able to make that optimization.
54    #[inline]
55    pub(super) fn extend_from_array<const N: usize>(&mut self, xs: &[u8; N]) {
56        if xs.len() > (self.capacity - self.len) {
57            let b = self.take();
58            *self = (b.reserve)(b, xs.len());
59        }
60        unsafe {
61            xs.as_ptr().copy_to_nonoverlapping(self.data.add(self.len), xs.len());
62            self.len += xs.len();
63        }
64    }
65
66    #[inline]
67    pub(super) fn extend_from_slice(&mut self, xs: &[u8]) {
68        if xs.len() > (self.capacity - self.len) {
69            let b = self.take();
70            *self = (b.reserve)(b, xs.len());
71        }
72        unsafe {
73            xs.as_ptr().copy_to_nonoverlapping(self.data.add(self.len), xs.len());
74            self.len += xs.len();
75        }
76    }
77
78    #[inline]
79    pub(super) fn push(&mut self, v: u8) {
80        // The code here is taken from Vec::push, and we know that reserve()
81        // will panic if we're exceeding isize::MAX bytes and so there's no need
82        // to check for overflow.
83        if self.len == self.capacity {
84            let b = self.take();
85            *self = (b.reserve)(b, 1);
86        }
87        unsafe {
88            *self.data.add(self.len) = v;
89            self.len += 1;
90        }
91    }
92}
93
94impl Drop for Buffer {
95    #[inline]
96    fn drop(&mut self) {
97        let b = self.take();
98        (b.drop)(b);
99    }
100}
101
102impl From<Vec<u8>> for Buffer {
103    fn from(v: Vec<u8>) -> Self {
104        let (data, len, capacity) = v.into_raw_parts();
105
106        // This utility function is nested in here because it can *only*
107        // be safely called on `Buffer`s created by *this* `proc_macro`.
108        fn to_vec(b: Buffer) -> Vec<u8> {
109            unsafe {
110                let b = ManuallyDrop::new(b);
111                Vec::from_raw_parts(b.data, b.len, b.capacity)
112            }
113        }
114
115        extern "C" fn reserve(b: Buffer, additional: usize) -> Buffer {
116            let mut v = to_vec(b);
117            v.reserve(additional);
118            Buffer::from(v)
119        }
120
121        extern "C" fn drop(b: Buffer) {
122            mem::drop(to_vec(b));
123        }
124
125        Buffer { data, len, capacity, reserve, drop }
126    }
127}