[PATCH v2 19/20] rust: pin-init: internal: pin_data: support shared projection

[Gary Guo <gary@garyguo.net>]

For fields that self-references, it is desirable that projection can
happen on shared references too, so lifetime between different fields can
be correlated. Add support for that with `with_project_ref`.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/pin-init/internal/src/pin_data.rs | 134 +++++++++++++++++++++++++++++++++
 1 file changed, 134 insertions(+)

diff --git a/rust/pin-init/internal/src/pin_data.rs b/rust/pin-init/internal/src/pin_data.rs
index 3abd06ca92dc..3d7726e12b3b 100644
--- a/rust/pin-init/internal/src/pin_data.rs
+++ b/rust/pin-init/internal/src/pin_data.rs
@@ -1150,10 +1150,17 @@ fn generate_projections(info: &StructInfo) -> TokenStream {
         &info.field_lts_outlive_chain,
         generics,
     ]);
+    let generics_with_this_field_ref_lt = CombinedGenerics(vec![
+        &this_lt_generics,
+        &info.field_lts_split_variance_outlive_chain,
+        generics,
+    ]);
     let (impl_generics, ty_generics, whr) = generics.split_for_impl();
     let (_, field_lt_ty_generics, _) = field_lts.split_for_impl();
     let (_, ty_generics_with_this_lt, _) = generics_with_this_lt.split_for_impl();
     let (_, ty_generics_with_this_field_lt, _) = generics_with_this_field_lt.split_for_impl();
+    let (_, ty_generics_with_this_field_ref_lt, _) =
+        generics_with_this_field_ref_lt.split_for_impl();
 
     let this = format_ident!("this");
 
@@ -1467,6 +1474,109 @@ fn generate_projections(info: &StructInfo) -> TokenStream {
         }
     }
 
+    let (fields_ref_decl_lt, fields_ref_proj_lt): (Vec<_>, Vec<_>) = info
+        .fields
+        .iter()
+        .map(|f| {
+            let vis = &f.field.vis;
+            let ident = f.member.as_ident();
+            let member = &f.member;
+            let name = (!info.is_tuple_struct).then(|| quote!(#ident:));
+
+            let ty = &f.field.ty;
+
+            let mut accessor = quote!(&#this.#member);
+            if !f.captures.is_empty() || f.borrowed.is_some() {
+                accessor = quote!(
+                    // SAFETY: we have `Erase<..>` which we know is layout compatible with `f.ty`.
+                    // We cannot include explicit type name here as the field lifetimes are nameable
+                    // in this context.
+                    unsafe { ::core::mem::transmute(#accessor) }
+                )
+            }
+
+            // In `with_project`, borrowed fields have their field lifetime available, so use it
+            // instead of `'__this`.
+            let lt = if f.borrowed.is_some() {
+                Lifetime::from_ident(&ident)
+            } else {
+                this_lt.clone()
+            };
+
+            if matches!(
+                f.borrowed,
+                Some(BorrowedInfo {
+                    kind: BorrowedKind::Mutable,
+                    ..
+                })
+            ) {
+                // Mutable borrow must be omitted for aliasing reason.
+                (
+                    quote!(
+                        #vis #name ::pin_init::__internal::NotVisible<&#lt #ty>,
+                    ),
+                    quote!(
+                        #name ::pin_init::__internal::NotVisible::new(),
+                    ),
+                )
+            } else if f.pinned {
+                (
+                    quote!(
+                        #vis #name ::core::pin::Pin<&#lt #ty>,
+                    ),
+                    quote!(
+                        // SAFETY: this field is structurally pinned.
+                        #name unsafe { ::core::pin::Pin::new_unchecked(#accessor) },
+                    ),
+                )
+            } else {
+                (
+                    quote!(
+                        #vis #name &#lt #ty,
+                    ),
+                    quote!(
+                        #name #accessor,
+                    ),
+                )
+            }
+        })
+        .collect();
+
+    let projection_ref_lt = format_ident!("__ProjectionRef");
+    let (projection_ref_lt_def, projection_ref_lt_init) = if info.is_tuple_struct {
+        (
+            quote!(
+                #vis struct #projection_ref_lt #generics_with_this_field_ref_lt(
+                    #(#fields_ref_decl_lt)*
+                    ::core::marker::PhantomData<&'__this #ident #ty_generics>,
+                ) #whr;
+            ),
+            quote!(
+                #projection_ref_lt(
+                    #(#fields_ref_proj_lt)*
+                    ::core::marker::PhantomData,
+                )
+            ),
+        )
+    } else {
+        (
+            quote! {
+                #vis struct #projection_ref_lt #generics_with_this_field_ref_lt
+                    #whr
+                {
+                    #(#fields_ref_decl_lt)*
+                    ___pin_phantom_data: ::core::marker::PhantomData<&'__this #ident #ty_generics>,
+                }
+            },
+            quote! {
+                #projection_ref_lt {
+                    #(#fields_ref_proj_lt)*
+                    ___pin_phantom_data: ::core::marker::PhantomData,
+                }
+            },
+        )
+    };
+
     quote! {
         #[doc = #docs]
         // Allow `non_snake_case` since the same warning will be emitted on
@@ -1482,6 +1592,12 @@ fn generate_projections(info: &StructInfo) -> TokenStream {
         #[doc(hidden)]
         #projection_lt_def
 
+        // Allow `non_snake_case` since the same warning will be emitted on
+        // the struct definition.
+        #[allow(dead_code, non_snake_case)]
+        #[doc(hidden)]
+        #projection_ref_lt_def
+
         impl #impl_generics #ident #ty_generics
             #whr
         {
@@ -1520,6 +1636,24 @@ impl #impl_generics #ident #ty_generics
                 f(#projection_lt_init)
             }
 
+            /// Pin-projects all fields of `Self` from a shared reference with proper lifetime.
+            ///
+            /// These fields are structurally pinned:
+            #(#[doc = #structurally_pinned_fields_docs])*
+            ///
+            /// These fields are **not** structurally pinned:
+            #(#[doc = #not_structurally_pinned_fields_docs])*
+            #[inline]
+            #vis fn with_project_ref<'__this, R>(
+                self: ::core::pin::Pin<&'__this Self>,
+                f: impl for #field_lt_ty_generics ::core::ops::FnOnce(
+                    #projection_ref_lt #ty_generics_with_this_field_ref_lt
+                ) -> R,
+            ) -> R {
+                let #this = ::core::pin::Pin::get_ref(self);
+                f(#projection_ref_lt_init)
+            }
+
             #(#accessors)*
         }
     }

-- 
2.54.0