Add Array::shrink_to_fit(&mut self)

arrow-buffer/src/buffer/immutable.rs

@@ -177,12 +177,27 @@ impl Buffer {
     ///
     /// The memory region will be reallocated using `std::alloc::realloc`.
     pub fn shrink_to_fit(&mut self) {
-        let desired_capacity = self.len();
+        let offset = self.ptr_offset();
+        let is_empty = self.is_empty();
+        let desired_capacity = if is_empty {
+            0
+        } else {
+            // For realloc to work, we cannot free the elements before the offset
+            offset + self.len()
+        };
         if desired_capacity < self.capacity() {
             if let Some(bytes) = Arc::get_mut(&mut self.data) {
-                // We can sefely ignore errors.
-                // We are best-effort, and if the realloc fails, we still have the old size.
-                bytes.try_realloc(desired_capacity).ok();
+                if bytes.try_realloc(desired_capacity).is_ok() {
+                    // Realloc complete - update our pointer into `bytes`:
+                    self.ptr = if is_empty {
+                        bytes.as_ptr()
+                    } else {
+                        // SAFETY: we kept all elements leading up to the offset
+                        unsafe { bytes.as_ptr().add(offset) }
+                    }
+                } else {
+                    // Failure to reallocate is fine; we just failed to free up memory.
+                }
             }
         }
     }
@@ -576,30 +591,30 @@ mod tests {
 
     #[test]
     fn test_shrink_to_fit() {
-        let original = Buffer::from(&[1, 2, 3, 4, 5, 6, 7]);
-        assert_eq!(original.len(), 7);
+        let original = Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7]);
+        assert_eq!(original.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7]);
         assert_eq!(original.capacity(), 64);
 
-        let slice = original.slice(3);
+        let slice = original.slice_with_length(2, 3);
         drop(original); // Make sure the buffer isn't shared (or shrink_to_fit won't work)
-        assert_eq!(slice.len(), 4);
+        assert_eq!(slice.as_slice(), &[2, 3, 4]);
         assert_eq!(slice.capacity(), 64);
 
         let mut shrunk = slice;
         shrunk.shrink_to_fit();
-        assert_eq!(shrunk.len(), 4);
-        assert_eq!(shrunk.capacity(), 4);
+        assert_eq!(shrunk.as_slice(), &[2, 3, 4]);
+        assert_eq!(shrunk.capacity(), 5); // shrink_to_fit is allowed to keep the elements before the offset
 
         // Test that we can handle empty slices:
-        let empty_slice = shrunk.slice(4);
+        let empty_slice = shrunk.slice_with_length(1, 0);
         drop(shrunk); // Make sure the buffer isn't shared (or shrink_to_fit won't work)
-        assert_eq!(empty_slice.len(), 0);
-        assert_eq!(empty_slice.capacity(), 4);
+        assert_eq!(empty_slice.as_slice(), &[]);
+        assert_eq!(empty_slice.capacity(), 5);
 
         let mut shrunk_empty = empty_slice;
         shrunk_empty.shrink_to_fit();
-        assert_eq!(shrunk_empty.len(), 0);
-        assert_eq!(shrunk_empty.capacity(), 1); // `Buffer` and `Bytes` doesn't support 0-capacity, so we shrink to 1
+        assert_eq!(shrunk_empty.as_slice(), &[]);
+        assert_eq!(shrunk_empty.capacity(), 1); // NOTE: `Buffer` and `Bytes` doesn't support 0-capacity
     }
 
     #[test]