SFT-7482: improve arena allocations

arena/src/lib.rs

@@ -26,20 +26,29 @@
 
 #![no_std]
 
-use core::{cell::RefCell, mem::MaybeUninit};
+use core::{
+    cell::{Cell, UnsafeCell},
+    mem::MaybeUninit,
+};
 
 pub mod boxed;
 
 /// An arena of objects of type `T`.
 pub struct Arena<T, const N: usize> {
-    storage: RefCell<Chunk<T, N>>,
+    storage: UnsafeCell<[MaybeUninit<T>; N]>,
+    len: Cell<usize>,
 }
 
 impl<T, const N: usize> Arena<T, N> {
     /// Construct a new arena.
     pub const fn new() -> Self {
+        // SAFETY: an array of `MaybeUninit<T>` may contain uninitialized
+        // elements regardless of `T`.
+        let storage = unsafe { MaybeUninit::<[MaybeUninit<T>; N]>::uninit().assume_init() };
+
         Self {
-            storage: RefCell::new(Chunk::new()),
+            storage: UnsafeCell::new(storage),
+            len: Cell::new(0),
         }
     }
 
@@ -48,48 +57,117 @@ impl<T, const N: usize> Arena<T, N> {
     ///
     /// If there's not enough space left in the arena, then the item is
     /// returned as-is.
-    #[allow(clippy::mut_from_ref)] // Sound: interior mutability via RefCell
+    ///
+    /// Values allocated directly are not dropped when the arena is dropped.
+    /// Use [`boxed::Box`] when the value needs drop glue.
+    ///
+    /// # Safety invariants
+    ///
+    /// - `len` increases monotonically, so no two successful allocations use
+    ///   the same slot.
+    /// - A mutable reference is created only for the newly initialized slot.
+    ///   Later allocations access other slots only through raw pointers and
+    ///   never create a mutable reference to the complete backing array.
+    /// - The returned reference is tied to the arena, so safe code cannot
+    ///   outlive the arena or reset a slot while that reference exists.
+    #[allow(clippy::mut_from_ref)] // SAFETY: see the invariants below.
     pub fn alloc(&self, item: T) -> Result<&mut T, T> {
-        let mut storage = self.storage.borrow_mut();
-        let len = storage.len();
-        storage.push(item)?;
-        Ok(unsafe { &mut *storage.as_mut_ptr().add(len) })
+        let slot = self.len.get();
+        if slot >= N {
+            return Err(item);
+        }
+
+        let ptr = self.slot_ptr(slot);
+        // SAFETY: `slot < N`, the slot is uninitialized, and the monotonic
+        // allocation index ensures no other reference can point at it.
+        unsafe { ptr.write(item) };
+
+        self.len.set(slot + 1);
+
+        // SAFETY: the slot was initialized above and is uniquely owned by
+        // this allocation for the lifetime of the arena.
+        Ok(unsafe { &mut *ptr })
     }
-}
 
-struct Chunk<T, const N: usize> {
-    buffer: [MaybeUninit<T>; N],
-    len: usize,
+    fn slot_ptr(&self, slot: usize) -> *mut T {
+        // This should compile out if invariants hold up, and if not,
+        // crashing is preferable to undefined behavior.
+        assert!(
+            slot < N,
+            "Arena allocation slot out of bounds: slot = {}, N = {}",
+            slot,
+            N
+        );
+
+        // SAFETY: callers ensure `slot < N`. Casting the raw pointer avoids
+        // creating a reference to the whole array.
+        unsafe {
+            self.storage
+                .get()
+                .cast::<MaybeUninit<T>>()
+                .add(slot)
+                .cast::<T>()
+        }
+    }
 }
 
-impl<T, const N: usize> Chunk<T, N> {
-    const ELEM: MaybeUninit<T> = MaybeUninit::uninit();
-    const INIT: [MaybeUninit<T>; N] = [Self::ELEM; N];
+#[cfg(test)]
+mod tests {
+    use core::cell::Cell;
 
-    pub const fn new() -> Self {
-        Self {
-            buffer: Self::INIT,
-            len: 0,
-        }
+    use super::{boxed::Box, Arena};
+
+    #[test]
+    fn allocates_distinct_slots_up_to_capacity() {
+        let arena: Arena<u32, 2> = Arena::new();
+
+        let first = arena.alloc(11).unwrap();
+        let second = arena.alloc(22).unwrap();
+
+        assert_eq!((*first, *second), (11, 22));
+        assert_eq!(arena.alloc(33), Err(33));
+    }
+
+    #[test]
+    fn earlier_allocation_remains_usable_after_later_allocation() {
+        let arena: Arena<u32, 2> = Arena::new();
+
+        let first = arena.alloc(11).unwrap();
+        let second = arena.alloc(22).unwrap();
+
+        *first += 1;
+        *second += 1;
+
+        assert_eq!((*first, *second), (12, 23));
     }
 
-    pub const fn len(&self) -> usize {
-        self.len
+    #[test]
+    fn counts_zero_sized_allocations_toward_capacity() {
+        let arena: Arena<(), 2> = Arena::new();
+
+        let first = arena.alloc(()).unwrap();
+        let second = arena.alloc(()).unwrap();
+
+        assert_eq!((*first, *second), ((), ()));
+        assert_eq!(arena.alloc(()), Err(()));
     }
 
-    pub fn push(&mut self, item: T) -> Result<(), T> {
-        if self.len < N {
-            unsafe {
-                *self.buffer.get_unchecked_mut(self.len) = MaybeUninit::new(item);
-                self.len += 1;
+    #[test]
+    fn boxed_value_is_dropped_once() {
+        #[derive(Debug)]
+        struct DropCounter<'a>(&'a Cell<u8>);
+
+        impl Drop for DropCounter<'_> {
+            fn drop(&mut self) {
+                self.0.set(self.0.get() + 1);
             }
-            Ok(())
-        } else {
-            Err(item)
         }
-    }
 
-    pub fn as_mut_ptr(&mut self) -> *mut T {
-        self.buffer.as_mut_ptr() as *mut T
+        let drops = Cell::new(0);
+        let arena: Arena<DropCounter, 1> = Arena::new();
+        let value = Box::new_in(DropCounter(&drops), &arena).unwrap();
+
+        drop(value);
+        assert_eq!(drops.get(), 1);
     }
 }

urtypes/src/value.rs

@@ -239,6 +239,31 @@ mod tests {
         assert!(matches!(decoded, Terminal::WitnessScriptHash(_)));
     }
 
+    #[test]
+    fn test_decode_output_descriptor_retains_nested_arena_nodes() {
+        const CBOR: &[u8] = &[
+            0xd9, 0x01, 0x90, // script-hash
+            0xd9, 0x01, 0x90, // script-hash
+            0xd9, 0x01, 0x91, // witness-script-hash
+            0xd9, 0x01, 0x98, // raw-script
+            0x41, 0x42, // byte string: 0x42
+        ];
+
+        let arena: TerminalContext<3> = TerminalContext::new();
+        let decoded = decode_output_descriptor("crypto-output", CBOR, &arena).unwrap();
+
+        let Terminal::ScriptHash(first) = decoded else {
+            panic!("expected outer script-hash");
+        };
+        let Terminal::ScriptHash(second) = &*first else {
+            panic!("expected nested script-hash");
+        };
+        let Terminal::WitnessScriptHash(third) = &**second else {
+            panic!("expected witness-script-hash");
+        };
+        assert!(matches!(&**third, Terminal::RawScript(&[0x42])));
+    }
+
     #[test]
     fn test_decode_output_descriptor_rejects_output_descriptor_alias() {
         // BCR-2023-010 `output-descriptor` uses a different CBOR shape (tag