Preserve original lifetime when performing pin-projection

pin-project-internal/src/pin_project/enums.rs

@@ -4,7 +4,7 @@ use syn::{parse::Nothing, Field, Fields, FieldsNamed, FieldsUnnamed, ItemEnum, R
 
 use crate::utils::VecExt;
 
-use super::{Context, PIN};
+use super::{Context, ProjTraitGenerics, PIN};
 
 pub(super) fn parse(cx: &mut Context, mut item: ItemEnum) -> Result<TokenStream> {
     if item.variants.is_empty() {
@@ -34,17 +34,23 @@ pub(super) fn parse(cx: &mut Context, mut item: ItemEnum) -> Result<TokenStream>
     let Context { proj_ident, proj_trait, orig_ident, lifetime, .. } = &cx;
     let proj_generics = cx.proj_generics();
     let proj_ty_generics = proj_generics.split_for_impl().1;
-    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
+
+    let ProjTraitGenerics { impl_generics, ty_generics, where_clause, orig_ty_generics } =
+        cx.proj_trait_generics();
 
     let mut proj_items = quote! {
         #[allow(dead_code)]
         enum #proj_ident #proj_generics #where_clause { #(#proj_variants,)* }
     };
+
+    let crate_path = crate::utils::crate_path();
+
     proj_items.extend(quote! {
-        impl #impl_generics #proj_trait #ty_generics for ::core::pin::Pin<&mut #orig_ident #ty_generics> #where_clause {
-            fn project<#lifetime>(&#lifetime mut self) -> #proj_ident #proj_ty_generics #where_clause {
+        impl #impl_generics #proj_trait #ty_generics for ::core::pin::Pin<&#lifetime mut #orig_ident #orig_ty_generics> #where_clause {
+            fn project(&mut self) -> #proj_ident #proj_ty_generics #where_clause {
+                use #crate_path::ProjectThrough;
                 unsafe {
-                    match self.as_mut().get_unchecked_mut() {
+                    match self.proj_through().get_unchecked_mut() {
                         #(#proj_arms,)*
                     }
                 }

pin-project-internal/src/pin_project/mod.rs

@@ -75,6 +75,15 @@ struct Context {
 
     unsafe_unpin: bool,
     pinned_drop: Option<Span>,
+
+    proj_trait_generics: Generics,
+}
+
+struct ProjTraitGenerics<'a> {
+    impl_generics: ImplGenerics<'a>,
+    ty_generics: TypeGenerics<'a>,
+    where_clause: Option<&'a WhereClause>,
+    orig_ty_generics: TypeGenerics<'a>,
 }
 
 impl Context {
@@ -104,6 +113,9 @@ impl Context {
             );
         }
 
+        let mut proj_trait_generics = generics.clone();
+        crate::utils::proj_generics(&mut proj_trait_generics, lifetime.clone());
+
         Ok(Self {
             orig_ident: orig_ident.clone(),
             proj_ident,
@@ -115,6 +127,7 @@ impl Context {
             pinned_fields: vec![],
             unsafe_unpin: unsafe_unpin.is_some(),
             pinned_drop,
+            proj_trait_generics,
         })
     }
 
@@ -125,6 +138,21 @@ impl Context {
         generics
     }
 
+    fn proj_trait_generics(&self) -> ProjTraitGenerics<'_> {
+        //let mut orig_generics_for_impl = orig.clone();
+        //crate::utils::proj_generics(&mut orig_generics_for_impl, self.lifetime.clone());
+        //let (impl_generics, modified_ty_generics, _) = orig_generics_for_impl.split_for_impl();
+        let (impl_generics, modified_ty_generics, _) = self.proj_trait_generics.split_for_impl();
+
+        let (_, ty_generics, where_clause) = self.generics.split_for_impl();
+        ProjTraitGenerics {
+            impl_generics,
+            ty_generics: modified_ty_generics,
+            where_clause,
+            orig_ty_generics: ty_generics,
+        }
+    }
+
     fn push_unpin_bounds(&mut self, ty: Type) {
         self.pinned_fields.push(ty);
     }
@@ -343,11 +371,15 @@ impl Context {
         let proj_generics = self.proj_generics();
         let proj_ty_generics = proj_generics.split_for_impl().1;
 
-        let (orig_generics, _, orig_where_clause) = self.generics.split_for_impl();
+        let mut trait_generics = self.generics.clone();
+        crate::utils::proj_generics(&mut trait_generics, lifetime.clone());
+        //trait_generics.params.push(GenericParam::Lifetime(LifetimeDef::new(lifetime.clone())));
+
+        let (orig_generics, _, orig_where_clause) = trait_generics.split_for_impl();
 
         quote! {
             trait #proj_trait #orig_generics {
-                fn project<#lifetime>(&#lifetime mut self) -> #proj_ident #proj_ty_generics #orig_where_clause;
+                fn project(&mut self) -> #proj_ident #proj_ty_generics #orig_where_clause;
             }
         }
     }

pin-project-internal/src/pin_project/structs.rs

@@ -4,7 +4,7 @@ use syn::{parse::Nothing, Field, Fields, FieldsNamed, FieldsUnnamed, Index, Item
 
 use crate::utils::VecExt;
 
-use super::{Context, PIN};
+use super::{Context, ProjTraitGenerics, PIN};
 
 pub(super) fn parse(cx: &mut Context, mut item: ItemStruct) -> Result<TokenStream> {
     let (proj_fields, proj_init) = match &mut item.fields {
@@ -31,17 +31,29 @@ pub(super) fn parse(cx: &mut Context, mut item: ItemStruct) -> Result<TokenStrea
     let Context { proj_ident, proj_trait, orig_ident, lifetime, .. } = &cx;
     let proj_generics = cx.proj_generics();
     let proj_ty_generics = proj_generics.split_for_impl().1;
-    let (impl_generics, ty_generics, where_clause) = item.generics.split_for_impl();
+
+    /*let mut orig_generics_for_impl = item.generics.clone();
+    crate::utils::proj_generics(&mut orig_generics_for_impl, lifetime.clone());
+    let (impl_generics, modified_ty_generics, _) = orig_generics_for_impl.split_for_impl();
+
+    let (_, ty_generics, where_clause) = item.generics.split_for_impl();*/
+
+    let ProjTraitGenerics { impl_generics, ty_generics, where_clause, orig_ty_generics } =
+        cx.proj_trait_generics();
 
     let mut proj_items = quote! {
         #[allow(dead_code)]
         struct #proj_ident #proj_generics #where_clause #proj_fields
     };
+
+    let crate_path = crate::utils::crate_path();
+
     proj_items.extend(quote! {
-        impl #impl_generics #proj_trait #ty_generics for ::core::pin::Pin<&mut #orig_ident #ty_generics> #where_clause {
-            fn project<#lifetime>(&#lifetime mut self) -> #proj_ident #proj_ty_generics #where_clause {
+        impl #impl_generics #proj_trait #ty_generics for ::core::pin::Pin<&#lifetime mut #orig_ident #orig_ty_generics> #where_clause {
+            fn project(&mut self) -> #proj_ident #proj_ty_generics #where_clause {
                 unsafe {
-                    let this = self.as_mut().get_unchecked_mut();
+                    use #crate_path::ProjectThrough;
+                    let this = self.proj_through().get_unchecked_mut();
                     #proj_ident #proj_init
                 }
             }

src/lib.rs

@@ -161,6 +161,78 @@ pub use pin_project_internal::project;
 #[allow(unsafe_code)]
 pub unsafe trait UnsafeUnpin {}
 
+use core::pin::Pin;
+
+/// A helper trait to allow projecting through a mutable reference,
+/// while preserving lifetimes.
+///
+/// Normally, `Pin::as_mut` can be used to convert a '&mut Pin' to a owned 'Pin'
+/// However, since `Pin::as_mut'`needs to work with generic `DerefMut` impls,
+/// it loses lifetime information. Specifically, it returns a `Pin<&'a mut P::Target>`,
+/// where 'a is tied to the lifetime of the original '&mut Pin'. If you started
+/// with a `&mut Pin<'pin &mut T>`, the 'pin lifetime is lost.
+///
+/// This can cause issues when trying to return the result of a pin projection
+/// from a method - e.g. `fn foo<'a>(mut self: Pin<&'a mut Self>) -> Pin<&'a mut T>`
+///
+/// The `ProjectThrough` trait was introduced to solve this issue.
+/// Normally, you'll never need to interact with this type
+/// directly - it's used internally by pin-project to allow expressing
+/// the proper lifetime.
+pub trait ProjectThrough<'a> {
+    type Target;
+    fn proj_through(&mut self) -> Pin<&'a mut Self::Target>;
+}
+
+impl<'a, T> ProjectThrough<'a> for Pin<&'a mut T> {
+    type Target = T;
+    fn proj_through(&mut self) -> Pin<&'a mut Self::Target> {
+        // This method is fairly tricky. We start out with
+        // a `&'p mut Pin<&'a mut T>`, which we want
+        // to convert into a `Pin<&'a mut T>`
+        //
+        // Effectively, we want to discard the outer '&mut'
+        // reference, and just work with the inner 'Pin'
+        // We accomplish this in four steps:
+        //
+        // First, call Pin::as_mut(self). This
+        // converts our '&'p mut Pin<&'a mut T>` to
+        // a `Pin<&mut T>`. This is *almost* what we want -
+        // however, the lifetime is wrong. Since 'as_mut'
+        // uses the `DerefMut` impl of `&mut T`, it loses
+        // the lifetime associated with the `&mut T`.
+        // However, we know that we're dealing with a
+        // `Pin<&mut T>`, so we can soundly recover the original
+        // lifetime.
+        //
+        // This relies on the following property:
+        // The lifetime of the reference returned by <&'a mut T as DerefMut>::deref_mut
+        // is the same as 'a. In order words, calling deref_mut on a mutable reference
+        // is a no-op.
+        let as_mut: Pin<&mut T> = Pin::as_mut(self);
+
+        // Next, we unwrap the Pin<&mut T> by calling 'Pin::get_unchecked_mut'.
+        // This gives us access to the underlying mutable reference
+        #[allow(unsafe_code)]
+        let raw_mut: &mut T = unsafe { Pin::get_unchecked_mut(as_mut) };
+
+        // NExt, we transmute the mutable reference, changing its lifetime to 'a.
+        // Here, we rely on the behavior of DerefMut for mutable references
+        // described above.
+        //
+        // This is the core of this method - everything else is just to deconstruct
+        // and reconstrut a Pin
+        #[allow(unsafe_code)]
+        let raw_transmuted: &'a mut T = unsafe { core::mem::transmute(raw_mut) };
+
+        // Finally, we construct a Pin from our 'upgraded' mutable reference
+        #[allow(unsafe_code)]
+        unsafe {
+            Pin::new_unchecked(raw_transmuted)
+        }
+    }
+}
+
 #[doc(hidden)]
 pub mod __private {
     use super::UnsafeUnpin;

tests/pin_project.rs

@@ -262,3 +262,19 @@ fn test_private_type_in_public_type() {
         OtherVariant(u8),
     }
 }
+
+#[test]
+fn test_lifetime_project() {
+    #[pin_project::pin_project]
+    struct Struct<T, U> {
+        #[pin]
+        pinned: T,
+        unpinned: U,
+    }
+
+    impl<T, U> Struct<T, U> {
+        fn get_pin_mut<'a>(mut self: Pin<&'a mut Self>) -> Pin<&'a mut T> {
+            self.project().pinned
+        }
+    }
+}