Specify in trait method that return type is a generic Self with different concrete type?

Question

I'm trying to implement a typestate system that is generic over its implementation/implementer. For a minimal reproduction, consider 2 states under State:

trait State {}
struct One {}
struct Two {}
impl State for One {}
impl State for Two {}

Now, I want to create 2 traits for transitioning between One and Two. The trait should enforce that its concrete type must be generic over State. When transitioning to Two, State must be One, and vice versa. Thus I came up with the following:

// trait that marks something as typestate-able over State
pub trait StateHolder<S: State> {}

// to transition to State = Two
trait ToTwo where Self: StateHolder<One> {
    fn two(self) -> impl StateHolder<Two>;
} 

// to transition to State = One
trait ToOne where Self: StateHolder<Two> {
    fn one(self) -> impl StateHolder<One>;
} 

Then someone could implement this like so:

// our concrete struct
struct Bar<S: State>(S);

// Bar is now a StateHolder for all S
impl<S: State> StateHolder<S> for Bar<S> {}

// Bar<One> can transition to Two
impl ToOne for Bar<Two> {
    fn one(self) -> Bar<One> {
        todo!()
    }
}

// and Bar<Two> can transition to One
impl ToTwo for Bar<One> {
    fn two(self) -> Bar<Two> {
        todo!()
    }
}

However, my issue with this is the traits' return type -> impl StateHolder<...> is not explicitly Self. Thus, I could also return a different concrete StateHolder:

struct Foo<S: State>(S);

impl<S: State> StateHolder<S> for Foo<S> {}

impl ToOne for Bar<Two> {
    fn one(self) -> Foo<One> { // this compiles but is wrong
        todo!()
    }
}

Is there a way to specify that the return type must be Self, but with a different concrete type for its generic parameter?

Answer 1

You can add a generic associated type to your StateHolder trait, via which it can project to a different state; by constraining the projection of the current state to Self, you enforce that only the expected type will be accepted:

trait StateHolder<S: State> {
    type NewState<T: State>: StateHolder<T, NewState<S> = Self>;
}

impl<S: State> StateHolder<S> for Bar<S> {
    type NewState<T: State> = Bar<T>; // only `Bar<T>` will be accepted here
}

Now you can use this associated type in your state changes:

trait ToTwo: StateHolder<One> {
    fn two(self) -> Self::NewState<Two>;
}

impl ToTwo for Bar<One> {
    fn two(self) -> Bar<Two> { // only `Bar<Two>` will be accepted here
        todo!()
    }
}

See it on the playground.