From d790c043a8e0b588d1c9f5a3120a63cb6d311da5 Mon Sep 17 00:00:00 2001
From: Thom Chiovoloni <thom@shift.click>
Date: Tue, 5 Sep 2023 12:46:04 -0700
Subject: [PATCH] Add functions for writing zeroed bytes to `&mut impl
 Zeroable` and `&mut [impl Zeroable]` (#193)

* Add functions for writing zeroed bytes to `&mut impl Zeroable` and `&mut [impl Zeroable]`

* Support `T: !Copy` in `fill_zero`/`write_zero`

* Zero bytes in the drop guard for write_zero

* Update src/lib.rs

Co-authored-by: Alphyr <47725341+a1phyr@users.noreply.github.com>

---------

Co-authored-by: Lokathor <zefria@gmail.com>
Co-authored-by: Alphyr <47725341+a1phyr@users.noreply.github.com>
---
 src/lib.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index a51ed8a..ebfc195 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -408,3 +408,47 @@ pub fn try_cast_slice_mut<
 ) -> Result<&mut [B], PodCastError> {
   unsafe { internal::try_cast_slice_mut(a) }
 }
+
+/// Fill all bytes of `target` with zeroes (see [`Zeroable`]).
+///
+/// This is similar to `*target = Zeroable::zeroed()`, but guarantees that any
+/// padding bytes in `target` are zeroed as well.
+///
+/// See also [`fill_zero`], if you have a slice rather than a single value.
+#[inline]
+pub fn write_zero<T: Zeroable>(target: &mut T) {
+  struct EnsureZeroWrite<T>(*mut T);
+  impl<T> Drop for EnsureZeroWrite<T> {
+    #[inline(always)]
+    fn drop(&mut self) {
+      unsafe {
+        core::ptr::write_bytes(self.0, 0u8, 1);
+      }
+    }
+  }
+  unsafe {
+    let guard = EnsureZeroWrite(target);
+    core::ptr::drop_in_place(guard.0);
+    drop(guard);
+  }
+}
+
+/// Fill all bytes of `slice` with zeroes (see [`Zeroable`]).
+///
+/// This is similar to `slice.fill(Zeroable::zeroed())`, but guarantees that any
+/// padding bytes in `slice` are zeroed as well.
+///
+/// See also [`write_zero`], which zeroes all bytes of a single value rather
+/// than a slice.
+#[inline]
+pub fn fill_zero<T: Zeroable>(slice: &mut [T]) {
+  if core::mem::needs_drop::<T>() {
+    // If `T` needs to be dropped then we have to do this one item at a time, in
+    // case one of the intermediate drops does a panic.
+    slice.iter_mut().for_each(write_zero);
+  } else {
+    // Otherwise we can be really fast and just fill everthing with zeros.
+    let len = core::mem::size_of_val::<[T]>(slice);
+    unsafe { core::ptr::write_bytes(slice.as_mut_ptr() as *mut u8, 0u8, len) }
+  }
+}