Demilade Sonuga's blog

All posts

Implementing Trait Objects I

2023-04-20 · 29 min read

This is the last thing we need to complete before we can fully get on with event handling (I promise this time). In this post, we're going to be implementing boxed trait objects for functions.

The Idea

The main idea behind the implementation I'm presenting here is that, in the C programming language, a pointer to a struct cast as a pointer to something else is guaranteed to keep pointing to the same address. To understand this, consider the following Rust-like pseudocode:

struct Base {
    drop(pointer to self);
    clone(pointer to self);
}

struct Child {
    Base base;
    some_other_func();
}

In the above, Base is a structure of two function pointers: drop and clone. While Child is a struct with the first field as a Base instance and the second field as another function pointer.

According to C's guarantees, casting a pointer to Child as a pointer to Base is valid because the first field in Child is a valid Base.

For example, suppose you had a pointer to a Child instance:

let child_ptr = &child as *mut Child;

When you do this, the compiler's view of child's memory will look something like this:

Memory:

0:  base - drop
1:       - clone
2:  some_other_func

A Base instance is at address 0 and together with the some_other_func function pointer is a Child instance. child_ptr will have the value 0 since that is the address of the first byte.

Casting this pointer as a pointer to Base, like so:

let base_ptr = child_ptr as *mut Base;

just tells the compiler: "Yo! Interpret the value starting at this address as a Base, not a Child anymore. Alright?"

The compiler then goes: "Okay, since only the bytes in addresses 0 and 1 make up the Base, I'll keep looking at that and forget about address 2 (and the other addresses that come after it)"

The compiler's view then becomes:

Memory:

0:  base - drop
1:       - clone

Like the some_other_func never existed. It just becomes a regular Base. The other really wonderful thing about it all is that (if we remember correctly), base_ptr was once a pointer to a Child, and the some_other_func never got overwritten. So, we can still cast the pointer base_ptr back to a Child pointer.

This is all just a very simple implementation of polymorphism. If you're interested in finding out more about this idea, check the references.

Aims

Our function trait object should be easily created like so:

let boxed_fn = BoxedFn::new(a_function);

And called like so:

boxed_fn(...arguments required);

Implementation

In this implementation, we're going to start from a structure called BaseFn. You can think of BaseFn as the Base struct in the previous section. It specifies functions that all our function trait objects must have:

Create a new file, boxed_fn.rs, and type this in:

use crate::event_hook::EventInfo;

// The polymorphic representation of a base function which can stand in
// for any event handler function
#[repr(C)]
struct BaseFn {
    call_boxed: fn(pointer to data, EventInfo) -> (),
    drop: fn(pointer to data) -> (),
    clone: fn(pointer to data) -> BoxedFn
}

In main.rs

mod vec;
mod boxed_fn;
mod event_hook; // NEW
mod allocator;

In event_hook.rs

use crate::keyboard::KeyEvent; // NEW

/* COMMENTING THIS OUT, FOR NOW TO AVOID FUTURE ERRORS
// Mediator between the game code and the interrupt service routines
pub struct EventHooker {
    // The functions that will be called when an event occurs
    handlers: Vec<Box<dyn FnMut(EventInfo)>>
}
*/

// The info about an event that will be passed into a handler
// when it is called
pub enum EventInfo {
    Timer,
    Keyboard(KeyEvent)
}

Remember from the previous post that a trait object is a structure holding a pointer to a method table and a pointer to data of a specific type that implements that trait. This BaseFn is our trait object's method table.

The call_boxed function is the method that will serve to call the actual functions themselves like boxed_fn(). The drop function is the method that will serve as the destructor. And the clone serves to, well, clone.

call_boxed takes an EventInfo because, earlier, we imposed the restriction that all event handler functions must take a single argument of the type EventInfo.

Instances of this BaseFn are going to have pointers to different implementations of drop and clone and call_boxed, for each separate function trait object.

Next, we'll create the trait object itself:

use crate::allocator::Allocator;
// A stand in for `Box<dyn FnMut>`
pub struct BoxedFn(*mut BaseFn, *mut dyn Allocator);

This is the data that gets passed to the BaseFn functions:

#[repr(C)]
struct BaseFn {
    call_boxed: fn(*const BoxedFn, EventInfo) -> (),
    drop: fn(*const BoxedFn) -> (),
    clone: fn(*const BoxedFn) -> BoxedFn
}

