Dependency Injection in Rust with Type-Maps

26 Feb 22 10 min read

In this article we will look at type-maps and how they can help enable dependency injection in Rust. We will go a few steps further to look at how we can provide auto wiring of arguments to a given callable, utilising generics and a few traits. This is a pattern seen in frameworks such as Bevy and Actix-Web to provide shared state back to the user.

Type Maps

It is common in dynamically typed languages, such as PHP, to use a service container to aid in dependency injection. Dependency injection containers are usually represented as some sort of map/dictionary where a string representation of the type is set as the key and an instance of the type is set as the value. We can do something similar at compile time in Rust using a type-map as the core of the container.

A type-map is a container that stores some value against its given unique type identifier in order for state to be shared within an application. They are usually represented as a form of HashMap where the key is std::any::TypeId and the value is std::any::Any. This can add a dynamic element to an otherwise statically and strongly-typed systems language.

I first stumbled upon this pattern when looking at systems in the Bevy game engine. At the time, I couldn’t understand how the game engine knew how to wire up components, queries, etc. to the systems, solely based on the type signatures of the function arguments. I then came across it again in the Actix Web framework where it is used to pass shared application state, like a database connection, to multiple handlers. It was after researching the intricacies of this pattern that I decided it would be best to keep some notes for myself, and thus I ended up with this article.

Let’s have a look at how we can create a basic dependency injection container in Rust that can auto wire function arguments.

A simple type-map

Let’s start by building a simple type-map to demonstrate what they are. A test is always a good place to start:

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_type_can_be_bound_and_resolved() {
        let mut container = TypeMap::default();
        container.bind::<i32>(42);
        assert_eq!(container.get::<i32>(), Some(&42));
    }
}

This should be fairly self-explanatory. First, we create a container and bind the value 42 to the type i32. Then, we assert that the value matching the given type resolves to the correct value. As is common practice in Rust, our get method returns a reference to the contained value. We will have a look at other ways to resolve our values towards the end of the article.

You may be wondering what the strange looking method_name::<T>() syntax is all about. This is commonly known as the turbofish operator and is Rust’s way of specifying the type for a generic parameter. Rust is smart enough to infer the types we are asking for, so we can actually remove the turbofish operators here.

#[test]
fn a_type_can_be_bound_and_resolved_through_inference() {
    let mut container = TypeMap::default();
    container.bind(42);
    assert_eq!(container.get(), Some(&42));
}

Let’s have a look at the implementation:

use std::{
    any::{Any, TypeId},
    collections::HashMap,
};

#[derive(Default)]
pub struct TypeMap {
    bindings: HashMap<TypeId, Box<dyn Any>>,
}

impl TypeMap {
    pub fn bind<T: Any>(&mut self, val: T) {
        self.bindings.insert(val.type_id(), Box::new(val));
    }

    pub fn get<T: Any>(&self) -> Option<&T> {
        self.bindings
            .get(&TypeId::of::<T>())
            .and_then(|boxed| boxed.downcast_ref())
    }
}

Things are starting to get interesting, but this is still a relatively straightforward implementation. We represent the contained bindings with a HashMap<TypeId, Box<dyn Any>>. The value is boxed in order for the size of the type implementing dyn Any to be known at compile time.

TypeId::of::<T>() returns the TypeId of the type the generic function has been instantiated with. We use it in get to represent the type when accessing the HashMap and implicitly through val.type_id() in bind. This is possible due to T implementing the std::any::Any trait. Finally, the value must be unboxed and downcast to a reference to allow shared access.

If we run our tests we should now see everything passing:

┌nick@Nicks-MBP ~/c/typemap (master +%)
└> cargo test
Finished test [unoptimized + debuginfo] target(s) in 0.01s
Running unittests (target/debug/deps/typemap-2c921a98b228541f)

running 2 tests
test tests::a_type_can_be_bound_and_resolved ... ok
test tests::a_type_can_be_bound_and_resolved_through_inference ... ok

