Skip to content

Commit 5bf391f

Browse files
committed
feat: Add optimized grow()/shrink() functions for Pool
If resizing is requested for the last allocation in the pool, it may be possible to adjust pool data and avoid any real allocations.
1 parent b796695 commit 5bf391f

File tree

1 file changed

+84
-0
lines changed

1 file changed

+84
-0
lines changed

src/core/pool.rs

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,54 @@ unsafe impl Allocator for Pool {
6767
ngx_pfree(self.0.as_ptr(), ptr.as_ptr().cast());
6868
}
6969
}
70+
71+
unsafe fn grow(
72+
&self,
73+
ptr: NonNull<u8>,
74+
old_layout: Layout,
75+
new_layout: Layout,
76+
) -> Result<NonNull<[u8]>, AllocError> {
77+
debug_assert!(
78+
new_layout.size() >= old_layout.size(),
79+
"`new_layout.size()` must be greater than or equal to `old_layout.size()`"
80+
);
81+
self.resize_in_place(ptr, old_layout, new_layout)
82+
}
83+
84+
unsafe fn grow_zeroed(
85+
&self,
86+
ptr: NonNull<u8>,
87+
old_layout: Layout,
88+
new_layout: Layout,
89+
) -> Result<NonNull<[u8]>, AllocError> {
90+
debug_assert!(
91+
new_layout.size() >= old_layout.size(),
92+
"`new_layout.size()` must be greater than or equal to `old_layout.size()`"
93+
);
94+
self.resize_in_place(ptr, old_layout, new_layout)
95+
.inspect(|new_ptr| {
96+
unsafe {
97+
ptr::write_bytes(
98+
new_ptr.as_ptr().cast::<u8>().byte_add(new_layout.size()),
99+
0,
100+
old_layout.size() - new_layout.size(),
101+
)
102+
};
103+
})
104+
}
105+
106+
unsafe fn shrink(
107+
&self,
108+
ptr: NonNull<u8>,
109+
old_layout: Layout,
110+
new_layout: Layout,
111+
) -> Result<NonNull<[u8]>, AllocError> {
112+
debug_assert!(
113+
new_layout.size() <= old_layout.size(),
114+
"`new_layout.size()` must be smaller than or equal to `old_layout.size()`"
115+
);
116+
self.resize_in_place(ptr, old_layout, new_layout)
117+
}
70118
}
71119

72120
impl AsRef<ngx_pool_t> for Pool {
@@ -229,6 +277,42 @@ impl Pool {
229277
p
230278
}
231279
}
280+
281+
/// Resizes a memory allocation in place if possible.
282+
///
283+
/// If resizing is requested for the last allocation in the pool, it may be
284+
/// possible to adjust pool data and avoid any real allocations.
285+
///
286+
/// # Safety
287+
/// This function is marked as unsafe because it involves raw pointer manipulation.
288+
unsafe fn resize_in_place(
289+
&self,
290+
ptr: NonNull<u8>,
291+
old_layout: Layout,
292+
new_layout: Layout,
293+
) -> Result<NonNull<[u8]>, AllocError> {
294+
if ptr.byte_add(old_layout.size()).as_ptr() == self.as_ref().d.last
295+
&& ptr.byte_add(new_layout.size()).as_ptr() <= self.as_ref().d.end
296+
&& ptr.align_offset(new_layout.align()) == 0
297+
{
298+
let pool = self.0.as_ptr();
299+
unsafe {
300+
(*pool).d.last = (*pool)
301+
.d
302+
.last
303+
.byte_offset(new_layout.size() as isize - old_layout.size() as isize)
304+
};
305+
Ok(NonNull::slice_from_raw_parts(ptr, new_layout.size()))
306+
} else {
307+
let size = core::cmp::min(old_layout.size(), new_layout.size());
308+
let new_ptr = self.allocate(new_layout)?;
309+
unsafe {
310+
ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_ptr().cast(), size);
311+
self.deallocate(ptr, old_layout);
312+
}
313+
Ok(new_ptr)
314+
}
315+
}
232316
}
233317

234318
/// Cleanup handler for a specific type `T`.

0 commit comments

Comments
 (0)