Skip to content

Commit

Permalink
Support shrink to empty
Browse files Browse the repository at this point in the history
  • Loading branch information
tustvold committed Nov 29, 2024
1 parent fdb06e8 commit 0ee84f4
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 10 deletions.
2 changes: 1 addition & 1 deletion arrow-buffer/src/buffer/immutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
2 changes: 1 addition & 1 deletion arrow-buffer/src/buffer/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -468,7 +468,7 @@ impl MutableBuffer {
}

#[inline]
fn dangling_ptr() -> NonNull<u8> {
pub(crate) 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
// calling new_unchecked() are respected.
Expand Down
25 changes: 17 additions & 8 deletions arrow-buffer/src/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down Expand Up @@ -104,20 +105,28 @@ 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 => {
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);
Expand Down

0 comments on commit 0ee84f4

Please sign in to comment.