RPIT and TAIT

While re-thinking the high level interface for the bip324 library, I stumbled upon a cool little type pattern, Return Position Impl Trait (RPIT), which was defined in Rust’s RFC 1522. This was the conservative precursor to the RPITIT (…In Trait) ones which recently landed and are more of a gamechanger for async interfaces. I am curious about the history of the simpler RPIT and how it is a bit at odds with common golang patterns.

First though, what was the scenario where I ran into it. I was refactoring the bip324 I/O handshake which is generic over a reader which implements the standard library’s Read trait. Just focusing on the reader half and ignoring the write half here.

pub fn handshake<R, W>(
    network: Network,
    role: Role,
    garbage: Option<&[u8]>,
    decoys: Option<&[&[u8]]>,
    reader: R,
    writer: W,
) -> Result<(InboundCipher, OutboundCipher), ProtocolError>
where
    R: Read,
    W: Write,
{}

Simplified example, but the idea is that handshake takes ownership of the reader.

The handshake operation performed in the new function is relatively complex and it is made much simpler to wrap the reader with a Chain<Cursor<Vec<u8>, R>>. Now, this big ol’ type is still a Read implementation, Chain is part of the standard library and implements Read, but it hides some of the internal logic to handle over-reads. In this first example of this method’s signature, this doesn’t really matter, all of it is internal. But sometimes the caller wants to “re-expose” the reader instance for whatever reason, so what if it’s returned to them in this wrapped state?

pub fn handshake<R, W>(
    network: Network,
    role: Role,
    garbage: Option<&[u8]>,
    decoys: Option<&[&[u8]]>,
    reader: R,
    writer: W,
) -> Result<(InboundCipher, OutboundCipher, Chain<Cursor<Vec<u8>, R>>, ProtocolError>
where
    R: Read,
    W: Write,
{}

Returning the wrapped reader so no bytes are lost.

This works, but it does seem a little gross to leak implementation details through the return type. The caller doesn’t care about the reader chain and why it is necessary. Also the caller is fully exposed the Chain type’s API which allows them to “split” the chain, which may lead to them dropping some bytes from the stream. Bit of an unnecessary foot-gun.

So my first instinct was to wrap this reader in a new type. I found this to be pretty clean and self documenting.

pub struct SessionReader<R> {
    inner: Chain<Cursor<Vec<u8>>, R>,
}

impl<R: Read> SessionReader<R> {
    fn new(buffer: Vec<u8>, reader: R) -> Self {
        Self {
            inner: std::io::Read::chain(Cursor::new(buffer), reader),
        }
    }
}

impl<R: Read> Read for SessionReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.inner.read(buf)
    }
}

Minimize the API surface of this new type.

This was all well and good until I tried to create an async version of the handshake with its own async version of SessionReader. Implementing AsyncRead is not nearly as straight forward since async interfaces still has warts and can’t always use the simple async/await syntax. So I thought fine, I’ll just toss back the async version of Chain<Cursor<Vec<u8>>, R> directly. However, tokio’s (I am coding with this runtime assumed) Chain equivalent is not exposed, so it can’t be named. So what to do?

Well, technically what I want is to just hide from the caller this gnarly type transformation and have them just work against a AsyncRead trait. And I would like to avoid dynamic dispatch. Turns out, this case is mentioned in RFC 1522 for RPIT as leaky APIs.

