From 5ad68f486d909df01929ecd14c3aca54c40533dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= <developer@crozet.re>
Date: Sun, 23 Jun 2024 11:46:20 +0200
Subject: [PATCH] Introduce Storage::forget_elements() to fix memory leak in
 Matrix::generic_resize() (#1416)

* Add Storage::forget

* Adjust implementations of Reallocator to use Storage::forget

* Fix formatting

* Rename forget to forget_elements and add safety comments

* Update comments in Reallocator implementations

---------

Co-authored-by: Nick Mertin <nickmertin@gmail.com>
---
 src/base/array_storage.rs     |  6 ++++++
 src/base/default_allocator.rs | 14 +++++++-------
 src/base/matrix_view.rs       |  5 +++++
 src/base/storage.rs           |  3 +++
 src/base/vec_storage.rs       | 22 ++++++++++++++++++++++
 5 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/src/base/array_storage.rs b/src/base/array_storage.rs
index 37019e6d6..15405a6a2 100644
--- a/src/base/array_storage.rs
+++ b/src/base/array_storage.rs
@@ -130,6 +130,12 @@ where
     {
         self.clone()
     }
+
+    #[inline]
+    fn forget_elements(self) {
+        // No additional cleanup required.
+        std::mem::forget(self);
+    }
 }
 
 unsafe impl<T, const R: usize, const C: usize> RawStorageMut<T, Const<R>, Const<C>>
diff --git a/src/base/default_allocator.rs b/src/base/default_allocator.rs
index 97278b548..3b5d12811 100644
--- a/src/base/default_allocator.rs
+++ b/src/base/default_allocator.rs
@@ -15,7 +15,7 @@ use crate::base::array_storage::ArrayStorage;
 use crate::base::dimension::Dim;
 #[cfg(any(feature = "alloc", feature = "std"))]
 use crate::base::dimension::{DimName, Dyn};
-use crate::base::storage::{RawStorage, RawStorageMut};
+use crate::base::storage::{RawStorage, RawStorageMut, Storage};
 #[cfg(any(feature = "std", feature = "alloc"))]
 use crate::base::vec_storage::VecStorage;
 use crate::base::Scalar;
@@ -210,8 +210,8 @@ where
 
         // Safety:
         // - We don’t care about dropping elements because the caller is responsible for dropping things.
-        // - We forget `buf` so that we don’t drop the other elements.
-        std::mem::forget(buf);
+        // - We forget `buf` so that we don’t drop the other elements, but ensure the buffer itself is cleaned up.
+        buf.forget_elements();
 
         res
     }
@@ -241,8 +241,8 @@ where
 
         // Safety:
         // - We don’t care about dropping elements because the caller is responsible for dropping things.
-        // - We forget `buf` so that we don’t drop the other elements.
-        std::mem::forget(buf);
+        // - We forget `buf` so that we don’t drop the other elements, but ensure the buffer itself is cleaned up.
+        buf.forget_elements();
 
         res
     }
@@ -272,8 +272,8 @@ where
 
         // Safety:
         // - We don’t care about dropping elements because the caller is responsible for dropping things.
-        // - We forget `buf` so that we don’t drop the other elements.
-        std::mem::forget(buf);
+        // - We forget `buf` so that we don’t drop the other elements, but ensure the buffer itself is cleaned up.
+        buf.forget_elements();
 
         res
     }
diff --git a/src/base/matrix_view.rs b/src/base/matrix_view.rs
index 47b2965a2..68e5d978e 100644
--- a/src/base/matrix_view.rs
+++ b/src/base/matrix_view.rs
@@ -229,6 +229,11 @@ macro_rules! storage_impl(
                 let it = MatrixIter::new(self).cloned();
                 DefaultAllocator::allocate_from_iterator(nrows, ncols, it)
             }
+
+            #[inline]
+            fn forget_elements(self) {
+                // No cleanup required.
+            }
         }
     )*}
 );
diff --git a/src/base/storage.rs b/src/base/storage.rs
index 93dbfcd8c..708a14402 100644
--- a/src/base/storage.rs
+++ b/src/base/storage.rs
@@ -149,6 +149,9 @@ pub unsafe trait Storage<T: Scalar, R: Dim, C: Dim = U1>: RawStorage<T, R, C> {
     fn clone_owned(&self) -> Owned<T, R, C>
     where
         DefaultAllocator: Allocator<R, C>;
+
+    /// Drops the storage without calling the destructors on the contained elements.
+    fn forget_elements(self);
 }
 
 /// Trait implemented by matrix data storage that can provide a mutable access to its elements.
diff --git a/src/base/vec_storage.rs b/src/base/vec_storage.rs
index 72f844998..f0a0593fa 100644
--- a/src/base/vec_storage.rs
+++ b/src/base/vec_storage.rs
@@ -281,6 +281,17 @@ where
     {
         self.clone()
     }
+
+    #[inline]
+    fn forget_elements(mut self) {
+        // SAFETY: setting the length to zero is always sound, as it does not
+        // cause any memory to be deemed initialized. If the previous length was
+        // non-zero, it is equivalent to using mem::forget to leak each element.
+        // Then, when this function returns, self.data is dropped, freeing the
+        // allocated memory, but the elements are not dropped because they are
+        // now considered uninitialized.
+        unsafe { self.data.set_len(0) };
+    }
 }
 
 unsafe impl<T, R: DimName> RawStorage<T, R, Dyn> for VecStorage<T, R, Dyn> {
@@ -332,6 +343,17 @@ where
     {
         self.clone()
     }
+
+    #[inline]
+    fn forget_elements(mut self) {
+        // SAFETY: setting the length to zero is always sound, as it does not
+        // cause any memory to be deemed initialized. If the previous length was
+        // non-zero, it is equivalent to using mem::forget to leak each element.
+        // Then, when this function returns, self.data is dropped, freeing the
+        // allocated memory, but the elements are not dropped because they are
+        // now considered uninitialized.
+        unsafe { self.data.set_len(0) };
+    }
 }
 
 /*