This struct will serve as our trait object. If you look very closely, you'll notice that this doesn't fully fit our picture of trait objects since it was mentioned that they are pointers to a method table and data. This struct only contains the pointer to the method table, but not the data.

This is really not a problem. It's just an implementation choice and you can do it another way if you want to.

Instead of keeping another pointer to the function/function-like thing (stuff like closures), we do this:

#[repr(C)]
struct Repr<F: FnMut(EventInfo)> {
    base: BaseFn,
    func: F
}

By annotating this with #[repr(C)], we tell Rust to behave like a C compiler when dealing with this struct. This enables us to get a pointer to BaseFn from a Repr by simply casting a *mut Repr to *mut BaseFn and forgetting about the func field. We add the trait bound F: FnMut(EventInfo) to make the compiler ensure that func is a function-like thing that can be called and takes an EventInfo as an argument.

Creating a new BoxedFn is now just a matter of doing the following:

  1. Accepting the function or function-like thing that we want boxed.
  2. Creating call_boxed, drop, and clone functions for the thing we want boxed and putting them together into a BaseFn instance.
  3. Creating a Repr from the new BaseFn and the function-like thing.
  4. Putting the Repr on the heap, so that it will outlive the function.
  5. Retrieving the pointer to the Repr and casting it to a *mut BaseFn.
  6. Use that *mut BaseFn to create the new BoxedFn.

The only step in this list that doesn't seem easy to resolve is number 2. How do we create call_boxed, drop, and clone functions for each of those things we want boxed? What if there are a hundred of them? Do we then start writing one hundred implementations each? No. That is so unprogrammerly.

You have to think: what mechanism does Rust offer to us that allows us to automatically create different versions of a function based on their different types? And how do we make use of it? Make sure you answer these questions before reading on.

Generics, of course. Remember that the Rust compiler monomorphizes generic functions, that is it creates multiple copies of the function for each type that takes the place of its generic.

We can use this to our advantage. Let's start with call_boxed. The role of this function is to call the boxed function thing.

// Calls the concrete function wrapped by the BoxedFn
fn call_boxed<F>(boxed_fn_ptr: *const BoxedFn, event: EventInfo) where F: FnMut(EventInfo) {
    unsafe {
        let concrete_repr_ptr = (*boxed_fn_ptr).0 as *mut Repr<F>;
        ((*concrete_repr_ptr).func)(event) 
    }
}

This code is very straightforward. First, it casts the BoxedFn's *mut BaseFn pointer as a Repr (We can do this because a *mut BaseFn in a BoxedFn points to the first field of a Repr). Then it simply calls the function thing.

For the drop implementation, all we need to do is deallocate the memory being used by the Repr:

use crate::boxed::Box;
// Drops the boxed function
fn drop<F>(boxed_fn_ptr: *const BoxedFn) where F: FnMut(EventInfo) {
    unsafe {
        let base_fn_ptr = (*boxed_fn_ptr).0;
        let concrete_ptr: *mut Repr<F> = base_fn_ptr as *mut Repr<F>;
        let allocator = (*boxed_fn_ptr).1;
        Box::<Repr<F>>::from_raw(concrete_ptr, allocator);
        // Box is dropped at the end of the scope
    }
}

The above simply casts the BoxedFn's BaseFn pointer as a *mut Repr, constructs Box from it and allows the Box to get dropped. The Box's Drop implementation then handles the memory deallocation.

The purpose of clone is to create a new identical BoxedFn:

// Clones the boxed function
//
// # Safety
//
// Cloning a BoxedFn is highly unsafe. The function may contains mutable
// references to the outer scope. Cloning the BoxedFn will result in cloning
// mutable references, defeating Rust's safety guarantees.
fn clone<F>(boxed_fn_ptr: *const BoxedFn) -> BoxedFn where F: FnMut(EventInfo) {
    unsafe {
        let base_fn_ptr = (*boxed_fn_ptr).0;
        let concrete_ptr: *mut Repr<F> = base_fn_ptr.cast::<Repr<F>>();
        let allocator = (*boxed_fn_ptr).1;
        let func_ptr = &(*concrete_ptr).func as *const F;
        BoxedFn::new(func_ptr.read(), allocator)
    }
}

Again, the pointer casting stuff, but the function thing gets copied out with read and a new BoxedFn is created from the copy. This function has some serious safety problems that just may screw us up in the future, so we'll have to be careful with it.

Now, with these, let's create BoxedFn::new.

impl BoxedFn {
    // Creates a new BoxedFn from the given function-thing
    pub fn new<F>(func: F, allocator: *mut dyn Allocator) -> Self where F: FnMut(EventInfo) {
        
    }
}

