From c1b10677528985664e88506d17654bd3f483e8ca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marin=20Ver=C5=A1i=C4=87?= <marin.versic101@gmail.com>
Date: Fri, 12 Jan 2024 07:50:47 +0300
Subject: [PATCH] Add implementation of Decode for Box<str> and Box<[T]>
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Marin Veršić <marin.versic101@gmail.com>
---
 src/codec.rs | 53 ++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 45 insertions(+), 8 deletions(-)

diff --git a/src/codec.rs b/src/codec.rs
index 08f4640b..8f44003c 100644
--- a/src/codec.rs
+++ b/src/codec.rs
@@ -537,9 +537,7 @@ impl<T> WrapperTypeDecode for Box<T> {
 		} else {
 
 			// SAFETY: Layout has a non-zero size so calling this is safe.
-			let ptr: *mut u8 = unsafe {
-				crate::alloc::alloc::alloc(layout)
-			};
+			let ptr = unsafe { crate::alloc::alloc::alloc(layout) };
 
 			if ptr.is_null() {
 				crate::alloc::alloc::handle_alloc_error(layout);
@@ -554,25 +552,42 @@ impl<T> WrapperTypeDecode for Box<T> {
 		//
 		//         Constructing a `Box` from `NonNull::dangling` is also always safe as long
 		//         as the underlying type is zero-sized.
-		let mut boxed: Box<MaybeUninit<T>> = unsafe { Box::from_raw(ptr) };
-
+		let mut boxed = unsafe { Box::from_raw(ptr) };
 		T::decode_into(input, &mut boxed)?;
 
 		// Decoding succeeded, so let's get rid of `MaybeUninit`.
 		//
 		// TODO: Use `Box::assume_init` once that's stable.
-		let ptr: *mut MaybeUninit<T> = Box::into_raw(boxed);
-		let ptr: *mut T = ptr.cast();
+		let ptr: *mut T = Box::into_raw(boxed).cast();
 
 		// SAFETY: `MaybeUninit` doesn't affect the memory layout, so casting the pointer back
 		//         into a `Box` is safe.
-		let boxed: Box<T> = unsafe { Box::from_raw(ptr) };
+		let boxed = unsafe { Box::from_raw(ptr) };
 
 		input.ascend_ref();
 		Ok(boxed)
 	}
 }
 
+impl<T: Decode> Decode for Box<[T]> {
+	fn decode<I: Input>(input: &mut I) -> Result<Self, Error> {
+		Ok(Vec::decode(input)?.into_boxed_slice())
+	}
+}
+
+impl Decode for Box<str> {
+	fn decode<I: Input>(input: &mut I) -> Result<Self, Error> {
+		// Guaranteed to create a Vec with capacity == len
+		let vec = Vec::from(Box::<[u8]>::decode(input)?);
+		// Guaranteed not to reallocate the vec, only transmute to String
+		let str = String::from_utf8(vec).map_err(|_| "Invalid utf8 sequence")?;
+
+		// At this point we have String with capacity == len,
+		// therefore String::into_boxed_str will not reallocate
+		Ok(str.into_boxed_str())
+	}
+}
+
 impl<T> WrapperTypeDecode for Rc<T> {
 	type Wrapped = T;
 
@@ -1612,6 +1627,28 @@ mod tests {
 		assert_eq!((x, y), Decode::decode(&mut &encoded[..]).unwrap());
 	}
 
+	#[test]
+	fn boxed_str_works() {
+		let s = "Hello world".to_owned();
+		let b = s.clone().into_boxed_str();
+
+		let encoded = b.encode();
+		assert_eq!(s.encode(), encoded);
+
+		assert_eq!(*b, *Box::<str>::decode(&mut &encoded[..]).unwrap());
+	}
+
+	#[test]
+	fn boxed_slice_works() {
+		let v = vec![1u32, 2, 3, 4, 5, 6];
+		let b = v.clone().into_boxed_slice();
+
+		let encoded = b.encode();
+		assert_eq!(v.encode(), encoded);
+
+		assert_eq!(*b, *Box::<[u32]>::decode(&mut &b.encode()[..]).unwrap());
+	}
+
 	#[test]
 	fn cow_works() {
 		let x = &[1u32, 2, 3, 4, 5, 6][..];