24 How to Use Trait Objects in Practical Applications of the Type System

24 Type System: How to Use Trait Object in Practical Situations? #

Hello, I am Chen Tian.

Today, let’s look at how trait objects are used in practical situations.

As usual, let’s first review what a trait object is. When we want a concrete type to exhibit only the behaviors of a certain trait at runtime, we can assign it to a dyn T, whether it’s &dyn T, Box, or Arc. Here, T is a trait implemented by the current data type. At that point, the original type is erased, Rust will create a trait object and allocate a vtable to fulfill the trait.

You can revisit this image from [lecture 13] to recall what a trait object is: - Image

When compiling dyn T, Rust will generate the corresponding vtable for implementations of the trait that used the trait object type and place it in the executable file (typically in the TEXT or RODATA section): - Image

Thus, when a trait object calls a method of a trait, it first finds the corresponding vtable from the vptr, and then finds the corresponding method to execute.

The benefit of using trait objects is that when a type that satisfies a certain trait is required in a context, and such types might be numerous, with the current context unable to determine which type it will get, we can use trait objects to uniformly handle behavior. Like generic parameters, trait objects are a form of late binding, allowing decisions to be delayed to runtime, thereby gaining the greatest flexibility.

Of course, there are trade-offs. Trait objects push decisions to runtime, resulting in reduced execution efficiency. In Rust, executing a function or method is a simple jump instruction, but executing a trait object method takes an extra step. It involves an additional memory access to find the jump destination, so the execution efficiency is slightly lower.

Furthermore, if a trait object is returned as a return value, or if a trait object is passed between threads, you can’t get away from using Box or Arc, which brings additional heap allocation overhead.

Ok, that’s all for the review of trait objects. If you’re still unclear about it, please review [lecture 13] and read Rust book’s: Using Trait Objects that allow for values of different types. Next, let’s talk about the main usage scenarios of trait objects in practice.

Using Trait Objects in Functions #

We can use trait objects in the parameters or return values of functions.

Let’s look at the use of trait objects in parameters first. The following code constructs an Executor trait, and compares how it’s used for static dispatch with impl Executor, and for dynamic dispatch with &dyn Executor and Box as different types of parameters:

use std::{error::Error, process::Command};

pub type BoxedError = Box<dyn Error + Send + Sync>;

pub trait Executor {
    fn run(&self) -> Result<Option<i32>, BoxedError>;
}

pub struct Shell<'a, 'b> {
    cmd: &'a str,
    args: &'b [&'a str],
}

impl<'a, 'b> Shell<'a, 'b> {
    pub fn new(cmd: &'a str, args: &'b [&'a str]) -> Self {
        Self { cmd, args }
    }
}

impl<'a, 'b> Executor for Shell<'a, 'b> {
    fn run(&self) -> Result<Option<i32>, BoxedError> {
        let output = Command::new(self.cmd).args(self.args).output()?;
        Ok(output.status.code())
    }
}

/// Using generic parameter
pub fn execute_generics(cmd: &impl Executor) -> Result<Option<i32>, BoxedError> {
    cmd.run()
}

/// Using trait object: &dyn T
pub fn execute_trait_object(cmd: &dyn Executor) -> Result<Option<i32>, BoxedError> {
    cmd.run()
}

/// Using trait object: Box<dyn T>
pub fn execute_boxed_trait_object(cmd: Box<dyn Executor>) -> Result<Option<i32>, BoxedError> {
    cmd.run()
}

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

    #[test]
    fn shell_should_work() {
        let cmd = Shell::new("ls", &[]);
        let result = cmd.run().unwrap();
        assert_eq!(result, Some(0));
    }

    #[test]
    fn execute_should_work() {
        let cmd = Shell::new("ls", &[]);

        let result = execute_generics(&cmd).unwrap();
        assert_eq!(result, Some(0));
        let result = execute_trait_object(&cmd).unwrap();
        assert_eq!(result, Some(0));
        let boxed = Box::new(cmd);
        let result = execute_boxed_trait_object(boxed).unwrap();
        assert_eq!(result, Some(0));
    }
}

In the example, impl Executor uses a simplified version of the generic parameter, while &dyn Executor and Box are trait objects, with the former on the stack and the latter allocated on the heap. It is worth noting that the trait object allocated on the heap can also be returned as a return value, such as, in the example, Result<BoxedError> where a trait object is used.

To simplify the code, I used the type keyword to create a BoxedError type, which is an alias for Box, it is a trait object of the Error trait, besides requiring that the type implements the Error trait, it also comes with additional constraints: the type must satisfy both the Send and Sync traits.

Using trait objects in parameters is relatively straightforward. Let’s solidify this with an example from real practice:

pub trait CookieStore: Send + Sync {
    fn set_cookies(
        &self, 
        cookie_headers: &mut dyn Iterator<Item = &HeaderValue>, 
        url: &Url
    );
    fn cookies(&self, url: &Url) -> Option<HeaderValue>;
}

