autoref specialization in rust

Rust has this neat thing called the orphan rules, which prevent implementing a foreign trait on a foreign type.

This means you can:

Implement a foreign trait on a local type.
Implement a local trait on a foreign type.

But you cannot implement a foreign trait on a foreign type.

In practice this exists to ensure coherence. I also recommend reading this article, which explains the reasoning better than I can.

I understand these rules, but I recently ran into a situation where they became a bit frustrating.

I was using the foreign arbitrary trait across a large number of types. Not all of those types implemented Arbitrary, and I wanted Rust’s type system to enforce that some way of generating an arbitrary value existed.

For local types this is easy enough, an impl Arbitrary works great. For foreign types that do not have it? Not so much. E.g.:

fn print_value<'a, T>(u: &mut Unstructured<'a>)
where
    T: arbitrary::Arbitrary<'a> + std::fmt::Debug,
{
    println!("{:?}", T::arbitrary(u));
}

print_value::<std::net::Shutdown>(&mut u);

This fails because std::net::Shutdown does not implement Arbitrary.

Normally the solution would be wrapper:

#[derive(Debug)]
struct WrapperShutdown(std::net::Shutdown);

impl arbitrary::Arbitrary<'_> for WrapperShutdown {
    fn arbitrary<'a>(
        u: &mut Unstructured<'a>,
    ) -> arbitrary::Result<Self> {
        use std::net::Shutdown;

        let chosen =
            u.choose(&[Shutdown::Write, Shutdown::Read, Shutdown::Both])?;

        Ok(WrapperShutdown(*chosen))
    }
}

This works, but I don’t really like it. Why ? The arbitrary crate has a derive feature which makes implementations on types you own much nicer:

#[derive(Arbitrary)]
pub enum Shutdown {
    Read,
    Write,
    Both,
}

If another variant gets added, the derived implementation updates with it.

With std::net::Shutdown we obviously can’t do that, and we also can’t implement Arbitrary ourselves because both the trait and type are foreign.

So I wanted something closer to:

if T implements Arbitrary:
    use Arbitrary
otherwise:
    use my fallback

Rust’s actual specialization feature is unstable, but there is a useful stable trick based on method resolution and autoref.

The basic setup looks like this:

use std::{marker::PhantomData, net::Shutdown};

use arbitrary::{Arbitrary, Unstructured};

trait FallbackArbitrary {
    fn fallback_arbitrary<'a>(
        u: &mut Unstructured<'a>,
    ) -> arbitrary::Result<Self>
    where
        Self: Sized;
}

struct Dispatch<T>(PhantomData<T>);

trait Call<T> {
    fn call<'a>(
        self,
        u: &mut Unstructured<'a>,
    ) -> arbitrary::Result<T>;
}

impl<T> Call<T> for &&Dispatch<T>
where
    T: for<'a> Arbitrary<'a>,
{
    fn call<'a>(
        self,
        u: &mut Unstructured<'a>,
    ) -> arbitrary::Result<T> {
        T::arbitrary(u)
    }
}

impl<T> Call<T> for &Dispatch<T>
where
    T: FallbackArbitrary,
{
    fn call<'a>(
        self,
        u: &mut Unstructured<'a>,
    ) -> arbitrary::Result<T> {
        T::fallback_arbitrary(u)
    }
}

macro_rules! arbitrary_or_fallback {
    ($ty:ty, $u:expr) => {{
        let marker = Dispatch::<$ty>(PhantomData);
        (&&marker).call($u)
    }};
}

Then we can provide a fallback for foreign types:

impl FallbackArbitrary for Shutdown {
    fn fallback_arbitrary<'a>(
        u: &mut Unstructured<'a>,
    ) -> arbitrary::Result<Self> {
        let chosen =
            u.choose(&[Shutdown::Write, Shutdown::Read, Shutdown::Both])?;

        Ok(*chosen)
    }
}

And use both cases through the same macro:

fn main() {
    let bytes = [1, 2, 3, 4, 5, 6];
    let mut u = Unstructured::new(&bytes);

    // Does not have Arbitrary
    let shutdown =
        arbitrary_or_fallback!(Shutdown, &mut u).unwrap();
    // Has Arbitrary
    let number =
        arbitrary_or_fallback!(u16, &mut u).unwrap();

    println!("{shutdown:?}");
    println!("{number:?}");
}

So why does the &&marker work?

The important part is that these implementations are on two different receiver types:

&&Dispatch<T>
&Dispatch<T>

The preferred implementation is:

impl<T> Call<T> for &&Dispatch<T>
where
    T: Arbitrary<'_>,

while the fallback is:

impl<T> Call<T> for &Dispatch<T>
where
    T: FallbackArbitrary,

When we write:

(&&marker).call(u)

Rust first tries to find a valid method for &&Dispatch<T>.

If T implements Arbitrary, that implementation matches directly and gets used.

If it doesn’t, Rust can dereference the receiver:

&&Dispatch<T> -> &Dispatch<T>

and the fallback implementation becomes available instead.

So the extra & gives one implementation higher priority than the other.

There is no runtime “if T implements Arbitrary” happening here, we’re abusing normal method resolution to get roughly the same effect.

The macro is also important.

You might want to turn this into:

fn arbitrary_or_fallback<T>(
    u: &mut Unstructured<'_>,
) -> arbitrary::Result<T> {
    let marker = Dispatch::<T>(PhantomData);
    (&&marker).call(u)
}

but this doesn’t give us the same behaviour.

Generic functions are type-checked using the bounds known about T. Rust doesn’t wait until monomorphization and then redo method resolution based on whatever traits the concrete type happens to implement.

The macro expands at the call-site, so:

arbitrary_or_fallback!(Shutdown, &mut u)

becomes roughly:

{
    let marker = Dispatch::<Shutdown>(PhantomData);
    (&&marker).call(&mut u)
}

At that point Rust knows the concrete type and can resolve the correct implementation.

This is usually referred to as autoref specialization, there’s a nice crate for it here

I wouldn’t recommend building a giant trait system out of increasingly cursed combinations of:

&T
&&T
&&&T
&&&&T

but for a small dispatch point like this, it is surprisingly useful.

The orphan rules still do exactly what they’re supposed to do. I still cannot implement Arbitrary for Shutdown – but this allows me to fanangle the types and use the #[derive(Arbitrary)] where possible, and fallback to a manual listing of every field where I cannot add Arbitrary to said type.