From 86412b747ae6b0ba5409c9afc015ef94b64df020 Mon Sep 17 00:00:00 2001
From: Philip Craig <philipjcraig@gmail.com>
Date: Tue, 4 May 2021 14:41:53 +1000
Subject: [PATCH] pod: add from_bytes_mut and slice_from_bytes_mut

---
 src/pod.rs | 53 +++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 49 insertions(+), 4 deletions(-)

diff --git a/src/pod.rs b/src/pod.rs
index 6a06457d..d9cf6f80 100644
--- a/src/pod.rs
+++ b/src/pod.rs
@@ -27,12 +27,12 @@ pub unsafe trait Pod: Copy + 'static {}
 /// Returns the type and the tail of the slice.
 #[inline]
 pub fn from_bytes<T: Pod>(data: &[u8]) -> Result<(&T, &[u8])> {
+    let size = mem::size_of::<T>();
+    let tail = data.get(size..).ok_or(())?;
     let ptr = data.as_ptr();
     if (ptr as usize) % mem::align_of::<T>() != 0 {
         return Err(());
     }
-    let size = mem::size_of::<T>();
-    let tail = data.get(size..).ok_or(())?;
     // Safety:
     // The alignment and size are checked by this function.
     // The Pod trait ensures the type is valid to cast from bytes.
@@ -40,17 +40,38 @@ pub fn from_bytes<T: Pod>(data: &[u8]) -> Result<(&T, &[u8])> {
     Ok((val, tail))
 }
 
+/// Cast a mutable byte slice to a `Pod` type.
+///
+/// Returns the type and the tail of the slice.
+#[inline]
+pub fn from_bytes_mut<T: Pod>(data: &mut [u8]) -> Result<(&mut T, &mut [u8])> {
+    let size = mem::size_of::<T>();
+    if size > data.len() {
+        return Err(());
+    }
+    let (data, tail) = data.split_at_mut(size);
+    let ptr = data.as_mut_ptr();
+    if (ptr as usize) % mem::align_of::<T>() != 0 {
+        return Err(());
+    }
+    // Safety:
+    // The alignment and size are checked by this function.
+    // The Pod trait ensures the type is valid to cast from bytes.
+    let val = unsafe { &mut *ptr.cast() };
+    Ok((val, tail))
+}
+
 /// Cast a byte slice to a slice of a `Pod` type.
 ///
 /// Returns the type slice and the tail of the byte slice.
 #[inline]
 pub fn slice_from_bytes<T: Pod>(data: &[u8], count: usize) -> Result<(&[T], &[u8])> {
+    let size = count.checked_mul(mem::size_of::<T>()).ok_or(())?;
+    let tail = data.get(size..).ok_or(())?;
     let ptr = data.as_ptr();
     if (ptr as usize) % mem::align_of::<T>() != 0 {
         return Err(());
     }
-    let size = count.checked_mul(mem::size_of::<T>()).ok_or(())?;
-    let tail = data.get(size..).ok_or(())?;
     // Safety:
     // The alignment and size are checked by this function.
     // The Pod trait ensures the type is valid to cast from bytes.
@@ -58,6 +79,30 @@ pub fn slice_from_bytes<T: Pod>(data: &[u8], count: usize) -> Result<(&[T], &[u8
     Ok((slice, tail))
 }
 
+/// Cast a mutable byte slice to a slice of a `Pod` type.
+///
+/// Returns the type slice and the tail of the byte slice.
+#[inline]
+pub fn slice_from_bytes_mut<T: Pod>(
+    data: &mut [u8],
+    count: usize,
+) -> Result<(&mut [T], &mut [u8])> {
+    let size = count.checked_mul(mem::size_of::<T>()).ok_or(())?;
+    if size > data.len() {
+        return Err(());
+    }
+    let (data, tail) = data.split_at_mut(size);
+    let ptr = data.as_mut_ptr();
+    if (ptr as usize) % mem::align_of::<T>() != 0 {
+        return Err(());
+    }
+    // Safety:
+    // The alignment and size are checked by this function.
+    // The Pod trait ensures the type is valid to cast from bytes.
+    let slice = unsafe { slice::from_raw_parts_mut(ptr.cast(), count) };
+    Ok((slice, tail))
+}
+
 /// Cast a `Pod` type to a byte slice.
 #[inline]
 pub fn bytes_of<T: Pod>(val: &T) -> &[u8] {