This is a trait for handling CookieStore from the reqwest library we’ve used before. In the set_cookies method, a trait object &mut dyn Iterator is used.

Using Trait Objects in Return Values #

Alright, you should have no problem understanding how to use trait objects in parameters by now, let’s look at using trait objects in return values, which is a scenario where trait objects are used more frequently.

It has already appeared many times before. For example, the previous lecture already explained in detail why the Storage trait in the KV server cannot use generic parameters to handle iterators and can only use Box:

pub trait Storage: Send + Sync + 'static {
    ...
    /// Iterate over the HashTable, returning an Iterator of kv pairs
    fn get_iter(&self, table: &str) -> Result<Box<dyn Iterator<Item = Kvpair>>, KvError>;
}

Let’s look at some real-world examples.

First is async_trait. It is a special kind of trait that contains async fn. Currently Rust does not support using async fn in traits, one workaround is to use the async_trait macro.

In our get hands dirty series, we used an async trait. Below is the Fetch trait defined in [lecture 6] for retrieving SQL query tool data sources:

// Rust's async trait is not stable yet, the async_trait macro can be used
#[async_trait]
pub trait Fetch {
    type Error;
    async fn fetch(&self) -> Result<String, Self::Error>;
}

The macro expansion would look like:

pub trait Fetch {
    type Error;
    fn fetch<'a>(&'a self) -> 
        Result<Pin<Box<dyn Future<Output = String> + Send + 'a>>, Self::Error>;
}

It uses a trait object as a return type. Thus, no matter what type of Future the fetch() implementation returns, it can all be unified into a trait object, allowing the caller to use it just like a normal Future.

Let’s look at another example from snow, the CryptoResolver:

/// An object that resolves the providers of Noise crypto choices
pub trait CryptoResolver {
    /// Provide an implementation of the Random trait or None if none available.
    fn resolve_rng(&self) -> Option<Box<dyn Random>>;

    /// Provide an implementation of the Dh trait for the given DHChoice or None if unavailable.
    fn resolve_dh(&self, choice: &DHChoice) -> Option<Box<dyn Dh>>;

    /// Provide an implementation of the Hash trait for the given HashChoice or None if unavailable.
    fn resolve_hash(&self, choice: &HashChoice) -> Option<Box<dyn Hash>>;

    /// Provide an implementation of the Cipher trait for the given CipherChoice or None if unavailable.
    fn resolve_cipher(&self, choice: &CipherChoice) -> Option<Box<dyn Cipher>>;

    /// Provide an implementation of the Kem trait for the given KemChoice or None if unavailable
    #[cfg(feature = "hfs")]
    fn resolve_kem(&self, _choice: &KemChoice) -> Option<Box<dyn Kem>> {
        None
    }
}

This trait deals with which encryption algorithms are used in the Noise Protocol. Each method of this trait returns a trait object, with each providing different capabilities required for encryption, such as random number generation algorithm (Random), DH algorithm (Dh), hash algorithm (Hash), symmetric cipher algorithm (Cipher), and key encapsulation algorithm (Kem).

All of these have a series of specific algorithm implementations. Through the CryptoResolver trait, you can obtain the trait object of the specific algorithm currently in use. In this way, in processing business logic, we don’t need to care about what algorithm is actually being used and can construct the corresponding implementations based on these trait objects. For example, below is generate_keypair:

pub fn generate_keypair(&self) -> Result<Keypair, Error> {
    // Get the current random number generation algorithm
    let mut rng = self.resolver.resolve_rng().ok_or(InitStage::GetRngImpl)?;
    // Get the current DH algorithm
    let mut dh = self.resolver.resolve_dh(&self.params.dh).ok_or(InitStage::GetDhImpl)?;
    let mut private = vec![0u8; dh.priv_len()];
    let mut public = vec![0u8; dh.pub_len()];
    // Use the random number generator and DH to generate a key pair
    dh.generate(&mut *rng);

    private.copy_from_slice(dh.privkey());
    public.copy_from_slice(dh.pubkey());

    Ok(Keypair { private, public })
}

By the way, if you want to learn more effectively about trait usage and trait objects, snow is an excellent source. You could follow the trail of CryptoResolver to see how they define traits, how they link various traits together, and ultimately how they connect traits with core data structures (tweeter’s Builder and HandshakeState are good starting points).

Using Trait Objects in Data Structures #

After understanding how to use trait objects in functions, let’s see how to use trait objects in data structures.

Continuing with snow’s code as an example, let’s examine which trait objects are used in the HandshakeState data structure for handling Noise Protocol handshake (see code):

