diff --git a/library/proc_macro/src/bridge/buffer.rs b/library/proc_macro/src/bridge/buffer.rs index aeecbd496621b..717201aef102b 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), } @@ -78,18 +78,44 @@ impl Buffer { mem::take(self) } + // We have the array method separate from extending from a slice. This is + // because in the case of small arrays, codegen can be more efficient + // (avoiding a memmove call). With extend_from_slice, LLVM at least + // currently is not able to make that optimization. + pub(super) fn extend_from_array(&mut self, xs: &[T; N]) { + if xs.len() > (self.capacity - 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 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 - 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)); } } @@ -131,9 +157,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) } @@ -141,6 +167,6 @@ impl From> for Buffer { mem::drop(to_vec(b)); } - Buffer { data, len, capacity, extend_from_slice, drop } + Buffer { data, len, capacity, reserve, drop } } } diff --git a/library/proc_macro/src/bridge/rpc.rs b/library/proc_macro/src/bridge/rpc.rs index 5c2f9ec9848dd..588e6ded0f480 100644 --- a/library/proc_macro/src/bridge/rpc.rs +++ b/library/proc_macro/src/bridge/rpc.rs @@ -27,7 +27,7 @@ macro_rules! rpc_encode_decode { (le $ty:ty) => { impl Encode for $ty { fn encode(self, w: &mut Writer, _: &mut S) { - w.write_all(&self.to_le_bytes()).unwrap(); + w.extend_from_array(&self.to_le_bytes()); } } @@ -114,7 +114,7 @@ impl DecodeMut<'_, '_, S> for () { impl Encode for u8 { fn encode(self, w: &mut Writer, _: &mut S) { - w.write_all(&[self]).unwrap(); + w.push(self); } }