Steps 1 and 2, putting the function thing's BaseFn functions together into a BaseFn instance:

impl BoxedFn {
    // Creates a new BoxedFn from the given function-thing
    pub fn new<F>(func: F, allocator: *mut dyn Allocator) -> Self where F: FnMut(EventInfo) {
        let base_fn = BaseFn { call_boxed: call_boxed::<F>, drop: drop::<F>, clone: clone::<F> };
    }
}

Notice that new is a generic function. Because of this, call_boxed, drop, and clone will have separate copies for each type that is passed as the func argument to BoxedFn::new.

Step 3, creating a Repr from the new BaseFn` and the function-thing:

impl BoxedFn {
    // Creates a new BoxedFn from the given function-thing
    pub fn new<F>(func: F, allocator: *mut dyn Allocator) -> Self where F: FnMut(EventInfo) {
        let base_fn = BaseFn { call_boxed: call_boxed::<F>, drop: drop::<F>, clone: clone::<F> };
        // NEW:
        let concrete_repr = Repr {
            base: base_fn,
            func
        };
    }
}

Steps 4, 5, and 6: putting the Repr on the heap, retrieving the pointer and casting it to *mut BaseFn, then creating a new BoxedFn from that *mut BaseFn.

impl BoxedFn {
    // Creates a new BoxedFn from the given function-thing
    pub fn new<F>(func: F, allocator: *mut dyn Allocator) -> Self where F: FnMut(EventInfo) {
        let base_fn = BaseFn { call_boxed: call_boxed::<F>, drop: drop::<F>, clone: clone::<F> };
        let concrete_repr = Repr {
            base: base_fn,
            func
        };
        // NEW:
        let concrete_repr_ptr: *mut Repr<F> = Box::<Repr<F>>::into_raw(Box::new(concrete_repr, allocator));
        let polymorphic_ptr: *mut BaseFn = concrete_repr_ptr as *mut BaseFn;
        BoxedFn(polymorphic_ptr, allocator)
    }
}

Notice that we're using the Box's into_raw function. This function first calls ManuallyDrop on the Box, so that the Box won't get dropped at the end of the function.

And that's it for creating a new BoxedFn.

As for Drop and Clone implementations:

impl Drop for BoxedFn {
    fn drop(&mut self) {
        let base_fn_ptr = self.0;
        unsafe { ((*base_fn_ptr).drop)(self as *const BoxedFn) };
    }
}

impl Clone for BoxedFn {
    fn clone(&self) -> Self {
        let base_fn_ptr = self.0;
        unsafe { ((*base_fn_ptr).clone)(self as *const BoxedFn) }
    }
}

We make use of the implementations in the BaseFn, which serve as the method table.

Finally, we need to make BoxedFn callable, so we can do something like this: boxed_fn(event).

To do this, we need to implement FnOnce, Fn, and FnMut, the traits that indicate that a structure is callable. If you need a refresher on these traits, check this post.

The thing about these function traits is that they're unstable, so we need to add modify main.rs:

#![cfg_attr(not(test), no_std)]
#![cfg_attr(not(test), no_main)]
#![feature(abi_efiapi, abi_x86_interrupt)]
#![feature(panic_info_message)]
#![feature(unboxed_closures, fn_traits)] // NEW
#![cfg_attr(test, feature(allocator_api))]

First, we'll implement Fn:

impl Fn<(EventInfo,)> for BoxedFn {
    extern "rust-call" fn call(&self, args: (EventInfo,)) -> Self::Output {
        let base_fn_ptr = self.0;
        unsafe { ((*base_fn_ptr).call_boxed)(self as *const BoxedFn, args.0) }
    }
}

This implementation simply calls the call_boxed function in the method table (in the BaseFn instance).

The remaining implementations will just call this call function.

impl FnMut<(EventInfo,)> for BoxedFn {
    extern "rust-call" fn call_mut(&mut self, args: (EventInfo,)) -> Self::Output {
        self.call(args)
    }
}

impl FnOnce<(EventInfo,)> for BoxedFn {
    type Output = ();
    extern "rust-call" fn call_once(self, args: (EventInfo,)) -> Self::Output {
        self.call(args)
    }
}

And that's it for our implementation.

Take Away

For the full code, go to the repo

In The Next Post

We'll be testing our boxed functions.

References

← Previous Post

A Few Things About Trait Objects

Next Post →

Implementing Trait Objects II