[PATCH v2 07/20] rust: pin-init: internal: pin_data: check covariance of self-referential fields

[Gary Guo <gary@garyguo.net>]

We implicitly infer covariance for fields that self-references. This needs
to be checked to ensure that the fields are really covariant, so the rest
of expansion code can rely on this fact.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/pin-init/internal/src/pin_data.rs | 75 +++++++++++++++++++++++++++++++++-
 rust/pin-init/internal/src/util.rs     | 24 ++++++++++-
 2 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/rust/pin-init/internal/src/pin_data.rs b/rust/pin-init/internal/src/pin_data.rs
index ec0b1aeefe7f..6a29ec358c30 100644
--- a/rust/pin-init/internal/src/pin_data.rs
+++ b/rust/pin-init/internal/src/pin_data.rs
@@ -76,7 +76,6 @@ enum Variance {
 }
 
 /// Information about field lifetimes captured in a type.
-#[expect(unused)]
 struct Capture {
     variance: Variance,
     /// Lifetime to be captured.
@@ -435,6 +434,7 @@ fn expand(
     let unpin_impl = generate_unpin_impl(&info);
     let drop_impl = generate_drop_impl(&info);
     let drop_order_check = generate_drop_order_check(dcx, &info);
+    let variance_check = generate_variance_check(&info);
     let projections = generate_projections(&info);
     let the_pin_data = generate_the_pin_data(&info);
 
@@ -444,6 +444,7 @@ fn expand(
         // outside.
         const _: () = {
             #drop_order_check
+            #variance_check
             #projections
             #the_pin_data
             #unpin_impl
@@ -770,6 +771,78 @@ fn __drop_order_check #impl_generics_with_field_lt (
     }
 }
 
+/// Produce variance checks, so we can ensure that the variance of lifetimes captured by field types
+/// actually match our expectation.
+fn generate_variance_check(info: &StructInfo) -> TokenStream {
+    if !info.self_referential {
+        return quote!();
+    }
+
+    let mut checks = Vec::new();
+
+    for f in info.fields.iter() {
+        let covariant_captures: Vec<_> = f
+            .captures
+            .iter()
+            .filter(|b| b.variance == Variance::Covariant)
+            .map(|b| &b.lifetime)
+            .collect();
+        if covariant_captures.is_empty() {
+            continue;
+        }
+
+        let ident = f.member.as_ident();
+        // Use the span of type for better error message.
+        let span = f.field.ty.span().resolved_at(Span::mixed_site());
+
+        let other_field_lifetimes = Generics {
+            lt_token: None,
+            params: f
+                .captures
+                .iter()
+                .filter(|b| b.variance != Variance::Covariant)
+                .map(|b| GenericParam::Lifetime(LifetimeParam::new(b.lifetime.clone())))
+                .collect(),
+            gt_token: None,
+            where_clause: None,
+        };
+
+        let long = Lifetime::new("'__long", span);
+        let long_ty = f
+            .field
+            .ty
+            .replace_lifetimes(&covariant_captures, &vec![&long; covariant_captures.len()]);
+
+        let short = Lifetime::new("'__short", span);
+        let short_ty = f
+            .field
+            .ty
+            .replace_lifetimes(&covariant_captures, &vec![&short; covariant_captures.len()]);
+
+        let check_name = format_ident!("__{ident}_covariance", span = span);
+
+        // Add `<'__long: '__short, 'short>` as additional generics.
+        let covariance_check_generics = parse_quote!(<#long: #short, #short>);
+        let combined_generics = CombinedGenerics(vec![
+            &covariance_check_generics,
+            &other_field_lifetimes,
+            &info.struct_.generics,
+        ]);
+        let (combined_impl_generics, _, whr) = combined_generics.split_for_impl();
+
+        checks.push(quote_spanned!(span =>
+            // Emit a check to ensure the type is *really* covariant for soundness.
+            fn #check_name #combined_impl_generics (long: #long_ty) -> #short_ty #whr {
+                long
+            }
+        ));
+    }
+
+    quote!(
+        #(#checks)*
+    )
+}
+
 fn generate_projections(info: &StructInfo) -> TokenStream {
     let ItemStruct {
         vis,
diff --git a/rust/pin-init/internal/src/util.rs b/rust/pin-init/internal/src/util.rs
index f19712a46a30..59054f5934fb 100644
--- a/rust/pin-init/internal/src/util.rs
+++ b/rust/pin-init/internal/src/util.rs
@@ -5,8 +5,8 @@
 use proc_macro2::{Ident, TokenStream};
 use quote::{format_ident, ToTokens};
 use syn::{
-    visit::Visit, Attribute, BoundLifetimes, GenericParam, Generics, Index, Lifetime, Member,
-    Token, TypePath,
+    parse_quote, visit::Visit, Attribute, BoundLifetimes, GenericParam, Generics, Index, Lifetime,
+    Member, Token, Type, TypePath,
 };
 
 use crate::DiagCtxt;
@@ -386,3 +386,23 @@ fn visit_type_path(&mut self, ty: &'a TypePath) {
         TypeParamVisitor(f)
     }
 }
+
+pub(crate) trait TypeExt {
+    fn replace_lifetimes(&self, needle: &[&Lifetime], replacement: &[&Lifetime]) -> Type;
+}
+
+impl TypeExt for Type {
+    fn replace_lifetimes(&self, needle: &[&Lifetime], replacement: &[&Lifetime]) -> Type {
+        if needle.is_empty() {
+            return self.clone();
+        }
+
+        parse_quote!(
+            <
+                for<#(#needle,)*> fn(#(&#needle (),)*) -> #self
+                    as
+                ::pin_init::__internal::FnOutput<(#(&#replacement (),)*)>
+            >::Output
+        )
+    }
+}

-- 
2.54.0