Proposal for a growing owned segment's allocator

[marmeladema]

For similar reasons as explained in https://github.com/capnproto/capnproto-rust/pull/243#discussion_r773027692, I need an owned version of ScratchSpaceHeapAllocator.

Here is my proposal for a growing and owned segments allocator:

use capnp::message::{Allocator, HeapAllocator};
use capnp::{word, Word};

const BYTES_PER_WORD: usize = 8;
const ZERO: Word = word(0, 0, 0, 0, 0, 0, 0, 0);

/// An Allocator whose first segment is backed by a growing vector.
///
/// Recall that an `Allocator` implementation must ensure that allocated segments are
/// initially *zeroed*. `GrowingScratchSpaceHeapAllocator` ensures that is the case by
/// zeroing the entire buffer upon initial construction, and then zeroing any
/// *potentially used* part of the buffer upon `deallocate_segment()`.
///
/// You can reuse a `GrowingScratchSpaceHeapAllocator` by calling `message::Builder::into_allocator()`,
/// or by initally passing it to `message::Builder::new()` as a `&mut GrowingScratchSpaceHeapAllocator`.
/// Such reuse can save significant amounts of zeroing.
pub struct GrowingScratchSpaceHeapAllocator {
    scratch_space: Vec<Word>,
    scratch_space_allocated: bool,
    allocator: HeapAllocator,
    currently_allocated: u32,
    maximum_allocated: u32,
}

impl GrowingScratchSpaceHeapAllocator {
    pub fn with_capacity(capacity: usize) -> Self {
        // We need to ensure that the buffer is zeroed.
        let scratch_space = Word::allocate_zeroed_vec(capacity);

        Self {
            scratch_space,
            scratch_space_allocated: false,
            allocator: HeapAllocator::new(),
            currently_allocated: 0,
            maximum_allocated: 0,
        }
    }
}

impl Default for GrowingScratchSpaceHeapAllocator {
    fn default() -> Self {
        Self {
            scratch_space: Vec::new(),
            scratch_space_allocated: false,
            allocator: HeapAllocator::new(),
            currently_allocated: 0,
            maximum_allocated: 0,
        }
    }
}

unsafe impl<'a> Allocator for GrowingScratchSpaceHeapAllocator {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut u8, u32) {
        let (ptr, len) = if (minimum_size as usize) < self.scratch_space.len()
            && !self.scratch_space_allocated
        {
            self.scratch_space_allocated = true;
            (
                Word::words_to_bytes_mut(&mut self.scratch_space).as_mut_ptr(),
                self.scratch_space.len() as u32,
            )
        } else {
            self.allocator.allocate_segment(minimum_size)
        };

        self.currently_allocated += len;
        if self.currently_allocated > self.maximum_allocated {
            self.maximum_allocated = self.currently_allocated;
        }

        (ptr, len)
    }

    fn deallocate_segment(&mut self, ptr: *mut u8, word_size: u32, words_used: u32) {
        self.currently_allocated -= word_size;
        if ptr == Word::words_to_bytes_mut(&mut self.scratch_space).as_mut_ptr() {
            // Rezero the slice to allow reuse of the allocator. We only need to write
            // words that we know might contain nonzero values.
            unsafe {
                core::ptr::write_bytes(ptr, 0u8, (words_used as usize) * BYTES_PER_WORD);
            }
            self.scratch_space_allocated = false;
            let maximum_allocated = self.maximum_allocated as usize;
            if maximum_allocated > self.scratch_space.len() {
                self.scratch_space.resize(maximum_allocated, ZERO);
            }
        } else {
            self.allocator
                .deallocate_segment(ptr, word_size, words_used);
        }
    }
}

This is mostly a verbatim copy of ScratchSpaceHeapAllocator but which uses an internal Vec<Word> as first segment instead of a user provided mutable slice.

I am sharing this to get some feedback and why not merge it upstream if deemed appropriate.

[dwrensha]

Hm... maybe we should make ScratchSpaceHeapAllocator generic, like this:

pub struct ScratchSpaceHeapAllocator<T> where T: DerefMut[u8] {
    scratch_space: T,
    scratch_space_allocated: bool,
    allocator: HeapAllocator,
}