Enum pin-projections

[Nemo157]

I'm not sure if there's any nicer API that pin-utils could provide, but I recently had to manually implement pin-projections into an enum, if there's some way to encapsulate this into a macro that ensures most of the safety it would be great. (Maybe a proc-macro-attribute like discussed in #20 could do this with annotated variants? @taiki-e). Here's an example from futures (I don't really care about the exact API here, just something that allows the same operations to be written without any unsafe):

#[derive(Debug)]
enum ElemState<F>
where
    F: Future,
{
    Pending(F),
    Done(Option<F::Output>),
}

impl<F> ElemState<F>
where
    F: Future,
{
    fn pending_pin_mut<'a>(self: Pin<&'a mut Self>) -> Option<Pin<&'a mut F>> {
        // Safety: Basic enum pin projection, no drop + optionally Unpin based
        // on the type of this variant
        match unsafe { self.get_unchecked_mut() } {
            ElemState::Pending(f) => Some(unsafe { Pin::new_unchecked(f) }),
            ElemState::Done(_) => None,
        }
    }

    fn take_done(self: Pin<&mut Self>) -> Option<F::Output> {
        // Safety: Going from pin to a variant we never pin-project
        match unsafe { self.get_unchecked_mut() } {
            ElemState::Pending(_) => None,
            ElemState::Done(output) => output.take(),
        }
    }
}

[taiki-e]

I think it is possible to implement it.
I have not started this work yet, but I think it will be as follows:

// I think it is possible to integrate this with `#[unsafe_project]`, 
// but I also feel that the explanation will be complicated.
#[unsafe_enum_project]
#[derive(Debug)]
enum ElemState<F>
where
    F: Future,
{
    #[pin]
    Pending(F),
    #[take]
    Done(Option<F::Output>),
}

impl<F> ElemState<F> 
where
    F: Future,
{
    fn foo(mut self: Pin<&mut Self>) {
        // the variant that uses `#[pin]` attribute:
        // method's name (provisional): VariantName.to_lowercase() + "_pin_mut"
        let _: Option<Pin<&mut F>> = self.as_mut().pending_pin_mut();

        // the variant that uses `#[take]` attribute:
        // method's name (provisional): "take_" + VariantName.to_lowercase()
        let _: Option<F::Output> = self.as_mut().take_done();
        
        // the other variants:
        // (assume that there is `Other(T)` variant.)
        // method's name (provisional): VariantName.to_lowercase() + "_mut"
        // let _: Option<&mut T> = self.as_mut().other_mut();

        // the variant with no value (e.g. `Option::None`) is skipped.
    }
}

// This is generated when `#[unsafe_enum_project(Unpin)]` is used.
// impl<F> Unpin for ElemState<T, U> where F: Future, F: Unpin {} // Conditional Unpin impl

Probably I think I can write the implementation of this within this week.

[Nemo157]

#[take] is probably too use-case specific, being able to do self.as_mut().done_mut().unwrap().take() would work fine here.

The other thing to consider is multi-tuple-variants and struct-variants, in both those cases you may want to pin some of the fields but not all.

My first thought on doing this as a proc-macro-attribute was to generate a corresponding enum representing the projected enum, e.g. for this case

enum ElemStateProjection<F>
where
    F: Future,
{
    Pending(Pin<&mut F>),
    Done(&mut Option<F::Output>),
}

unfortunately this would not be as nice as it is for structs as you will have to actually refer to it by name when matching

match state.project() {
    ElemStateProjection::Pending(future) => { ... },
    ElemStateProjection::Done(item) => { ... },
}

That would retain the matchability of the enum though, returning Option from methods works for the usecase linked above because of how the enum is used, but it would get complex to use in more complicated usecases.

[taiki-e]

The other thing to consider is multi-tuple-variants and struct-variants, in both those cases you may want to pin some of the fields but not all.

I think it can support with this by changing to the implementation that uses attributes to fields instead of variants.

