allow BuilderArena to be used in no-alloc mode (#442)

capnp/src/lib.rs

@@ -623,28 +623,31 @@ impl ::std::error::Error for Error {
 
 /// Helper struct that allows `MessageBuilder::get_segments_for_output()` to avoid heap allocations
 /// in the single-segment case.
-#[cfg(feature = "alloc")]
 pub enum OutputSegments<'a> {
     SingleSegment([&'a [u8]; 1]),
+
+    #[cfg(feature = "alloc")]
     MultiSegment(Vec<&'a [u8]>),
 }
 
-#[cfg(feature = "alloc")]
 impl<'a> core::ops::Deref for OutputSegments<'a> {
     type Target = [&'a [u8]];
     fn deref(&self) -> &[&'a [u8]] {
         match self {
             OutputSegments::SingleSegment(s) => s,
+
+            #[cfg(feature = "alloc")]
             OutputSegments::MultiSegment(v) => v,
         }
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<'s> message::ReaderSegments for OutputSegments<'s> {
     fn get_segment(&self, id: u32) -> Option<&[u8]> {
         match self {
             OutputSegments::SingleSegment(s) => s.get(id as usize).copied(),
+
+            #[cfg(feature = "alloc")]
             OutputSegments::MultiSegment(v) => v.get(id as usize).copied(),
         }
     }

capnp/src/message.rs

@@ -73,15 +73,12 @@ use alloc::vec::Vec;
 use core::convert::From;
 
 use crate::any_pointer;
-#[cfg(feature = "alloc")]
 use crate::private::arena::{BuilderArena, BuilderArenaImpl};
 use crate::private::arena::{ReaderArena, ReaderArenaImpl};
 use crate::private::layout;
 use crate::private::units::BYTES_PER_WORD;
-#[cfg(feature = "alloc")]
 use crate::traits::{FromPointerBuilder, SetPointerBuilder};
 use crate::traits::{FromPointerReader, Owned};
-#[cfg(feature = "alloc")]
 use crate::OutputSegments;
 use crate::Result;
 
@@ -344,7 +341,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<A, T> From<Builder<A>> for TypedReader<Builder<A>, T>
 where
     A: Allocator,
@@ -356,7 +352,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<A, T> From<TypedBuilder<T, A>> for TypedReader<Builder<A>, T>
 where
     A: Allocator,
@@ -404,18 +399,15 @@ pub unsafe trait Allocator {
 }
 
 /// A container used to build a message.
-#[cfg(feature = "alloc")]
 pub struct Builder<A>
 where
     A: Allocator,
 {
     arena: BuilderArenaImpl<A>,
 }
 
-#[cfg(feature = "alloc")]
 unsafe impl<A> Send for Builder<A> where A: Send + Allocator {}
 
-#[cfg(feature = "alloc")]
 fn _assert_kinds() {
     fn _assert_send<T: Send>() {}
     fn _assert_reader<S: ReaderSegments + Send>() {
@@ -426,7 +418,6 @@ fn _assert_kinds() {
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<A> Builder<A>
 where
     A: Allocator,
@@ -533,7 +524,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<A> ReaderSegments for Builder<A>
 where
     A: Allocator,
@@ -563,6 +553,19 @@ where
     message: Builder<A>,
 }
 
+// Defined separately because the A=HeapAllocator default type
+// argument is not allowed in no-alloc mode.
+// TODO(apibump): remove the A=HeapAllocator thing above?
+#[cfg(not(feature = "alloc"))]
+pub struct TypedBuilder<T, A>
+where
+    T: Owned,
+    A: Allocator,
+{
+    marker: ::core::marker::PhantomData<T>,
+    message: Builder<A>,
+}
+
 #[cfg(feature = "alloc")]
 impl<T> TypedBuilder<T, HeapAllocator>
 where
@@ -573,7 +576,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<T, A> TypedBuilder<T, A>
 where
     T: Owned,
@@ -623,7 +625,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<T, A> From<Builder<A>> for TypedBuilder<T, A>
 where
     T: Owned,
@@ -951,7 +952,6 @@ unsafe impl<'a> Allocator for SingleSegmentAllocator<'a> {
     }
 }
 
-#[cfg(feature = "alloc")]
 unsafe impl<'a, A> Allocator for &'a mut A
 where
     A: Allocator,

capnp/src/private/arena.rs

@@ -20,17 +20,14 @@
 
 #[cfg(feature = "alloc")]
 use alloc::vec::Vec;
-#[cfg(feature = "alloc")]
 use core::slice;
 use core::u64;
 
 use crate::message;
-#[cfg(feature = "alloc")]
 use crate::message::Allocator;
 use crate::message::ReaderSegments;
 use crate::private::read_limiter::ReadLimiter;
 use crate::private::units::*;
-#[cfg(feature = "alloc")]
 use crate::OutputSegments;
 use crate::{Error, ErrorKind, Result};
 
@@ -163,7 +160,6 @@ pub trait BuilderArena: ReaderArena {
 }
 
 /// A wrapper around a memory segment used in building a message.
-#[cfg(feature = "alloc")]
 struct BuilderSegment {
     /// Pointer to the start of the segment.
     ptr: *mut u8,
@@ -177,25 +173,71 @@ struct BuilderSegment {
 }
 
 #[cfg(feature = "alloc")]
+type BuilderSegmentArray = Vec<BuilderSegment>;
+
+#[cfg(not(feature = "alloc"))]
+#[derive(Default)]
+struct BuilderSegmentArray {
+    // In the no-alloc case, we only allow a single segment.
+    segment: Option<BuilderSegment>,
+}
+
+#[cfg(not(feature = "alloc"))]
+impl BuilderSegmentArray {
+    fn len(&self) -> usize {
+        match self.segment {
+            Some(_) => 1,
+            None => 0,
+        }
+    }
+
+    fn push(&mut self, segment: BuilderSegment) {
+        if self.segment.is_some() {
+            panic!("multiple segments are not supported in no-alloc mode")
+        }
+        self.segment = Some(segment);
+    }
+}
+
+#[cfg(not(feature = "alloc"))]
+impl core::ops::Index<usize> for BuilderSegmentArray {
+    type Output = BuilderSegment;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        assert_eq!(index, 0);
+        match &self.segment {
+            Some(s) => s,
+            None => panic!("no segment"),
+        }
+    }
+}
+
+#[cfg(not(feature = "alloc"))]
+impl core::ops::IndexMut<usize> for BuilderSegmentArray {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        assert_eq!(index, 0);
+        match &mut self.segment {
+            Some(s) => s,
+            None => panic!("no segment"),
+        }
+    }
+}
+
 pub struct BuilderArenaImplInner<A>
 where
     A: Allocator,
 {
     allocator: Option<A>, // None if has already be deallocated.
-
-    // TODO(perf): Try using smallvec to avoid heap allocations in the single-segment case?
-    segments: Vec<BuilderSegment>,
+    segments: BuilderSegmentArray,
 }
 
-#[cfg(feature = "alloc")]
 pub struct BuilderArenaImpl<A>
 where
     A: Allocator,
 {
     inner: BuilderArenaImplInner<A>,
 }
 
-#[cfg(feature = "alloc")]
 impl<A> BuilderArenaImpl<A>
 where
     A: Allocator,
@@ -204,7 +246,7 @@ where
         Self {
             inner: BuilderArenaImplInner {
                 allocator: Some(allocator),
-                segments: Vec::new(),
+                segments: Default::default(),
             },
         }
     }
@@ -227,18 +269,25 @@ where
             };
             OutputSegments::SingleSegment([slice])
         } else {
-            let mut v = Vec::with_capacity(reff.segments.len());
-            for seg in &reff.segments {
-                // See safety argument in above branch.
-                let slice = unsafe {
-                    slice::from_raw_parts(
-                        seg.ptr as *const _,
-                        seg.allocated as usize * BYTES_PER_WORD,
-                    )
-                };
-                v.push(slice);
+            #[cfg(feature = "alloc")]
+            {
+                let mut v = Vec::with_capacity(reff.segments.len());
+                for seg in &reff.segments {
+                    // See safety argument in above branch.
+                    let slice = unsafe {
+                        slice::from_raw_parts(
+                            seg.ptr as *const _,
+                            seg.allocated as usize * BYTES_PER_WORD,
+                        )
+                    };
+                    v.push(slice);
+                }
+                OutputSegments::MultiSegment(v)
+            }
+            #[cfg(not(feature = "alloc"))]
+            {
+                panic!("invalid number of segments: {}", reff.segments.len());
             }
-            OutputSegments::MultiSegment(v)
         }
     }
 
@@ -258,7 +307,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<A> ReaderArena for BuilderArenaImpl<A>
 where
     A: Allocator,
@@ -290,7 +338,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<A> BuilderArenaImplInner<A>
 where
     A: Allocator,
@@ -341,11 +388,19 @@ where
 
     fn deallocate_all(&mut self) {
         if let Some(a) = &mut self.allocator {
+            #[cfg(feature = "alloc")]
             for seg in &self.segments {
                 unsafe {
                     a.deallocate_segment(seg.ptr, seg.capacity, seg.allocated);
                 }
             }
+
+            #[cfg(not(feature = "alloc"))]
+            if let Some(seg) = &self.segments.segment {
+                unsafe {
+                    a.deallocate_segment(seg.ptr, seg.capacity, seg.allocated);
+                }
+            }
         }
     }
 
@@ -355,7 +410,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<A> BuilderArena for BuilderArenaImpl<A>
 where
     A: Allocator,
@@ -377,7 +431,6 @@ where
     }
 }
 
-#[cfg(feature = "alloc")]
 impl<A> Drop for BuilderArenaImplInner<A>
 where
     A: Allocator,

capnp/tests/single_segment_allocator.rs

@@ -1,10 +1,8 @@
-#![cfg(feature = "alloc")]
-
 use capnp::message;
 
 #[test]
 pub fn single_segment_allocator() {
-    let mut buffer = capnp::Word::allocate_zeroed_vec(200);
+    let mut buffer = [capnp::word(0, 0, 0, 0, 0, 0, 0, 0); 200];
     let allocator =
         message::SingleSegmentAllocator::new(capnp::Word::words_to_bytes_mut(&mut buffer[..]));
     let mut msg = message::Builder::new(allocator);