add SingleSegmentAllocator (#441)

Add SingleSegmentAllocator, to support message building in contexts where heap allocation is not available.

capnp/src/message.rs

@@ -849,11 +849,16 @@ unsafe impl<'a> Allocator for ScratchSpaceHeapAllocator<'a> {
     }
 
     unsafe fn deallocate_segment(&mut self, ptr: *mut u8, word_size: u32, words_used: u32) {
-        if ptr == self.scratch_space.as_mut_ptr() {
+        let seg_ptr = self.scratch_space.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.
             unsafe {
-                core::ptr::write_bytes(ptr, 0u8, (words_used as usize) * BYTES_PER_WORD);
+                core::ptr::write_bytes(
+                    seg_ptr, // miri isn't happy if we use ptr instead
+                    0u8,
+                    (words_used as usize) * BYTES_PER_WORD,
+                );
             }
             self.scratch_space_allocated = false;
         } else {
@@ -863,6 +868,89 @@ unsafe impl<'a> Allocator for ScratchSpaceHeapAllocator<'a> {
     }
 }
 
+/// An Allocator whose first and only segment is a backed by a user-provided buffer.
+/// If the segment fill up, subsequent allocations trigger panics.
+///
+/// The main purpose of this struct is to be used in situations where heap allocation
+/// is not available.
+///
+/// Recall that an `Allocator` implementation must ensure that allocated segments are
+/// initially *zeroed*. `SingleSegmentAllocator` 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 `SingleSegmentAllocator` by calling `message::Builder::into_allocator()`,
+/// or by initally passing it to `message::Builder::new()` as a `&mut SingleSegmentAllocator`.
+/// Such reuse can save significant amounts of zeroing.
+pub struct SingleSegmentAllocator<'a> {
+    segment: &'a mut [u8],
+    segment_allocated: bool,
+}
+
+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(segment: &'a mut [u8]) -> SingleSegmentAllocator<'a> {
+        #[cfg(not(feature = "unaligned"))]
+        {
+            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"
+                );
+            }
+        }
+
+        // We need to ensure that the buffer is zeroed.
+        for b in &mut segment[..] {
+            *b = 0;
+        }
+        SingleSegmentAllocator {
+            segment,
+            segment_allocated: false,
+        }
+    }
+}
+
+unsafe impl<'a> Allocator for SingleSegmentAllocator<'a> {
+    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut u8, u32) {
+        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, \
+                    but only {available_word_count} are available."
+            )
+        } else if self.segment_allocated {
+            panic!("Tried to allocated two segments in a SingleSegmentAllocator.")
+        } else {
+            self.segment_allocated = true;
+            (
+                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_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.
+            unsafe {
+                core::ptr::write_bytes(
+                    seg_ptr, // miri isn't happy if we use ptr instead
+                    0u8,
+                    (words_used as usize) * BYTES_PER_WORD,
+                );
+            }
+            self.segment_allocated = false;
+        }
+    }
+}
+
 #[cfg(feature = "alloc")]
 unsafe impl<'a, A> Allocator for &'a mut A
 where

capnp/tests/single_segment_allocator.rs

@@ -0,0 +1,15 @@
+#![cfg(feature = "alloc")]
+
+use capnp::message;
+
+#[test]
+pub fn single_segment_allocator() {
+    let mut buffer = capnp::Word::allocate_zeroed_vec(200);
+    let allocator =
+        message::SingleSegmentAllocator::new(capnp::Word::words_to_bytes_mut(&mut buffer[..]));
+    let mut msg = message::Builder::new(allocator);
+    msg.set_root("hello world!").unwrap();
+
+    let s: capnp::text::Reader = msg.get_root_as_reader().unwrap();
+    assert_eq!("hello world!", s);
+}