pub async fn handshake<R, W>(
    network: Network,
    role: Role,
    garbage: Option<&[u8]>,
    decoys: Option<&[&[u8]]>,
    mut reader: R,
    writer: &mut W,
) -> Result<(InboundCipher, OutboundCipher, impl AsyncRead + Unpin + Send), ProtocolError>
where
    R: AsyncRead + Send + Unpin,
    W: AsyncWrite + Unpin,
{

RPIT form of the async version of handshake.

A constraint of the RPIT feature is that all return values must be the same type, but that is fine here (would need dynamic dispatch to support that). Kinda interesting how the RFC calls these types “unboxed” abstract types. I see “boxed” as an extra wrapper, so that “un” just sounds funny to me, but I guess it is the technical definition.

This looks clean, but it does go against the golang principle “accept interfaces, return structs”. Specially the second half about returning concrete types. While this is rust and not go, I find go’s rules to usually apply well to any language, so curious what complexity I am taking on here. Returning a concrete type in go is about giving the caller the most flexibility. Go uses structural typing, so the caller can now use that returned type anywhere which accepts its interface. A returned database implementation could be used in functions which only work on the read methods or write methods. Both halves of go’s principle are focused on caller flexibility within go’s type system, which is dynamic dispatch backed by default.

Life isn’t so simple in rust land, so “return concrete type” might not be as consistently a good thing. Since rust is nominally typed, a concrete type doesn’t provide the caller with as much flexibility as go. And as seen above, the concrete type might just expose the caller to more complexity than they need.

TAIT

Ok, I’ll admit I got excited with RPIT and tried it out in bip324’s v0.8.0. And it is not ideal. I under-estimated how painful it would be that the RPIT’s returned type cannot be named. This means you can’t use the type in a struct or enum. The anonymous return type doesn’t exist until compile time (an existential type). The compile does know the existential type, but there is no way for you the developer to refer to it yet. That is where type alias impl trait (TAIT) comes in.

The feature is still (after many years) still only on the nightly toolchain. It is enabled with type_alias_impl_trait. It still has some edges to iron out unfortunately. And it is a little weird, like how it needs a defining function so the compiler can infer the type. There is tension between the “hide the implementation” and “we need to name this thing”.

type ProtocolSessionReader<R> = impl AsyncRead + Unpin;

// A defining functoin? Maybe? I haven't tested this, but would be cool.
pub async fn handshake<R, W>(
    network: Network,
    role: Role,
    garbage: Option<&[u8]>,
    decoys: Option<&[&[u8]]>,
    mut reader: R,
    writer: &mut W,
) -> Result<(InboundCipher, OutboundCipher, ProtocolSessionReader<R>), ProtocolError>
where
    R: AsyncRead + Send + Unpin,
    W: AsyncWrite + Unpin,
{
    // ... handshake logic ...
    
    let leftover_bytes = garbage_buffer[garbage_bytes..].to_vec();
    let session_reader = Cursor::new(leftover_bytes).chain(reader);
    
    Ok((inbound_cipher, outbound_cipher, session_reader))
}

In the future, using a TAIT to hide the session reader implementation.

I can’t use the TAIT pattern yet since it is only on nightly. But, knowing that it is in the pipeline makes me more willing to code up something ugly if it matches the future interface. So I’ll take on the async grossness.

/// An async reader that chains unconsumed handshake data with the underlying stream.
///
/// This type is returned from the async handshake process and ensures that any
/// unread bytes from the handshake (such as partial packets in the garbage buffer)
/// are read before data from the underlying stream.
pub struct ProtocolSessionReader<R> {
    // This should be simplified with TAIT (Type Alias Impl Trait) once stable.
    /// Leftover bytes from the handshake that need to be read first.
    leftover: Vec<u8>,
    /// Current position in the leftover buffer.
    leftover_pos: usize,
    /// The underlying async reader.
    reader: R,
}

impl<R> ProtocolSessionReader<R> {
    /// Create a new async session reader from leftover handshake bytes and the underlying reader.
    fn new(leftover: Vec<u8>, reader: R) -> Self {
        Self {
            leftover,
            leftover_pos: 0,
            reader,
        }
    }
}

impl<R: AsyncRead + Unpin> AsyncRead for ProtocolSessionReader<R> {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        // Manually implement AsyncRead chaining instead of using
        // tokio's chain() because chain() returns an unnameable `impl AsyncRead` type.

        // Get a normal mutable reference from the Pin.
        // This is safe because R: Unpin.
        let this = self.get_mut();

        // First, drain any leftover bytes from the handshake
        if this.leftover_pos < this.leftover.len() {
            let remaining = &this.leftover[this.leftover_pos..];
            let to_copy = remaining.len().min(buf.remaining());
            buf.put_slice(&remaining[..to_copy]);
            this.leftover_pos += to_copy;
            return Poll::Ready(Ok(()));
        }

        // Once leftover is drained, delegate to the underlying reader.
        // We need to re-pin the reader for the async read call.
        Pin::new(&mut this.reader).poll_read(cx, buf)
    }
}

Implementing poll_read and having to deal with Pin.

So yea, if you want to maintain zero-cost abstractions, the type transformation will always be visible to the caller. Protocol::new(..., reader: R, writer: W) produces a Protocol<ProtocolSessionReader<R>, W>. The best we can do is make it ergonomic (like with a type alias) but we can’t make it invisible without introducing runtime cost (boxing).