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

* Resource allocation override option

* custom resource allocation

* cut the dirty tricks out of the arena

* hide that it is an Option (typedef), put underscore to the end

* Some review comments of mine and refactorings

* Fix doc example

---------

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

Author: cpetig

crates/guest-rust/src/examples/_4_exported_resources.rs

@@ -23,4 +23,5 @@ crate::generate!({
             }
         }
     "#,
+    runtime_path: "crate::rt", // only needed for this in-crate example.
 });

crates/guest-rust/src/lib.rs

@@ -887,6 +887,8 @@ pub mod examples;
 #[doc(hidden)]
 pub mod rt;
 
+pub mod resource;
+
 #[cfg(feature = "inter-task-wakeup")]
 pub use rt::async_support::UnitStreamOps;
 #[cfg(feature = "async-spawn")]

crates/guest-rust/src/resource.rs

@@ -0,0 +1,57 @@
+//! Helper traits, types, and utilities for managing resources in the component
+//! model.
+
+/// A trait implemented by all resources that a component might export.
+///
+/// This is an implementation detail primarily for the code generated by
+/// exported resources. The primary purpose of this trait is to serve as an
+/// abstraction for the in-memory storage of a resource.
+pub trait Resource: Sized + 'static {
+    /// The type which is actually stored in-memory for this resource.
+    ///
+    /// By default this is `Option<T>`.
+    type Rep: ResourceRep<Self>;
+}
+
+impl<T: 'static> Resource for T {
+    type Rep = Option<T>;
+}
+
+/// A trait used to define how to access the underlying data `T` from an
+/// in-memory representation.
+///
+/// This is used as a bound on the [`Resource::Rep`] associated type which is in
+/// turn used to access data within a resources.
+pub unsafe trait ResourceRep<T> {
+    /// Creates a new instance of `Self` which wraps the provided data.
+    fn rep_new(inner: T) -> Self;
+
+    /// Acquires `&T` from a raw pointer to `Self`.
+    unsafe fn rep_as_ref<'a>(ptr: *const Self) -> &'a T;
+
+    /// Acquires `&mut T` from a raw pointer to `Self`.
+    unsafe fn rep_as_mut<'a>(ptr: *mut Self) -> &'a mut T;
+
+    /// Takes the value out of `Self` at the provided pointer.
+    ///
+    /// Note that `ptr` will later be deallocated meaning that it must not run
+    /// the destructor of `T` after this method is called. This is guaranteed to
+    /// be called at most once, however. Additionally after calling this method
+    /// it's guaranteed that the `rep_as_*` method above will not be called.
+    unsafe fn rep_take<'a>(ptr: *mut Self) -> T;
+}
+
+unsafe impl<T> ResourceRep<T> for Option<T> {
+    fn rep_new(inner: T) -> Option<T> {
+        Some(inner)
+    }
+    unsafe fn rep_as_ref<'a>(ptr: *const Option<T>) -> &'a T {
+        unsafe { (*ptr).as_ref().unwrap() }
+    }
+    unsafe fn rep_as_mut<'a>(ptr: *mut Option<T>) -> &'a mut T {
+        unsafe { (*ptr).as_mut().unwrap() }
+    }
+    unsafe fn rep_take(ptr: *mut Option<T>) -> T {
+        unsafe { (*ptr).take().unwrap() }
+    }
+}

crates/guest-rust/src/rt/mod.rs

@@ -1,6 +1,8 @@
 use core::alloc::Layout;
 use core::ptr::{self, NonNull};
 
+pub use crate::resource::{Resource, ResourceRep};
+
 // Re-export `bitflags` so that we can reference it from macros.
 #[cfg(feature = "bitflags")]
 pub use bitflags;

crates/rust/src/bindgen.rs

@@ -406,7 +406,9 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 let result = if is_own {
                     format!("{name}::from_handle({op} as u32)")
                 } else if self.r#gen.is_exported_resource(*resource) {