#[unsafe_enum_project]
enum Foo<A: Future, B, C, D> {
    Bar(#[pin] A, B),
    Baz { #[pin] field1: C, field2: D },
}

impl<A: Future, B, C, D> Foo<A, B, C, D> {
    fn foo(mut self: Pin<&mut Self>) {
        // multi-tuple-variants
        let _: Option<(Pin<&mut A>, &mut B)> = self.as_mut().bar();

        // struct-variants
        // extract it as a tuple in the order of arrangement of fields.
        let _: Option<(Pin<&mut C>, &mut D)> = self.as_mut().baz();
        // or create a new projection struct
        // let baz = self.as_mut().baz().unwrap();
        // let _: Pin<&mut C> = baz.field1;
        // let _: &mut D = baz.field2;
    }
}

As a way to handle all variants at the same time, I thought of the following two things.

The first is to add methods like Either::either and Either::either_with.

The second is to create a projected enum and used together with proc_macro(proc_macro_hack) or proc_macro_attribute.

#[unsafe_project]
enum Foo<A: Future, B, C> {
    Bar(#[pin] A, B),
    Baz { field: C },
}

impl<A: Future, B, C> Foo<A, B, C> {
    fn foo(mut self: Pin<&mut Self>) {
        project!(match self {
            Foo::Bar(a, b) => {
                let _: Pin<&mut A> = a;
                let _: &mut B = b;
            }
            Foo::Baz { field } => {
                let _: &mut C = field;
            }
        });
        // or `proc_macro_attribute`
        // #[project]
        // match self {
        //     Foo::Bar(a, b) => {
        //         let _: Pin<&mut A> = a;
        //         let _: &mut B = b;
        //     }
        //     Foo::Baz { field } => {
        //         let _: &mut C = field;
        //     }
        // }
    }
}

It is expanded as follows:

// create a new projection enum and the `.project()` method.
enum __FooProjection<A: Future, B, C> {
    Bar(Pin<&mut A>, &mut B),
    Baz { field: &mut C },
}

impl<A: Future, B, C> Foo<A, B, C> {
    fn foo(mut self: Pin<&mut Self>) {
        match self.project() {
            // the name is rewritten by the macro.
            __FooProjection::Bar(a, b) => {
                let _: Pin<&mut A> = a;
                let _: &mut B = b;
            }
            // the name is rewritten by the macro.
            __FooProjection::Baz { field } => {
                let _: &mut C = field;
            }
        }
    }
}

// This is generated when `#[unsafe_project(Unpin)]` is used.
// impl<A: Future, B, C> Unpin for Foo<A, B, C> where A: Unpin {} // Conditional Unpin impl

Probably, it can support match and if let expressions.

[taiki-e]

I wrote one of the ideas for enum projection: details
This is a minimal implementation, such as not supporting struct-variants, but I think this would improve the current situation that we are writing everything manually.

#[unsafe_variants(Unpin)] // `Unpin` is optional (create the appropriate conditional Unpin implementation)
enum Foo<A, B, C> {
    Variant1(#[pin] A, B),
    Variant2(C),
}

impl<A, B, C> Foo<A, B, C> {
    fn bar(mut self: Pin<&mut Self>) {
        let _: Option<(Pin<&mut A>, &mut B)> = self.as_mut().variant1();
        let _: Option<&mut C> = self.as_mut().variant2();
    }
}
// Automatically create the appropriate conditional Unpin implementation (optional).
// impl<A, B, C> Unpin for Foo<A, B, C> where A: Unpin {} // Conditional Unpin impl

As a next step, I'd like to implement an idea that uses attributes to create a projected enum and macros to support pattern matching.

I also believe macros to support pattern matching can handle projection struct more flexible.

#[unsafe_project]
strust Foo<A: Future, B> {
    #[pin]
    future: A,
    field: B,
}

impl<A: Future, B> Foo<A, B> {
    fn bar(mut self: Pin<&mut Self>) {
        #[project] // or `project!`
        let Foo { future, field } = self

        let _: Pin<&mut A> = future;
        let _: &mut B = field;
    }
}