Rust: Provide a trait method to (optionally) control resource allocation

crates/rust/src/interface.rs

@@ -239,6 +239,30 @@ fn _resource_rep(handle: u32) -> *mut u8
 
                     "#
             );
+            let box_path = self.path_to_box();
+            uwriteln!(
+                self.src,
+                r#"#[doc(hidden)]
+/// Place the value on the heap or in an arena, return the raw pointer.
+/// Override for custom resource allocators.
+fn _resource_into_raw(val: Option<Self>) -> *mut Option<Self> where Self: Sized
+{{
+    {box_path}::into_raw({box_path}::new(val))
+}}
+
+#[doc(hidden)]
+/// Consumes the value from the handle, handle is invalid afterwards.
+/// 
+/// # Safety
+/// 
+/// See Box::from_raw
+unsafe fn _resource_from_raw(handle: *mut Option<Self>) -> Option<Self> where Self: Sized
+{{
+    *unsafe {{ {box_path}::from_raw(handle) }}
+}}
+
+                    "#
+            );
             for method in methods {
                 self.src.push_str(method);
             }
@@ -2716,7 +2740,6 @@ impl<'a> wit_bindgen_core::InterfaceGenerator<'a> for InterfaceGenerator<'a> {
                 Identifier::World(_) => unimplemented!("resource exports from worlds"),
                 Identifier::StreamOrFuturePayload => unreachable!(),
             };
