[PATCH v2 02/20] rust: pin-init: internal: pin_data: infer self-referential struct

[Gary Guo <gary@garyguo.net>]

As a first step towards adding self-referential data structures in
pin-init, add parsing support.

Scan all field types for unbounded lifetimes, and if the names that of
fields, it is inferred as a self-referential field lifetime. No explicit
annotations are supported yet.

Acked-by: Benno Lossin <lossin@kernel.org>
Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/pin-init/internal/src/pin_data.rs | 162 ++++++++++++++++++++++++++++++++-
 rust/pin-init/internal/src/util.rs     |  94 ++++++++++++++++++-
 2 files changed, 253 insertions(+), 3 deletions(-)

diff --git a/rust/pin-init/internal/src/pin_data.rs b/rust/pin-init/internal/src/pin_data.rs
index 03e893cf5475..cf6142cd656d 100644
--- a/rust/pin-init/internal/src/pin_data.rs
+++ b/rust/pin-init/internal/src/pin_data.rs
@@ -1,5 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
+use std::collections::{BTreeMap, BTreeSet};
+
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote, ToTokens};
 use syn::{
@@ -7,8 +9,10 @@
     parse_quote, parse_quote_spanned,
     punctuated::Punctuated,
     spanned::Spanned,
+    visit::Visit,
     visit_mut::VisitMut,
-    Field, Fields, Generics, Index, Item, ItemStruct, Member, PathSegment, Type, TypePath,
+    Field, Fields, Generics, Ident, Index, Item, ItemStruct, Lifetime, Member, PathSegment, Type,
+    TypePath,
 };
 
 use crate::{
@@ -48,10 +52,69 @@ fn to_tokens(&self, tokens: &mut TokenStream) {
     }
 }
 
+/// Description of how a field is borrowed.
+#[derive(Clone, Copy, Default, PartialEq, Eq)]
+enum BorrowedKind {
+    /// Implicitly inferreed.
+    #[default]
+    Shared,
+}
+
+/// Information about a borrowed field.
+#[expect(unused)]
+struct BorrowedInfo {
+    kind: BorrowedKind,
+    /// Field lifetime for this field.
+    lifetime: Lifetime,
+}
+
+#[derive(Clone, Copy, Default, PartialEq, Eq)]
+enum Variance {
+    /// Implicitly inferred variance.
+    #[default]
+    Covariant,
+}
+
+/// Information about field lifetimes captured in a type.
+#[expect(unused)]
+struct Capture {
+    variance: Variance,
+    /// Lifetime to be captured.
+    lifetime: Lifetime,
+}
+
+impl std::borrow::Borrow<Lifetime> for Capture {
+    fn borrow(&self) -> &Lifetime {
+        &self.lifetime
+    }
+}
+
+impl PartialEq for Capture {
+    fn eq(&self, other: &Self) -> bool {
+        self.lifetime == other.lifetime
+    }
+}
+
+impl Eq for Capture {}
+
+impl PartialOrd for Capture {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for Capture {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.lifetime.cmp(&other.lifetime)
+    }
+}
+
 struct FieldInfo {
     field: Field,
     member: Member,
     pinned: bool,
+    borrowed: Option<BorrowedInfo>,
+    captures: BTreeSet<Capture>,
 }
 
 struct StructInfo {
@@ -59,6 +122,7 @@ struct StructInfo {
     struct_: ItemStruct,
     fields: Vec<FieldInfo>,
     is_tuple_struct: bool,
+    self_referential: bool,
 }
 
 pub(crate) fn expand_with_cfg(
@@ -147,7 +211,22 @@ fn expand(
     replacer.visit_fields_mut(&mut struct_.fields);
 
     let is_tuple_struct = matches!(struct_.fields, Fields::Unnamed(_));
-    let fields: Vec<FieldInfo> = struct_
+
+    // Collect all bound lifetimes from generics.
+    let bound_lifetimes: BTreeSet<&Lifetime> =
+        struct_.generics.lifetimes().map(|x| &x.lifetime).collect();
+    // Collect all fields.
+    let field_idx_map: BTreeMap<Ident, usize> = struct_
+        .fields
+        .iter()
+        .enumerate()
+        .filter_map(|(index, field)| Some((field.ident.clone()?, index)))
+        .collect();
+
+    // Keep track on fields being implicitly borrowed by being mentioned.
+    let mut implicitly_borrowed = BTreeSet::new();
+
+    let mut fields: Vec<FieldInfo> = struct_
         .fields
         .into_iter()
         .enumerate()
@@ -166,16 +245,88 @@ fn expand(
                 }),
             };
 
+            let mut captures = BTreeSet::new();
+            let wildcard_variance = Variance::default();
+
+            // Infer lifetime based on the field referenced.
+            // Bound lifetimes from struct generics take priority.
+            //
+            // For example,
+            // ```
+            // struct Foo<'a> {
+            //     bar: &'a (),
+            //     a: u32,
+            // }
+            // ```
+            // would not be inferred as self-referential because `'a` is already bound by the
+            // struct generics.
+            Lifetime::visitor(|lt| {
+                if bound_lifetimes.contains(lt) || captures.contains(lt) {
+                    return;
+                }
+
+                if !field_idx_map.contains_key(&lt.ident) {
+                    dcx.error(
+                        lt,
+                        format!("`{lt}` is neither a lifetime in generics nor a field name"),
+                    );
+                    return;
+                }
+
+                captures.insert(Capture {
+                    variance: wildcard_variance,
+                    lifetime: lt.clone(),
+                });
+            })
+            .visit_type(&field.ty);
+
+            for capture in captures.iter() {
+                implicitly_borrowed.insert(capture.lifetime.ident.clone());
+            }
+
             FieldInfo {
                 field,
                 member,
                 pinned,
+                borrowed: None,
+                captures,
             }
         })
         .collect();
 
