Skip to content

Commit

Permalink
revert unsound generalization of SingleSegmentAllocator
Browse files Browse the repository at this point in the history
  • Loading branch information
dwrensha committed Feb 19, 2024
1 parent 3180a44 commit fbf6619
Showing 1 changed file with 11 additions and 12 deletions.
23 changes: 11 additions & 12 deletions capnp/src/message.rs
Original file line number Diff line number Diff line change
Expand Up @@ -902,21 +902,21 @@ unsafe impl<'a> Allocator for ScratchSpaceHeapAllocator<'a> {
/// You can reuse a `SingleSegmentAllocator` by calling `message::Builder::into_allocator()`,
/// or by initially passing it to `message::Builder::new()` as a `&mut SingleSegmentAllocator`.
/// Such reuse can save significant amounts of zeroing.
pub struct SingleSegmentAllocator<T: AsMut<[u8]>> {
segment: T,
pub struct SingleSegmentAllocator<'a> {
segment: &'a mut [u8],
segment_allocated: bool,
}

impl<T: AsMut<[u8]>> SingleSegmentAllocator<T> {
impl<'a> SingleSegmentAllocator<'a> {
/// Writes zeroes into the entire buffer and constructs a new allocator from it.
///
/// If the buffer is large, this operation could be relatively expensive. If you want to reuse
/// the same scratch space in a later message, you should reuse the entire
/// `SingleSegmentAllocator`, to avoid paying this full cost again.
pub fn new(mut segment: T) -> SingleSegmentAllocator<T> {
pub fn new(segment: &'a mut [u8]) -> SingleSegmentAllocator<'a> {
#[cfg(not(feature = "unaligned"))]
{
if segment.as_mut().as_ptr() as usize % BYTES_PER_WORD != 0 {
if segment.as_ptr() as usize % BYTES_PER_WORD != 0 {
panic!(
"Segment must be 8-byte aligned, or you must enable the \"unaligned\" \
feature in the capnp crate"
Expand All @@ -925,7 +925,7 @@ impl<T: AsMut<[u8]>> SingleSegmentAllocator<T> {
}

// We need to ensure that the buffer is zeroed.
for b in segment.as_mut() {
for b in &mut segment[..] {
*b = 0;
}
SingleSegmentAllocator {
Expand All @@ -935,9 +935,9 @@ impl<T: AsMut<[u8]>> SingleSegmentAllocator<T> {
}
}

unsafe impl<T: AsMut<[u8]>> Allocator for SingleSegmentAllocator<T> {
unsafe impl<'a> Allocator for SingleSegmentAllocator<'a> {
fn allocate_segment(&mut self, minimum_size: u32) -> (*mut u8, u32) {
let available_word_count = self.segment.as_mut().len() / BYTES_PER_WORD;
let available_word_count = self.segment.len() / BYTES_PER_WORD;
if (minimum_size as usize) > available_word_count {
panic!(
"Allocation too large: asked for {minimum_size} words, \
Expand All @@ -947,16 +947,15 @@ unsafe impl<T: AsMut<[u8]>> Allocator for SingleSegmentAllocator<T> {
panic!("Tried to allocated two segments in a SingleSegmentAllocator.")
} else {
self.segment_allocated = true;
let segment_mut = self.segment.as_mut();
(
segment_mut.as_mut_ptr(),
(segment_mut.len() / BYTES_PER_WORD) as u32,
self.segment.as_mut_ptr(),
(self.segment.len() / BYTES_PER_WORD) as u32,
)
}
}

unsafe fn deallocate_segment(&mut self, ptr: *mut u8, _word_size: u32, words_used: u32) {
let seg_ptr = self.segment.as_mut().as_mut_ptr();
let seg_ptr = self.segment.as_mut_ptr();
if ptr == seg_ptr {
// Rezero the slice to allow reuse of the allocator. We only need to write
// words that we know might contain nonzero values.
Expand Down

0 comments on commit fbf6619

Please sign in to comment.