-            let box_path = self.path_to_box();
             uwriteln!(
                 self.src,
                 r#"
@@ -2737,8 +2760,7 @@ impl {camel} {{
     pub fn new<T: Guest{camel}>(val: T) -> Self {{
         Self::type_guard::<T>();
         let val: _{camel}Rep<T> = Some(val);
-        let ptr: *mut _{camel}Rep<T> =
-            {box_path}::into_raw({box_path}::new(val));
+        let ptr: *mut _{camel}Rep<T> = T::_resource_into_raw(val);
         unsafe {{
             Self::from_handle(T::_resource_new(ptr.cast()))
         }}
@@ -2797,9 +2819,9 @@ impl {camel} {{
     }}
 
     #[doc(hidden)]
-    pub unsafe fn dtor<T: 'static>(handle: *mut u8) {{
+    pub unsafe fn dtor<T: 'static + Guest{camel}>(handle: *mut u8) {{
         Self::type_guard::<T>();
-        let _ = unsafe {{ {box_path}::from_raw(handle as *mut _{camel}Rep<T>) }};
+        let _ = unsafe {{ T::_resource_from_raw(handle as *mut _{camel}Rep<T>) }};
     }}
 
     fn as_ptr<T: Guest{camel}>(&self) -> *mut _{camel}Rep<T> {{

tests/runtime/rust/arena-allocated-resources/runner.rs

@@ -0,0 +1,16 @@
+include!(env!("BINDINGS"));
+
+use crate::test::arena_allocated_resources::to_test::Thing;
+
+struct Component;
+
+export!(Component);
+
+impl Guest for Component {
+    fn run() {
+        let thing1 = Thing::new(3);
+        let thing2 = Thing::new(5);
+        assert_eq!(3, thing1.get());
+        assert_eq!(5, thing2.get());
+    }
+}

tests/runtime/rust/arena-allocated-resources/test.rs

@@ -0,0 +1,127 @@
+include!(env!("BINDINGS"));
+
+use crate::exports::test::arena_allocated_resources::to_test::{Guest, GuestThing};
+
+export!(Component);
+
+struct Component;
+
+impl Guest for Component {
+    type Thing = MyThing;
+}
+
+mod arena {
+
+    use core::{cell::UnsafeCell, mem::MaybeUninit};
+
+    /// A simple no_std arena allocator for fixed-size allocations.
+    ///
+    /// The arena allocates items of type T sequentially from a pre-allocated buffer
+    /// and does not support individual deallocation. Memory is reclaimed
+    /// only when the entire arena is reset.
+    pub struct Arena<T, const SIZE: usize> {
+        buffer: [MaybeUninit<T>; SIZE],
+        offset: usize,
+    }
+
+    impl<T, const SIZE: usize> Arena<T, SIZE> {
+        /// Allocates space for a single item of type T.
+        /// Returns a mutable reference to the allocated memory, or None if there's insufficient space.
+        pub fn alloc_one(&mut self) -> Option<&mut T> {
+            if self.offset < SIZE {
+                let ptr = self.buffer[self.offset].as_mut_ptr();
+                self.offset += 1;
+                Some(unsafe { &mut *ptr })
+            } else {
+                None
+            }
+        }
+    }
+
+    /// A static-safe wrapper for Arena that uses interior mutability.
+    ///
+    /// This allows an Arena to be stored in a static variable and accessed safely
+    /// in single-threaded contexts without requiring std or alloc.
+    ///
+    /// # Safety
+    ///
+    /// This type is safe to use in single-threaded environments. In multi-threaded
+    /// contexts, external synchronization is required.
+    pub struct StaticArena<T, const SIZE: usize> {
+        arena: UnsafeCell<Arena<T, SIZE>>,
+    }
+
+    // SAFETY: StaticArena is Sync because we enforce single-threaded access through
+    // the API. It can be safely shared across threads as long as only one thread
+    // accesses it at a time (which is the responsibility of the user).
+    unsafe impl<T, const SIZE: usize> Sync for StaticArena<T, SIZE> where T: Sync {}
+
+    // SAFETY: StaticArena is Send because the underlying Arena can be moved between
+    // threads, and T itself must be Send.
+    unsafe impl<T, const SIZE: usize> Send for StaticArena<T, SIZE> where T: Send {}
+
+    impl<T, const SIZE: usize> StaticArena<T, SIZE> {
+        /// Creates a new static arena.
+        pub const fn new() -> Self {
+            StaticArena {
+                arena: UnsafeCell::new(Arena {
+                    buffer: [const { MaybeUninit::uninit() }; SIZE],
+                    offset: 0,
+                }),
+            }
+        }
+
+        /// Gets mutable access to the arena.
+        ///
+        /// # Safety
+        ///
+        /// This is safe in single-threaded contexts. In multi-threaded contexts,
+        /// the caller must ensure exclusive access.
+        #[inline]
+        pub fn get_mut(&self) -> &mut Arena<T, SIZE> {
+            unsafe { &mut *self.arena.get() }
+        }
+
+        /// Allocates a single item.
+        pub fn alloc_one(&self) -> Option<&mut T> {
+            self.get_mut().alloc_one()
+        }
+    }
+}
+
+use arena::StaticArena;
+
+#[derive(Clone)]
+struct MyThing {
+    contents: u32,
+}
+
+static ARENA: StaticArena<Option<MyThing>, 4> = StaticArena::new();
+
+impl GuestThing for MyThing {
+    fn new(v: u32) -> MyThing {
+        MyThing { contents: v }
+    }
+    fn get(&self) -> u32 {
+        self.contents
+    }
+    fn _resource_into_raw(val: Option<Self>) -> *mut Option<Self>
+    where
+        Self: Sized,
+    {
+        val.and_then(|v| {
+            ARENA.alloc_one().map(|x| {
+                *x = Some(v);
+                x as *mut _
+            })
+        })
+        .unwrap_or(core::ptr::null_mut())
+    }
+    unsafe fn _resource_from_raw(handle: *mut Option<Self>) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        let res = unsafe { &mut *handle }.take();
+        res
+    }
+}

tests/runtime/rust/arena-allocated-resources/test.wit

@@ -0,0 +1,18 @@
+package test:arena-allocated-resources;
+
+interface to-test {
+  resource thing {
+    constructor(v: u32);
+    get: func() -> u32;
+  }
+}
+
+world test {
+  export to-test;
+}
+
+world runner {
+  import to-test;
+
+  export run: func();
+}