diff --git a/library/proc_macro/src/bridge/buffer.rs b/library/proc_macro/src/bridge/buffer.rs index 26d210a01c5..a2030b9b8bf 100644 --- a/library/proc_macro/src/bridge/buffer.rs +++ b/library/proc_macro/src/bridge/buffer.rs @@ -39,7 +39,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), } @@ -79,32 +79,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(); } - let b = self.take(); - *self = (b.extend_from_slice)(b, Slice::from(xs)); } pub(super) fn push(&mut self, v: T) { - // Fast path to avoid going through an FFI call. - if let Some(final_len) = self.len.checked_add(1) { - if final_len <= self.capacity { - unsafe { - *self.data.add(self.len) = v; - } - self.len = final_len; - return; - } + // 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(std::slice::from_ref(&v))); } } @@ -146,9 +142,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) } @@ -156,6 +152,6 @@ impl From> for Buffer { mem::drop(to_vec(b)); } - Buffer { data, len, capacity, extend_from_slice, drop } + Buffer { data, len, capacity, reserve, drop } } }