Refactoring in Rust: Introducing Traits

In the same codebase as last time, we extract data from a HashMap<String, String> called headers, presumably dealing with something similar to HTTP headers.

if headers.contains_key("x-user-id") {
    headers
        .get("x-user-id")
        .map(ToOwned::to_owned)
        .unwrap()
}

Again, for a few lines of code, there is a lot going on here. We check if a certain property is available, and if that’s the case we take an owned value from the hash map. get returns an Option<&String>, but since we did a contains_key check first, we can be sure that the unwrap() call won’t panic.

This procedure occurs a few times in our codebase. Sometimes together with some sort of error handling, which indicates that that "x-user-id" is a required header.

let user_id = if headers.contains_key("x-user-id") {
    headers
        .get("x-user-id")
        .map(ToOwned::to_owned)
        .unwrap()
} else {
    return Err(HeaderError {});
}

HeaderError is a custom error, we like to hold on to that for later.

#[derive(Debug)]
struct HeaderError;

impl Display for HeaderError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Header not found")
    }
}

impl std::error::Error for HeaderError {}

Sometimes we don’t unwrap the output from headers.get. This tells us that this header is optional.

if headers.contains_key("x-authentication") {
    headers
        .get("x-authentication")
        .map(ToOwned::to_owned)
}

It works, which is good, but there are some hidden problems that might bite us. The code itself doesn’t tell us anything about what we expect to happen. Knowing that we deal with optional or required headers is something that we mostly know because, well, I told you so a few paragraphs earlier.

Also, the unwrap() call is something that sticks out like a sore thumb. Yes, it should be safe and not panic because we did a conditional check on the existence of a key earlier, but since it’s not clearly defined what this code is supposed to do it’s just too easy for somebody to get rid of this condition at a later stage.

And last, we check for the existence of a key twice. Once when we explicitly ask if the headers struct contains a certain key, and then again through the return type of headers.get. headers.get returns an Option<&String>. This type tells us everything there is to know: This call might yield no value.

Most of the problems can be addressed by introducing context and semantics. The easiest way to do so is by introducing a new type. Let’s call this type Headers, and let it wrap HashMap<String, String>.

struct Headers(HashMap<String, String>);

impl Headers {
    pub fn new(headers: HashMap<String, String>) -> Self {
        Self(headers)
    }

    pub fn get_required_header(&self, name: &str)
        -> Result<String, HeaderError> {
        self.0
            .get(name)
            .map(ToOwned::to_owned)
            .ok_or(HeaderError {})
    }

    pub fn get_optional_header(&self, name: &str) -> Option<String> {
        self.0.get(name).map(ToOwned::to_owned)
    }
}

We introduce two methods:

• One is called get_required_header. It errors when a header is not available, which we can see immediately from the return type.
• The other is called get_optional_header. It returns the same as headers.get, but an owned value instead of a reference.

In both cases, we only need to look at the method signature to understand what’s going on. The method name gives us context, the return type tells us what to expect. We also got rid of any conditionals that might surround the actual task and are able to properly use error propagation, unwrap_or_default, or similar. The characteristics of the data structure move to the back, and our goals of what we want to achieve with them come forward.

There is just one thing that is odd about this implementation. The functionality we just described, the semantics of our program, are more or fewer convenience methods for an already existing data structure. Shouldn’t we wish that our original data structure would have those characteristics, or dare I say it … traits?

Instead of a new type, we introduce a trait called Headers that describes the methods we just implemented earlier. It tells us everything we need to know. This is what we care about.

trait Headers {
    fn get_required_header(&self, name: &str) -> Result<String, HeaderError>;
    fn get_optional_header(&self, name: &str) -> Option<String>;
}

For our function, we only want to get something that is impl Headers and implements the right methods. The actual data structure is now completely irrelevant to the semantics of our software. If those methods are associated with a HashMap is an implementation detail for some other part of our program.

Speaking of the implementation, for hashmaps, it’s almost the same as it was before.

impl Headers for HashMap<String, String> {
    fn get_required_header(&self, name: &str) -> Result<String, HeaderError> {
        self.get(name).map(ToOwned::to_owned).ok_or(HeaderError {})
    }

    fn get_optional_header(&self, name: &str) -> Option<String> {
        self.get(name).map(ToOwned::to_owned)
    }
}

Using implemented Headers is a delight. When we look at the following lines of code we see which parts are required, and which ones are optional, and we know how to handle errors and stay on the happy path.

fn do_something(headers: &impl Headers) -> Result<(), HeaderError> {
    let user_id = headers.get_required_header("x-user-id")?;
    let token = headers.get_required_header("x-token")?;
    let tenant_id = headers.get_optional_header("x-tenant");

    // ...
    Ok(())
}

It’s also a lot less for us to parse and understand. We don’t hide semantics by looking at basic data types and control flow, the code tells us what’s up. Nothing to keep in our mind maps, everything is laid out and spelled out clearly. Every intention is visible.

We also stay flexible with implementations. Imagine that headers will become optional at some point in time, but you still need to get header information at some other part of your code. Implement the same trait for Option<HashMap<String, String>> and don’t change a single thing of your actual logic.

impl Headers for Option<HashMap<String, String>> {
    fn get_required_header(&self, name: &str) -> Result<String, HeaderError> {
        match self {
            Some(headers_map) => headers_map.get_required_header(name),
            None => Err(HeaderError {}),
        }
    }

    fn get_optional_header(&self, name: &str) -> Option<String> {
        self.as_ref()
            .and_then(|headers_map| headers_map.get_optional_header(name))
    }
}

You might argue that a new type might be more fitting since not every HashMap<String, String> might contain header information. And that’s true! If you expose your traits to other parts of your software or have other hash maps with the same format around, names might be confusing.

Since refactoring is a lot about naming things, you might want to opt for something more generic, like a StringMapExt that does the required/optional handling for you.

trait StringMapExt {
    fn get_required_value(&self, name: &str) -> Result<String, Unavailable>;
    fn get_optional_value(&self, name: &str) -> Option<String>;
}

impl StringMapExt for HashMap<String, String> {
    fn get_required_value(&self, name: &str) -> Result<String, Unavailable> {
        self.get(name).map(ToOwned::to_owned).ok_or(Unavailable {})
    }

    fn get_optional_value(&self, name: &str) -> Option<String> {
        self.get(name).map(ToOwned::to_owned)
    }
}

Introducing traits when refactoring is a fantastic way to state your intentions. They make your code more readable and maintainable, and won’t keep you busy parsing built-in types.