test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s
...

Injecting a single dependency

Some service containers that handle automatic injection of dependencies (auto-wiring) will provide a method that the user can pass a callable to in order for its dependencies to be injected via the container. Let’s see how that might work in Rust. We can define another test to demonstrate what we are trying to achieve:

#[test]
fn injects_dependency_based_on_argument_type() {
    let mut container = TypeMap::default();
    container.bind(Data::new(42));
    container.call(|data: Data<i32>| {
        assert_eq!(data.get_ref(), &42);
    })
}

As you can see, there is something that looks pretty “magic” going on here. How does the call method know to pass our Data<i32> into the closure? You may also be wondering why we now need a Data<T> at all – hopefully that will become clear as we demystify this.

Let’s start by looking at the implementation of our call method so that we can get some context.

pub fn call<F, Args>(&self, callable: F)
where
    F: Callable<Args>,
    Args: FromTypeMap,
{
    callable.call(Args::from_type_map(self));
}

So, we take in some function F that implements a Callable trait which allows us to call the function with some arguments Args. Our arguments Args implement a FromTypeMap trait that allows us to build up the arguments from our type-map. These traits are pretty straight forward too. Let’s take a look:

pub trait Callable<Args> {
    fn call(&self, args: Args);
}

pub trait FromTypeMap {
    fn from_type_map(type_map: &TypeMap) -> Self;
}

Okay… so how do these traits allow automatic injection of our argument based on a given type? Hopefully the implementation of FromTypeMap will clear this up. Let’s see how we can implement that for an i32:

impl FromTypeMap for i32 {
    fn from_type_map(type_map: &TypeMap) -> Self {
        type_map.get::<Self>().expect("type not found")
    }
}

We use the get method on the type-map to resolve our type. We then clone the value so that we can return an owned version to the caller. This works fine for an i32 as it implements Copy, but what if we wanted to allow for any type? This is where our Data<T> comes in. Let’s have a look at the definition of this:

pub struct Data<T: ?Sized>(Arc<T>);

impl<T> Data<T> {
    pub fn new(val: T) -> Self {
        Data(Arc::new(val))
    }
}

impl<T: ?Sized> Data<T> {
    pub fn get(&self) -> &T {
        self.0.as_ref()
    }
}

impl<T: ?Sized> Clone for Data<T> {
    fn clone(&self) -> Data<T> {
        Data(self.0.clone())
    }
}

impl<T: ?Sized> Deref for Data<T> {
    type Target = Arc<T>;

    fn deref(&self) -> &Arc<T> {
        &self.0
    }
}

impl<T: ?Sized + 'static> FromTypeMap for Data<T> {
    fn from_type_map(type_map: &TypeMap) -> Self {
        type_map.get::<Self>().expect("type not found").clone()
    }
}

If we start by looking at the implementation of the FromTypeMap trait for our Data<T> we can see that the method body is the same as our i32 version. The interesting stuff comes from the defined methods and the implementation of the Clone trait that, together, make this possible. We wrap the T in an Arc which provides us with two benefits – the first being that it gives us thread safety and the second being that it allows us to clone our value easily and safely. Clone is required so that we can actually clone the Data<T> which we do by just cloning the inner Arc<T> and returning a new Data with the cloned value inside. Implementing Deref gives us some syntactic sugar so that we can access the inner value’s methods via the dot notation (Arc also implements Defef). Let’s write a test to show this in action and to also prove that we can pass any type:

#[test]
fn the_values_methods_can_be_accessed_through_deref() {
    let mut container = TypeMap::default();
    container.bind(Data::new(String::from("test test 123")));
    container.call(|data: Data<String>| {
        assert_eq!(data.as_str(), "test test 123");
    });
}

With this in place, our Data<T> container will allow us to have any optionally Sized type injected into our function if we wrap it in our Data<T> container, which is pretty nice.

Now for the last piece of the puzzle. How do we get our Data<T> into our closure? Well, that’s what the Callable trait is for:

