From a10f9b7524a1e233eceb1d24ecf3a2addd2869e0 Mon Sep 17 00:00:00 2001
From: nameexhaustion <simonlin.rqmmw@slmail.me>
Date: Thu, 1 Aug 2024 20:42:13 +1000
Subject: [PATCH] perf: Shrink `MemSliceInner` enum (#17991)

---
 crates/polars-utils/src/mmap.rs | 36 ++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 19 deletions(-)

diff --git a/crates/polars-utils/src/mmap.rs b/crates/polars-utils/src/mmap.rs
index 98e36bb76d67..5bd8e2df12a5 100644
--- a/crates/polars-utils/src/mmap.rs
+++ b/crates/polars-utils/src/mmap.rs
@@ -11,10 +11,10 @@ mod private {
 
     use crate::mem::prefetch_l2;
 
-    /// A read-only reference to a slice of memory.
+    /// A read-only reference to a slice of memory that can potentially be memory-mapped.
     ///
-    /// This maintains a reference count to the underlying buffer to ensure the memory is kept
-    /// alive. [`MemSlice::slice`] can be used to slice the memory in a zero-copy manner.
+    /// A reference count is kept to the underlying buffer to ensure the memory is kept alive.
+    /// [`MemSlice::slice`] can be used to slice the memory in a zero-copy manner.
     ///
     /// This still owns the all the original memory and therefore should probably not be a long-lasting
     /// structure.
@@ -33,10 +33,8 @@ mod private {
     #[derive(Clone, Debug)]
     #[allow(unused)]
     enum MemSliceInner {
-        Vec(Arc<Vec<u8>>),
         Bytes(bytes::Bytes),
         Mmap(Arc<Mmap>),
-        Unbacked,
     }
 
     impl Deref for MemSlice {
@@ -50,10 +48,7 @@ mod private {
 
     impl Default for MemSlice {
         fn default() -> Self {
-            Self {
-                slice: &[],
-                inner: MemSliceInner::Unbacked,
-            }
+            Self::from_bytes(bytes::Bytes::new())
         }
     }
 
@@ -67,12 +62,7 @@ mod private {
         /// Construct a `MemSlice` from an existing `Vec<u8>`. This is zero-copy.
         #[inline]
         pub fn from_vec(v: Vec<u8>) -> Self {
-            let arc_vec = Arc::new(v);
-
-            Self {
-                slice: unsafe { std::mem::transmute::<&[u8], &'static [u8]>(arc_vec.as_slice()) },
-                inner: MemSliceInner::Vec(arc_vec),
-            }
+            Self::from_bytes(bytes::Bytes::from(v))
         }
 
         /// Construct a `MemSlice` from [`bytes::Bytes`]. This is zero-copy.
@@ -97,10 +87,7 @@ mod private {
         /// Construct a `MemSlice` that simply wraps around a `&[u8]`.
         #[inline]
         pub fn from_slice(slice: &'static [u8]) -> Self {
-            Self {
-                slice,
-                inner: MemSliceInner::Unbacked,
-            }
+            Self::from_bytes(bytes::Bytes::from_static(slice))
         }
 
         /// Attempt to prefetch the memory belonging to to this [`MemSlice`]
@@ -267,6 +254,17 @@ mod tests {
             assert_eq!(ptr_out, ptr);
         }
 
+        {
+            let mut vec = vec![1u8, 2, 3, 4, 5];
+            vec.truncate(2);
+            let ptr = vec.as_ptr();
+
+            let mem_slice = MemSlice::from_vec(vec);
+            let ptr_out = mem_slice.as_ptr();
+
+            assert_eq!(ptr_out, ptr);
+        }
+
         {
             let bytes = bytes::Bytes::from(vec![1u8, 2, 3, 4, 5]);
             let ptr = bytes.as_ptr();