-                    format!("{name}Borrow::lift({op} as u32 as usize)")
+                    format!(
+                        "{name}Borrow::lift(core::ptr::with_exposed_provenance({op} as u32 as usize))"
+                    )
                 } else {
                     let tmp = format!("handle{}", self.tmp());
                     self.handle_decls.push(format!("let {tmp};"));

crates/rust/src/interface.rs

@@ -196,9 +196,10 @@ impl<'i> InterfaceGenerator<'i> {
                     .map(|(resource, (trait_name, ..))| (resource.unwrap(), trait_name.as_str())),
             )
         }
+        let rt = self.r#gen.runtime_path().to_string();
 
         for (resource, (trait_name, methods)) in traits.iter() {
-            uwriteln!(self.src, "pub trait {trait_name}: 'static {{");
+            uwriteln!(self.src, "pub trait {trait_name}: {rt}::Resource {{");
             let resource = resource.unwrap();
             let resource_name = self.resolve.types[resource].name.as_ref().unwrap();
             let (_, interface_name) = interface.unwrap();
@@ -239,6 +240,55 @@ fn _resource_rep(handle: u32) -> *mut u8
 
                     "#
             );
+            let box_path = self.path_to_box();
+            uwriteln!(
+                self.src,
+                r#"
+/// Place this resource's representation into a location with a stable pointer,
+/// returning a pointer to that location.
+///
+/// This method is used to place `val` on the heap, for example, or possibly in
+/// an arena. The returned pointer must remain valid for the lifetime of this
+/// resource. The default implementation of this metho will place `val` onto the
+/// heap with `Box`. The returned pointer will be deallocated by the sibling
+/// `resource_from_raw_` function.
+///
+/// # Safety
+///
+/// Note that this method is not `unsafe` to call, but it is `unsafe` to
+/// define. When overriding this method you must additionally override the
+/// `resource_from_raw_` to insert an appropriate deallocation for the returned
+/// pointer. In the future this might become a default associated trait bound,
+/// but that's not stable in Rust right now.
+#[doc(hidden)]
+unsafe fn resource_into_raw_(val: Self::Rep) -> *mut Self::Rep
+{{
+    {box_path}::into_raw({box_path}::new(val))
+}}
+
+/// Consumes and deallocates a pointer previously returned by
+/// `resource_into_raw_`.
+///
+/// This function is used to deallocate resources and allocations associated
+/// with the allocation routine when creating a resource. The default
+/// implementation of this method will read `handle`'s value and deallocate it
+/// as a `Box`.
+///
+/// Note that when overriding this method you'll almost surely want to override
+/// `resource_into_raw_` as well.
+///
+/// # Safety
+///
+/// This method is only safe to call with pointers previously created by
+/// `resource_into_raw_`. For more information see `Box::from_raw` for examples.
+#[doc(hidden)]
+unsafe fn resource_from_raw_(handle: *mut Self::Rep) -> Self::Rep
+{{
+    *unsafe {{ {box_path}::from_raw(handle) }}
+}}
+
+                    "#
+            );
             for method in methods {
                 self.src.push_str(method);
             }