+    for field_name in implicitly_borrowed.into_iter() {
+        let field = &mut fields[field_idx_map[&field_name]];
+
+        // If field is not explicit marked as borrowed, infer a shared borrow.
+        if field.borrowed.is_none() {
+            field.borrowed = Some(BorrowedInfo {
+                kind: BorrowedKind::Shared,
+                // Obtaining from `field` instead of `field_name` for the correct span.
+                lifetime: Lifetime::from_ident(&field.member.as_ident()),
+            });
+        }
+    }
+
+    // Check that field lifetimes do not appear in the bounds.
+    Lifetime::visitor(|lt| {
+        if bound_lifetimes.contains(&lt) {
+            return;
+        }
+
+        if field_idx_map.contains_key(&lt.ident) {
+            // Forbid the use of field lifetimes within bounds.
+            dcx.error(lt, "field lifetimes cannot be used in bounds");
+        }
+
+        // Otherwise this is completely unbound. Let Rust compiler produce that error instead.
+    })
+    .visit_generics(&struct_.generics);
+
     struct_.fields = Fields::Unit;
     let info = StructInfo {
+        self_referential: fields
+            .iter()
+            .any(|f| !f.captures.is_empty() || f.borrowed.is_some()),
         args,
         struct_,
         fields,
@@ -195,6 +346,13 @@ fn expand(
         }
     }
 
+    if info.self_referential {
+        dcx.error(
+            &info.struct_.ident,
+            "self-referential support is not fully implemented",
+        );
+    }
+
     let struct_def = generate_struct_def(&info);
     let unpin_impl = generate_unpin_impl(&info);
     let drop_impl = generate_drop_impl(&info);
diff --git a/rust/pin-init/internal/src/util.rs b/rust/pin-init/internal/src/util.rs
index 3ce498cea754..67ebb333710f 100644
--- a/rust/pin-init/internal/src/util.rs
+++ b/rust/pin-init/internal/src/util.rs
@@ -1,8 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
+use std::collections::BTreeSet;
+
 use proc_macro2::{Ident, TokenStream};
 use quote::{format_ident, ToTokens};
-use syn::{Attribute, GenericParam, Generics, Index, Member, Token};
+use syn::{
+    visit::Visit, Attribute, BoundLifetimes, GenericParam, Generics, Index, Lifetime, Member, Token,
+};
 
 use crate::DiagCtxt;
 
@@ -237,3 +241,91 @@ fn to_tokens(&self, tokens: &mut TokenStream) {
             .to_tokens(tokens);
     }
 }
+
+pub(crate) trait LifetimeExt {
+    /// Get a visitor that call the provided function for all unbound lifetimes.
+    fn visitor<'a>(f: impl FnMut(&'a Lifetime)) -> impl Visit<'a>;
+
+    /// Obtain a lifetime from a identifier.
+    ///
+    /// The created lifetime has the same span.
+    fn from_ident(ident: &Ident) -> Self;
+}
+
+impl LifetimeExt for Lifetime {
+    fn visitor<'a>(f: impl FnMut(&'a Lifetime)) -> impl Visit<'a> {
+        LifetimeVisitor {
+            bound: BTreeSet::new(),
+            visit: f,
+        }
+    }
+
+    fn from_ident(ident: &Ident) -> Self {
+        Lifetime {
+            apostrophe: ident.span(),
+            ident: ident.clone(),
+        }
+    }
+}
+
+struct LifetimeVisitor<'a, F> {
+    bound: BTreeSet<&'a Lifetime>,
+    visit: F,
+}
+
+impl<'a, F> LifetimeVisitor<'a, F> {
+    fn with_bound_lifetimes(
+        &mut self,
+        bound: Option<&'a BoundLifetimes>,
+        f: impl FnOnce(&mut Self),
+    ) {
+        // In case the type includes a lifetime binder, e.g. `dyn for<'a> Foo`,
+        // the lifetimes in the binder are bound and should not be visited.
+
+        let mut to_remove = Vec::new();
+        if let Some(bound) = bound {
+            for lt in &bound.lifetimes {
+                let GenericParam::Lifetime(lt) = lt else {
+                    continue;
+                };
+                if !self.bound.contains(&&lt.lifetime) {
+                    self.bound.insert(&lt.lifetime);
+                    to_remove.push(&lt.lifetime);
+                }
+            }
+        }
+
+        f(self);
+
+        for lt in to_remove {
+            self.bound.remove(lt);
+        }
+    }
+}
+
+impl<'a, F: FnMut(&'a Lifetime)> Visit<'a> for LifetimeVisitor<'a, F> {
+    fn visit_lifetime(&mut self, lt: &'a Lifetime) {
+        if lt.ident == "static" {
+            return;
+        }
+
+        if !self.bound.contains(lt) {
+            (self.visit)(lt);
+        }
+    }
+
+    fn visit_trait_bound(&mut self, bound: &'a syn::TraitBound) {
+        self.with_bound_lifetimes(bound.lifetimes.as_ref(), |this| {
+            this.visit_path(&bound.path)
+        });
+    }
+
+    fn visit_type_bare_fn(&mut self, bare_fn: &'a syn::TypeBareFn) {
+        self.with_bound_lifetimes(bare_fn.lifetimes.as_ref(), |this| {
+            for input in bare_fn.inputs.iter() {
+                this.visit_bare_fn_arg(input);
+            }
+            this.visit_return_type(&bare_fn.output);
+        });
+    }
+}

-- 
2.54.0