[PERF] directly manage Buffer memory without using Vec

library/proc_macro/src/bridge/buffer.rs

@@ -1,17 +1,16 @@
 //! Buffer management for same-process client<->server communication.
 
+use std::alloc::{self, Layout};
 use std::io::{self, Write};
-use std::mem::{self, ManuallyDrop};
 use std::ops::{Deref, DerefMut};
-use std::slice;
+use std::ptr::null_mut;
+use std::{mem, slice};
 
 #[repr(C)]
 pub struct Buffer {
     data: *mut u8,
     len: usize,
     capacity: usize,
-    reserve: extern "C" fn(Buffer, usize) -> Buffer,
-    drop: extern "C" fn(Buffer),
 }
 
 unsafe impl Sync for Buffer {}
@@ -20,7 +19,7 @@ unsafe impl Send for Buffer {}
 impl Default for Buffer {
     #[inline]
     fn default() -> Self {
-        Self::from(vec![])
+        Self { data: null_mut(), len: 0, capacity: 0 }
     }
 }
 
@@ -62,8 +61,7 @@ impl Buffer {
     #[inline]
     pub(super) fn extend_from_array<const N: usize>(&mut self, xs: &[u8; N]) {
         if xs.len() > (self.capacity - self.len) {
-            let b = self.take();
-            *self = (b.reserve)(b, xs.len());
+            self.reserve(xs.len());
         }
         unsafe {
             xs.as_ptr().copy_to_nonoverlapping(self.data.add(self.len), xs.len());
@@ -74,8 +72,7 @@ impl Buffer {
     #[inline]
     pub(super) fn extend_from_slice(&mut self, xs: &[u8]) {
         if xs.len() > (self.capacity - self.len) {
-            let b = self.take();
-            *self = (b.reserve)(b, xs.len());
+            self.reserve(xs.len());
         }
         unsafe {
             xs.as_ptr().copy_to_nonoverlapping(self.data.add(self.len), xs.len());
@@ -89,14 +86,34 @@ impl Buffer {
         // will panic if we're exceeding isize::MAX bytes and so there's no need
         // to check for overflow.
         if self.len == self.capacity {
-            let b = self.take();
-            *self = (b.reserve)(b, 1);
+            self.reserve(1);
         }
         unsafe {
             *self.data.add(self.len) = v;
             self.len += 1;
         }
     }
+
+    #[inline]
+    fn layout(&self) -> Layout {
+        Layout::array::<u8>(self.capacity).unwrap()
+    }
+
+    #[inline(never)]
+    fn reserve(&mut self, amount: usize) {
+        debug_assert!(amount > 0);
+        self.capacity += amount;
+        assert!(self.capacity < isize::MAX as usize);
+        let layout = self.layout();
+        self.data = if self.data.is_null() {
+            unsafe { alloc::alloc(layout) }
+        } else {
+            unsafe { alloc::realloc(self.data, layout, self.capacity) }
+        };
+        if self.data.is_null() {
+            alloc::handle_alloc_error(layout);
+        }
+    }
 }
 
 impl Write for Buffer {
@@ -121,35 +138,10 @@ impl Write for Buffer {
 impl Drop for Buffer {
     #[inline]
     fn drop(&mut self) {
-        let b = self.take();
-        (b.drop)(b);
-    }
-}
-
-impl From<Vec<u8>> for Buffer {
-    fn from(v: Vec<u8>) -> Self {
-        let mut v = ManuallyDrop::new(v);
-        let (data, len, capacity) = (v.as_mut_ptr(), v.len(), v.capacity());
-
-        // This utility function is nested in here because it can *only*
-        // be safely called on `Buffer`s created by *this* `proc_macro`.
-        fn to_vec(b: Buffer) -> Vec<u8> {
+        if !self.data.is_null() {
             unsafe {
-                let b = ManuallyDrop::new(b);
-                Vec::from_raw_parts(b.data, b.len, b.capacity)
+                alloc::dealloc(self.data, self.layout());
             }
         }
-
-        extern "C" fn reserve(b: Buffer, additional: usize) -> Buffer {
-            let mut v = to_vec(b);
-            v.reserve(additional);
-            Buffer::from(v)
-        }
-
-        extern "C" fn drop(b: Buffer) {
-            mem::drop(to_vec(b));
-        }
-
-        Buffer { data, len, capacity, reserve, drop }
     }
 }