@@ -2716,7 +2766,7 @@ 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();
+            let rt = self.r#gen.runtime_path();
             uwriteln!(
                 self.src,
                 r#"
@@ -2726,8 +2776,6 @@ pub struct {camel} {{
     handle: {resource}<{camel}>,
 }}
 
-type _{camel}Rep<T> = Option<T>;
-
 impl {camel} {{
     /// Creates a new resource from the specified representation.
     ///
@@ -2736,31 +2784,30 @@ impl {camel} {{
     /// create a handle. The owned handle is then returned as `{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 rep = <T::Rep as {rt}::ResourceRep<T>>::rep_new(val);
         unsafe {{
+            let ptr = T::resource_into_raw_(rep);
             Self::from_handle(T::_resource_new(ptr.cast()))
         }}
     }}
 
     /// Gets access to the underlying `T` which represents this resource.
     pub fn get<T: Guest{camel}>(&self) -> &T {{
-        let ptr = unsafe {{ &*self.as_ptr::<T>() }};
-        ptr.as_ref().unwrap()
+        let ptr = self.as_ptr::<T>();
+        unsafe {{ <T::Rep as {rt}::ResourceRep<T>>::rep_as_ref(ptr) }}
     }}
 
     /// Gets mutable access to the underlying `T` which represents this
     /// resource.
     pub fn get_mut<T: Guest{camel}>(&mut self) -> &mut T {{
-        let ptr = unsafe {{ &mut *self.as_ptr::<T>() }};
-        ptr.as_mut().unwrap()
+        let ptr = self.as_ptr::<T>();
+        unsafe {{ <T::Rep as {rt}::ResourceRep<T>>::rep_as_mut(ptr) }}
     }}
 
     /// Consumes this resource and returns the underlying `T`.
     pub fn into_inner<T: Guest{camel}>(self) -> T {{
-        let ptr = unsafe {{ &mut *self.as_ptr::<T>() }};
-        ptr.take().unwrap()
+        let ptr = self.as_ptr::<T>();
+        unsafe {{ <T::Rep as {rt}::ResourceRep<T>>::rep_take(ptr) }}
     }}
 
     #[doc(hidden)]
@@ -2783,7 +2830,7 @@ impl {camel} {{
     // It's theoretically possible to implement the `Guest{camel}` trait twice
     // so guard against using it with two different types here.
     #[doc(hidden)]
-    fn type_guard<T: 'static>() {{
+    fn type_guard<T: {rt}::Resource>() {{
         use core::any::TypeId;
         static mut LAST_TYPE: Option<TypeId> = None;
         unsafe {{
@@ -2797,12 +2844,14 @@ impl {camel} {{
     }}
 
     #[doc(hidden)]
-    pub unsafe fn dtor<T: 'static>(handle: *mut u8) {{
+    pub unsafe fn dtor<T: Guest{camel}>(handle: *mut u8) {{
         Self::type_guard::<T>();
-        let _ = unsafe {{ {box_path}::from_raw(handle as *mut _{camel}Rep<T>) }};
+        unsafe {{
+            let _rep = T::resource_from_raw_(handle.cast());
+        }}
     }}
 
-    fn as_ptr<T: Guest{camel}>(&self) -> *mut _{camel}Rep<T> {{
+    fn as_ptr<T: Guest{camel}>(&self) -> *mut T::Rep {{
        {camel}::type_guard::<T>();
        T::_resource_rep(self.handle()).cast()
     }}
@@ -2813,29 +2862,29 @@ impl {camel} {{
 #[derive(Debug)]
 #[repr(transparent)]
 pub struct {camel}Borrow<'a> {{
-    rep: *mut u8,
+    rep: *const u8,
     _marker: core::marker::PhantomData<&'a {camel}>,
 }}
 
 impl<'a> {camel}Borrow<'a>{{
     #[doc(hidden)]
-    pub unsafe fn lift(rep: usize) -> Self {{
+    pub unsafe fn lift(rep: *const u8) -> Self {{
         Self {{
-            rep: rep as *mut u8,
+            rep,
             _marker: core::marker::PhantomData,
         }}
     }}
 
     /// Gets access to the underlying `T` in this resource.
     pub fn get<T: Guest{camel}>(&self) -> &'a T {{
-       let ptr = unsafe {{ &mut *self.as_ptr::<T>() }};
-       ptr.as_ref().unwrap()
+       let ptr = self.as_ptr::<T>();
+       unsafe {{ <T::Rep as {rt}::ResourceRep<T>>::rep_as_ref(ptr) }}
     }}
 
     // NB: mutable access is not allowed due to the component model allowing
     // multiple borrows of the same resource.
 
-    fn as_ptr<T: 'static>(&self) -> *mut _{camel}Rep<T> {{
+    fn as_ptr<T: Guest{camel}>(&self) -> *const T::Rep {{
        {camel}::type_guard::<T>();
        self.rep.cast()
     }}

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,96 @@
+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::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
+    }
+
+    unsafe fn resource_into_raw_(val: Self::Rep) -> *mut Self::Rep {
+        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 Self::Rep) -> Self::Rep {
+        unsafe { &mut *handle }.take()
+    }
+}