proc_macro/bridge/
buffer.rs

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

proc_macro/bridge/
buffer.rs
