ruk·si

🦀 Rust
Typestate

Updated at 2024-02-24 06:26

Typestate makes incorrect usage of a stateful entity a compile-time error.

Typestate is a Rust-specific design pattern that encodes information about object's runtime state in its compile-time type.

The following examples use an HTTP response with three states:

  1. an empty response that expects a status line
  2. now it expects zero or more headers and a body text
  3. now it's complete and sent or ready for sending

And the state transitions must happen in that order.

Naive Typestate

Naive typestate hard-separates each state apart from shared inner details.

You can use Box<T> smart pointer for the inner details to avoid potentially expensive copying on move, but it's not essential with simple internal data like this. It's usually recommended to start with no smart pointer and only add them if profiling shows the moving being an issue.

Here we use boxes just to show how it would work.

struct ResponseInner {
    pub code: Option<u8>,
    pub headers: Vec<[String; 2]>,
    pub body: Option<String>,
}

struct Response(Box<ResponseInner>);
struct ResponseAfterStatus(Box<ResponseInner>);
struct ResponseAfterBody(Box<ResponseInner>);

// you could make the API nicer with Deref here if wanted,
// it should help getting rid of the `.0` here and there

impl Response {
    fn new() -> Self {
        Response(
            Box::new(ResponseInner {
                code: None,
                headers: vec![],
                body: None,
            })
        )
    }
    fn status(mut self, code: u8) -> ResponseAfterStatus {
        self.0.code = Some(code);
        ResponseAfterStatus(self.0)
    }
}

impl ResponseAfterStatus {
    fn header(mut self, key: &str, value: &str) -> Self {
        self.0.headers.push([key.to_string(), value.to_string()]);
        ResponseAfterStatus(self.0)
    }
    fn body(mut self, text: &str) -> ResponseAfterBody {
        self.0.body = Some(text.to_string());
        ResponseAfterBody(self.0)
    }
}

impl ResponseAfterBody {
    fn send(self) {
        // TODO: send it 📩
    }
}

fn main() {
    let res = Response::new()                  // Response
        .status(200)                           // ResponseAfterStatus
        .header("Content-Length", "6")         // ResponseAfterStatus
        .header("Content-Type", "text/plain")  // ResponseAfterStatus
        .body("Hello World!");                 // ResponseAfterBody

    assert_eq!(res.0.body, Some("Hello World!".into()));

    res.send();
    // sending consumes the response; ownership is gone now 😶‍🌫️
    // trying to access `res` here would not compile
}

Generic Typestate with Shared Details

Generic typestate allows having functions on all or a limited set of states.

struct Response<S: ResponseState> {
    inner: Box<ResponseInner>,
    marker: std::marker::PhantomData<S>,
}

// it's nicer to use empty enums instead of empty struct
// as you cannot create instances of empty enums but both 
// would work
enum Start {}   // expecting status line
enum Headers {} // expecting headers or body

trait ResponseState {}
impl ResponseState for Start {}
impl ResponseState for Headers {}

// _all_ states allow these
impl<S> Response<S> {
    fn bytes_so_far(&self) -> usize {}
}

// `Start` state allows these
impl Response<Start> {
    fn new() -> Self {}
    fn status(self, code: u8, message: &str) -> Response<Headers> {}
}

// `Headers` state allows these
impl Response<Headers> {
    fn header(&mut self, key: &str, value: &str) {}
    fn body(self, contents: &str) {}
}

/// _some_ states allow these
trait SendingState {}
impl SendingState for Headers {}
impl<S: SendingState> Response<S> {
    fn spam_spam_spam(&mut self);
}

Generic Typestate with Individual Details

If states have distinctly different inner details with less sharing, you can use this.

struct Response<S: ResponseState> {
    inner: Box<ResponseInner>,
    extra: S,
}

// now each state can have it's own data

struct Start;
struct Headers {
    code: u8,
}

trait ResponseState {}
impl ResponseState for Start {}
impl ResponseState for Headers {}

impl Response<Start> {
    fn status(self, code: u8, message: &str) -> Response<Headers> {
        Response {
            state: self.state,
            extra: Headers {
                code,
            },
        }
    }
}

impl Response<Headers> {
    fn code(&self) -> u8 {
        self.extra.code
    }
}

Sources