impl<Func, Arg1> Callable<Arg1> for Func
where
    Func: Fn(Arg1),
{
    fn call(&self, arg1: Arg1) {
        (self)(arg1);
    }
}

Here, we are implementing the Callable trait for a closure that takes one argument (for now) and returns nothing. In the body of our call method, we just call the closure with the given argument. This is where our limitation of injecting a single dependency comes into play as we have to be explicit about the number of arguments we are passing in.

Injecting multiple dependencies

In order to allow for multiple arguments being passed into our Callables, we need to define a single type that we can implement FromTypeMap on in order to inject our Args into the call method as a single value. We can do this with a tuple, with the only downside being that we will still have to be explicit about the number of arguments we are allowing in the Callable.

Let’s update our code to use a tuple for the Callable with a single argument:

impl<Func, A> Callable<(A,)> for Func
where
    Func: Fn(A),
{
    fn call(&self, (a,): (A,)) {
        (self)(a)
    }
}

Next, we will need to implement the FromTypeMap trait for a single value tuple so that we can build up its members automatically from the type-map:

impl<A: FromTypeMap> FromTypeMap for (A,) {
    fn from_type_map(type_map: &TypeMap) -> Self {
        (A::from_type_map(type_map),)
    }
}

Here, we are saying that we want to implement the FromTypeMap trait for a tuple that has a T which also implements the FromTypeMap trait. This allows us to then call T::from_type_map(type_map) to build the argument for the tuple that we are building from the container.

This may be more clear when we look at how we would do that for a two-argument tuple and its Callable implementation:

impl<Func, A, B> Callable<(A, B)> for Func
where
    Func: Fn(A, B),
{
    fn call(&self, (a, b): (A, B)) {
        (self)(a, b)
    }
}

impl<A: FromTypeMap, B: FromTypeMap> FromTypeMap for (A, B) {
    fn from_type_map(type_map: &TypeMap) -> Self {
        (A::from_type_map(type_map), B::from_type_map(type_map))
    }
}

We could then carry on implementing this for as many arguments as we want to allow in the API of our type-map. However, this is a lot of repeated code that is more difficult to maintain. A common approach to solving this is to allow up to a reasonably-sized fixed number of arguments to be passed to our callable for dependency resolution, and to generate this using a macro:

macro_rules! callable_tuple ({ $($param:ident)* } => {
    impl<Func, $($param,)*> Callable<($($param,)*)> for Func
    where
        Func: Fn($($param),*),
    {
        #[inline]
        #[allow(non_snake_case)]
        fn call(&self, ($($param,)*): ($($param,)*)) {
            (self)($($param,)*)
        }
    }
});

callable_tuple! {}
callable_tuple! { A }
callable_tuple! { A B }
callable_tuple! { A B C }

macro_rules! tuple_from_tm {
        ( $($T: ident )+ ) => {
            impl<$($T: FromTypeMap),+> FromTypeMap for ($($T,)+)
            {
                #[inline]
                fn from_type_map(type_map: &TypeMap) -> Self {
                    ($($T::from_type_map(type_map),)+)
                }
            }
        };
    }

tuple_from_tm! { A }
tuple_from_tm! { A B }
tuple_from_tm! { A B C }

The details on how macros work is out of the scope of this article, but the above is essentially saying:

Generate the code to implement the Callable and FromTypeMap trait for tuples with the argument length of the given number of arguments in the macro.

The letters in the macro call represent the name of the type of each individual argument in the tuple. We can then specify up to a set number of arguments allowed in our API. In most cases this is done for up to twelve arguments as seen in the final code below.

The final code

Here is a link to the final code where you can find everything we have implemented above, as well as some comments and additional tests to help explain what’s going on.

Further reading

  • Bevy - The bevy_reflect crate uses a much more advanced version of this pattern to build up the arguments for systems.
  • Actix Web - I took a lot of information from Actix Web as it uses this pattern to build the extractors for requests.