pub struct HandshakeState {
    pub(crate) rng:              Box<dyn Random>,
    pub(crate) symmetricstate:   SymmetricState,
    pub(crate) cipherstates:     CipherStates,
    pub(crate) s:                Toggle<Box<dyn Dh>>,
    pub(crate) e:                Toggle<Box<dyn Dh>>,
    pub(crate) fixed_ephemeral:  bool,
    pub(crate) rs:               Toggle<[u8; MAXDHLEN]>,
    pub(crate) re:               Toggle<[u8; MAXDHLEN]>,
    pub(crate) initiator:        bool,
    pub(crate) params:           NoiseParams,
    pub(crate) psks:             [Option<[u8; PSKLEN]>; 10],
    #[cfg(feature = "hfs")]
    pub(crate) kem:              Option<Box<dyn Kem>>,
    #[cfg(feature = "hfs")]
    pub(crate) kem_re:           Option<[u8; MAXKEMPUBLEN]>,
    pub(crate) my_turn:          bool,
    pub(crate) message_patterns: MessagePatterns,
    pub(crate) pattern_position: usize,
}

You don’t need to understand Noise protocol to roughly understand the roles of these three trait objects, Random, Dh, and Kem: they provide the maximum flexibility for the encryption protocol used during the handshake.

Think about it, how would we deal with this data structure if we didn’t use trait objects?

We could use generic parameters, which means:

pub struct HandshakeState<R, D, K>
where
    R: Random,
    D: Dh,
    K: Kem
{
  ...
}

This is what we mostly choose when handling such data structures. However, too many generic parameters lead to two problems: firstly, all related interfaces become very bulky in the implementation process, and wherever you use HandshakeState, you must carry all these generic parameters and their constraints. Secondly, all the cases where these parameters are used, when combined, will generate a lot of code.

Using trait objects, we sacrifice a bit of performance to remove these generic parameters, make the implementation code cleaner and more concise, and ensure there is only one copy of the code.

Another typical scenario for using trait objects in data structures is with closures.

Since in Rust, closures appear as anonymous types, we can’t directly use their types in data structures and can only use generic parameters. However, after using generic parameters with closures, if the captured data is too large, it may cause the data structure itself to be too large. Sometimes, we don’t mind a small performance loss and prefer to make the code easier to handle.

For example, the oso library used for RBAC (oso) includes an Fn in AttributeGetter:

#[derive(Clone)]
pub struct AttributeGetter(
    Arc<dyn Fn(&Instance, &mut Host) -> crate::Result<PolarValue> + Send + Sync>,
);

If you’re interested in how to implement Python’s getattr in Rust, you can check out the oso code.

Another example is dialoguer for interactive CLI, where the validator of an Input is an FnMut:

pub struct Input<'a, T> {
    prompt: String,
    default: Option<T>,
    show_default: bool,
    initial_text: Option<String>,
    theme: &'a dyn Theme,
    permit_empty: bool,
    validator: Option<Box<dyn FnMut(&T) -> Option<String> + 'a>>,
    #[cfg(feature = "history")]
    history: Option<&'a mut dyn History<T>>,
}

Using Trait Objects to Handle KV Server’s Service Structure #

Good, now that we’ve covered several scenarios where trait objects are used for dynamic dispatching, let’s practice writing some code with it.

Take the KV server’s Service structure we wrote before—we’ll try to handle it to use trait objects internally.

Actually, for the KV server, using generics is a better choice because generics don’t introduce too much complexity here, and we don’t want to lose even a tiny bit of performance. However, for the purpose of learning, we can see what the code would look like if the store used a trait object. You can try it yourself first, then see the example below (see code):

use std::{error::Error, sync::Arc};

// Define types to make KV server's traits compile
pub type KvError = Box<dyn Error + Send + Sync>;
pub struct Value(i32);
pub struct Kvpair(i32, i32);

/// An abstraction of storage, we don't care where the data is stored, but need to define how the outside world interacts with the storage
pub trait Storage: Send + Sync + 'static {
    fn get(&self, table: &str, key: &str) -> Result<Option<Value>, KvError>;
    fn set(&self, table: &str, key: String, value: Value) -> Result<Option<Value>, KvError>;
    fn contains(&self, table: &str, key: &str) -> Result<bool, KvError>;
    fn del(&self, table: &str, key: &str) -> Result<Option<Value>, KvError>;
    fn get_all(&self, table: &str) -> Result<Vec<Kvpair>, KvError>;
    fn get_iter(&self, table: &str) -> Result<Box<dyn Iterator<Item = Kvpair>>, KvError>;
}

// Using trait object, no need for generic parameters nor ServiceInner
pub struct Service {
    pub store: Arc<dyn Storage>,
}

// Implementing the trait becomes slightly simpler
impl Service {
    pub fn new<S: Storage>(store: S) -> Self {
        Self {
            store: Arc::new(store),
        }
    }
}

// Implementing the trait also doesn't need to carry generic parameters
impl Clone for Service {
    fn clone(&self) -> Self {
        Self {
            store: Arc::clone(&self.store),
        }
    }
}

From this code, you can see that by sacrificing a bit of performance, we made the overall code writing and usage much more convenient.

Summary #

Whether it’s generics from the previous lecture or today’s trait objects, both are Rust’s means of handling polymorphism. When the system needs to use polymorphism to solve complex and varying requirements, allowing the same interface to show different behaviors, we need to decide whether static dispatch at compile time is better, or dynamic dispatch at runtime is better.

In general, as Rust developers, we don’t mind the slightly more