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

crates/rust/src/interface.rs

@@ -239,6 +239,34 @@ fn _resource_rep(handle: u32) -> *mut u8
 
                     "#
             );
+            let box_path = self.path_to_box();
+            let camel = resource_name.to_pascal_case();
+            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.
+/// 
+/// The pointed object needs to live for the entire lifecycle of the resource and
+/// should only be freed once via the matching resource_from_raw_ function.
+fn resource_into_raw_(val: {camel}Storage<Self>) -> *mut {camel}Storage<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 (call exactly once and only on pointers received from resource_into_raw).
+unsafe fn resource_from_raw_(handle: *mut {camel}Storage<Self>) -> {camel}Storage<Self> where Self: Sized
+{{
+    *unsafe {{ {box_path}::from_raw(handle) }}
+}}
+
+                    "#
+            );
             for method in methods {
                 self.src.push_str(method);
             }
@@ -2716,7 +2744,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#"
@@ -2727,6 +2754,8 @@ pub struct {camel} {{
 }}
 
 type _{camel}Rep<T> = Option<T>;
+/// Data type for arena allocation of resources
+pub type {camel}Storage<T> = Option<T>;
 
 impl {camel} {{
     /// Creates a new resource from the specified representation.
@@ -2737,8 +2766,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 +2825,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,103 @@
+include!(env!("BINDINGS"));
+
+use crate::exports::test::arena_allocated_resources::to_test::{Guest, GuestThing, ThingStorage};
+
+export!(Component);
+
+struct Component;
+
+impl Guest for Component {
+    type Thing = MyThing;
+}
+
+mod arena {
+
+    use core::sync::atomic::{AtomicUsize, Ordering};
+    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: [UnsafeCell<MaybeUninit<T>>; SIZE],
+        offset: AtomicUsize,
+    }
+
+    // Element allocation is atomic and elements are exclusively handed out after allocation,
+    // so the arena can be send to other threads and simultaneosly accessed by multiple threads
+    unsafe impl<T: Sync, const SIZE: usize> Sync for Arena<T, SIZE> {}
+    unsafe impl<T: Send, const SIZE: usize> Send for Arena<T, SIZE> {}
+
+    impl<T: Default, const SIZE: usize> Arena<T, SIZE> {
+        pub const fn new() -> Self {
+            Self {
+                buffer: [const { UnsafeCell::new(MaybeUninit::uninit()) }; SIZE],
+                offset: AtomicUsize::new(0),
+            }
+        }
+
+        /// 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(&self) -> Option<&mut T> {
+            // short circuit the exhausted state (don't increment if full)
+            if self.offset.load(Ordering::Relaxed) >= SIZE {
+                None
+            } else {
+                // now try to allocate for real
+                let pos = self.offset.fetch_add(1, Ordering::Acquire);
+                if pos >= SIZE {
+                    // now self.offset is already beyond SIZE, reduce our increment and return none
+                    self.offset.fetch_sub(1, Ordering::Release);
+                    None
+                } else {
+                    let ptr = self.buffer[pos].get();
+                    // SAFETY: we demand exclusive ownership of the item in the arena
+                    let uninit = unsafe { &mut *ptr };
+                    Some(uninit.write(Default::default()))
+                }
+            }
+        }
+    }
+}
+
+use arena::Arena;
+
+#[derive(Clone)]
+struct MyThing {
+    contents: u32,
+}
+
+static ARENA: Arena<Option<MyThing>, 4> = Arena::new();
+
+impl GuestThing for MyThing {
+    fn new(v: u32) -> MyThing {
+        MyThing { contents: v }
+    }
+
+    fn get(&self) -> u32 {
+        self.contents
+    }
+
+    fn resource_into_raw_(val: ThingStorage<Self>) -> *mut ThingStorage<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 ThingStorage<Self>) -> ThingStorage<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();
+}