From 7247e6bb50853d962e4990cb684f2f94198a0128 Mon Sep 17 00:00:00 2001
From: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
Date: Sun, 1 Dec 2024 14:20:01 +0000
Subject: [PATCH] Support shrink to empty (#6817)

* Support shrink to empty

* Docs

* Format
---
 arrow-buffer/src/buffer/immutable.rs |  2 +-
 arrow-buffer/src/buffer/mutable.rs   |  9 ++++++---
 arrow-buffer/src/bytes.rs            | 26 ++++++++++++++++++--------
 3 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs
index be87d726e86..dc0de7736bd 100644
--- a/arrow-buffer/src/buffer/immutable.rs
+++ b/arrow-buffer/src/buffer/immutable.rs
@@ -614,7 +614,7 @@ mod tests {
         let mut shrunk_empty = empty_slice;
         shrunk_empty.shrink_to_fit();
         assert_eq!(shrunk_empty.as_slice(), &[]);
-        assert_eq!(shrunk_empty.capacity(), 1); // NOTE: `Buffer` and `Bytes` doesn't support 0-capacity
+        assert_eq!(shrunk_empty.capacity(), 0);
     }
 
     #[test]
diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs
index 5471bc0e5f6..c4315a1d64c 100644
--- a/arrow-buffer/src/buffer/mutable.rs
+++ b/arrow-buffer/src/buffer/mutable.rs
@@ -467,10 +467,13 @@ impl MutableBuffer {
     }
 }
 
+/// Creates a non-null pointer with alignment of [`ALIGNMENT`]
+///
+/// This is similar to [`NonNull::dangling`]
 #[inline]
-fn dangling_ptr() -> NonNull<u8> {
-    // SAFETY: ALIGNMENT is a non-zero usize which is then casted
-    // to a *mut T. Therefore, `ptr` is not null and the conditions for
+pub(crate) fn dangling_ptr() -> NonNull<u8> {
+    // SAFETY: ALIGNMENT is a non-zero usize which is then cast
+    // to a *mut u8. Therefore, `ptr` is not null and the conditions for
     // calling new_unchecked() are respected.
     #[cfg(miri)]
     {
diff --git a/arrow-buffer/src/bytes.rs b/arrow-buffer/src/bytes.rs
index 924b7a72027..77724137aef 100644
--- a/arrow-buffer/src/bytes.rs
+++ b/arrow-buffer/src/bytes.rs
@@ -24,6 +24,7 @@ use std::ptr::NonNull;
 use std::{fmt::Debug, fmt::Formatter};
 
 use crate::alloc::Deallocation;
+use crate::buffer::dangling_ptr;
 
 /// A continuous, fixed-size, immutable memory region that knows how to de-allocate itself.
 ///
@@ -104,20 +105,29 @@ impl Bytes {
     /// In case of `Err`, the [`Bytes`] will remain as it was (i.e. have the old size).
     pub fn try_realloc(&mut self, new_len: usize) -> Result<(), ()> {
         if let Deallocation::Standard(old_layout) = self.deallocation {
-            let new_len = new_len.max(1); // realloc requires a non-zero size
             if old_layout.size() == new_len {
                 return Ok(()); // Nothing to do
             }
+
             if let Ok(new_layout) = std::alloc::Layout::from_size_align(new_len, old_layout.align())
             {
                 let old_ptr = self.ptr.as_ptr();
-                // SAFETY: the call to `realloc` is safe if all of the following holds (from https://doc.rust-lang.org/stable/std/alloc/trait.GlobalAlloc.html#method.realloc):
-                // * `old_ptr` must be currently allocated via this allocator (guaranteed by the invariant/contract of `Bytes`)
-                // * `old_layout` must be the same layout that was used to allocate that block of memory (same)
-                // * `new_len` must be greater than zero (ensured by the `max` call earlier)
-                // * `new_len`, when rounded up to the nearest multiple of `layout.align()`, must not overflow `isize` (guaranteed by the success of `Layout::from_size_align`)
-                let new_ptr = unsafe { std::alloc::realloc(old_ptr, old_layout, new_len) };
-                if let Some(ptr) = NonNull::new(new_ptr) {
+
+                let new_ptr = match new_layout.size() {
+                    0 => {
+                        // SAFETY: Verified that old_layout.size != new_len (0)
+                        unsafe { std::alloc::dealloc(self.ptr.as_ptr(), old_layout) };
+                        Some(dangling_ptr())
+                    }
+                    // SAFETY: the call to `realloc` is safe if all the following hold (from https://doc.rust-lang.org/stable/std/alloc/trait.GlobalAlloc.html#method.realloc):
+                    // * `old_ptr` must be currently allocated via this allocator (guaranteed by the invariant/contract of `Bytes`)
+                    // * `old_layout` must be the same layout that was used to allocate that block of memory (same)
+                    // * `new_len` must be greater than zero
+                    // * `new_len`, when rounded up to the nearest multiple of `layout.align()`, must not overflow `isize` (guaranteed by the success of `Layout::from_size_align`)
+                    _ => NonNull::new(unsafe { std::alloc::realloc(old_ptr, old_layout, new_len) }),
+                };
+
+                if let Some(ptr) = new_ptr {
                     self.ptr = ptr;
                     self.len = new_len;
                     self.deallocation = Deallocation::Standard(new_layout);