diff --git a/crates/proc_macro_srv/src/proc_macro/bridge/buffer.rs b/crates/proc_macro_srv/src/proc_macro/bridge/buffer.rs index dae6ff1d16c3..41beeb98c4a0 100644 --- a/crates/proc_macro_srv/src/proc_macro/bridge/buffer.rs +++ b/crates/proc_macro_srv/src/proc_macro/bridge/buffer.rs @@ -42,7 +42,7 @@ pub struct Buffer { data: *mut T, len: usize, capacity: usize, - extend_from_slice: extern "C" fn(Buffer, Slice<'_, T>) -> Buffer, + reserve: extern "C" fn(Buffer, usize) -> Buffer, drop: extern "C" fn(Buffer), } @@ -82,17 +82,28 @@ impl Buffer { } pub(super) fn extend_from_slice(&mut self, xs: &[T]) { - // Fast path to avoid going through an FFI call. - if let Some(final_len) = self.len.checked_add(xs.len()) { - if final_len <= self.capacity { - let dst = unsafe { slice::from_raw_parts_mut(self.data, self.capacity) }; - dst[self.len..][..xs.len()].copy_from_slice(xs); - self.len = final_len; - return; - } + if xs.len() > self.capacity.wrapping_sub(self.len) { + let b = self.take(); + *self = (b.reserve)(b, xs.len()); + } + unsafe { + xs.as_ptr().copy_to_nonoverlapping(self.data.add(self.len), xs.len()); + self.len += xs.len(); + } + } + + pub(super) fn push(&mut self, v: T) { + // The code here is taken from Vec::push, and we know that reserve() + // 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); + } + unsafe { + *self.data.add(self.len) = v; + self.len += 1; } - let b = self.take(); - *self = (b.extend_from_slice)(b, Slice::from(xs)); } } @@ -134,9 +145,9 @@ impl From> for Buffer { } } - extern "C" fn extend_from_slice(b: Buffer, xs: Slice<'_, T>) -> Buffer { + extern "C" fn reserve(b: Buffer, additional: usize) -> Buffer { let mut v = to_vec(b); - v.extend_from_slice(&xs); + v.reserve(additional); Buffer::from(v) } @@ -144,6 +155,6 @@ impl From> for Buffer { mem::drop(to_vec(b)); } - Buffer { data, len, capacity, extend_from_slice, drop } + Buffer { data, len, capacity, reserve, drop } } }