From 9591eafba61c8ff8fcfb0a3358854a7a42e63b2b Mon Sep 17 00:00:00 2001 From: A4-Tacks Date: Mon, 8 Sep 2025 14:09:48 +0800 Subject: [PATCH] Add nested lifetime support for add_lifetime_to_type Changes: - Add nested lifetime support - Add explicit infer lifetime support - Change assist type to `quickfix` Example --- ```rust struct Foo { a: &$0i32, b: &'_ i32, c: (&i32, Bar<'_>), } ``` **Before this PR**: ```rust struct Foo<'a> { a: &'a i32, b: &'_ i32, c: (&i32, Bar<'_>), } ``` **After this PR**: ```rust struct Foo<'a> { a: &'a i32, b: &'a i32, c: (&'a i32, Bar<'a>), } ``` --- .../src/handlers/add_lifetime_to_type.rs | 115 +++++++++++------- 1 file changed, 68 insertions(+), 47 deletions(-) diff --git a/crates/ide-assists/src/handlers/add_lifetime_to_type.rs b/crates/ide-assists/src/handlers/add_lifetime_to_type.rs index 27dbdcf2c4d5..265ee3d2d4e7 100644 --- a/crates/ide-assists/src/handlers/add_lifetime_to_type.rs +++ b/crates/ide-assists/src/handlers/add_lifetime_to_type.rs @@ -1,4 +1,7 @@ -use syntax::ast::{self, AstNode, HasGenericParams, HasName}; +use syntax::{ + SyntaxKind, SyntaxNode, SyntaxToken, + ast::{self, AstNode, HasGenericParams, HasName}, +}; use crate::{AssistContext, AssistId, Assists}; @@ -21,7 +24,7 @@ use crate::{AssistContext, AssistId, Assists}; // ``` pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { let ref_type_focused = ctx.find_node_at_offset::()?; - if ref_type_focused.lifetime().is_some() { + if ref_type_focused.lifetime().is_some_and(|lifetime| lifetime.text() != "'_") { return None; } @@ -34,10 +37,10 @@ pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_>) - return None; } - let ref_types = fetch_borrowed_types(&node)?; + let changes = fetch_borrowed_types(&node)?; let target = node.syntax().text_range(); - acc.add(AssistId::generate("add_lifetime_to_type"), "Add lifetime", target, |builder| { + acc.add(AssistId::quick_fix("add_lifetime_to_type"), "Add lifetime", target, |builder| { match node.generic_param_list() { Some(gen_param) => { if let Some(left_angle) = gen_param.l_angle_token() { @@ -51,16 +54,21 @@ pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_>) - } } - for ref_type in ref_types { - if let Some(amp_token) = ref_type.amp_token() { - builder.insert(amp_token.text_range().end(), "'a "); + for change in changes { + match change { + Change::Replace(it) => { + builder.replace(it.text_range(), "'a"); + } + Change::Insert(it) => { + builder.insert(it.text_range().end(), "'a "); + } } } }) } -fn fetch_borrowed_types(node: &ast::Adt) -> Option> { - let ref_types: Vec = match node { +fn fetch_borrowed_types(node: &ast::Adt) -> Option> { + let ref_types: Vec<_> = match node { ast::Adt::Enum(enum_) => { let variant_list = enum_.variant_list()?; variant_list @@ -79,55 +87,50 @@ fn fetch_borrowed_types(node: &ast::Adt) -> Option> { } ast::Adt::Union(un) => { let record_field_list = un.record_field_list()?; - record_field_list - .fields() - .filter_map(|r_field| { - if let ast::Type::RefType(ref_type) = r_field.ty()? - && ref_type.lifetime().is_none() - { - return Some(ref_type); - } - - None - }) - .collect() + find_ref_types_from_field_list(&record_field_list.into())? } }; if ref_types.is_empty() { None } else { Some(ref_types) } } -fn find_ref_types_from_field_list(field_list: &ast::FieldList) -> Option> { - let ref_types: Vec = match field_list { - ast::FieldList::RecordFieldList(record_list) => record_list - .fields() - .filter_map(|f| { - if let ast::Type::RefType(ref_type) = f.ty()? - && ref_type.lifetime().is_none() - { - return Some(ref_type); - } - - None - }) - .collect(), - ast::FieldList::TupleFieldList(tuple_field_list) => tuple_field_list - .fields() - .filter_map(|f| { - if let ast::Type::RefType(ref_type) = f.ty()? - && ref_type.lifetime().is_none() - { - return Some(ref_type); - } - - None - }) - .collect(), +fn find_ref_types_from_field_list(field_list: &ast::FieldList) -> Option> { + let ref_types: Vec<_> = match field_list { + ast::FieldList::RecordFieldList(record_list) => { + record_list.fields().flat_map(|f| infer_lifetimes(f.syntax())).collect() + } + ast::FieldList::TupleFieldList(tuple_field_list) => { + tuple_field_list.fields().flat_map(|f| infer_lifetimes(f.syntax())).collect() + } }; if ref_types.is_empty() { None } else { Some(ref_types) } } +enum Change { + Replace(SyntaxToken), + Insert(SyntaxToken), +} + +fn infer_lifetimes(node: &SyntaxNode) -> Vec { + node.children() + .filter(|it| !matches!(it.kind(), SyntaxKind::FN_PTR_TYPE | SyntaxKind::TYPE_BOUND_LIST)) + .flat_map(|it| { + infer_lifetimes(&it) + .into_iter() + .chain(ast::Lifetime::cast(it.clone()).and_then(|lt| { + lt.lifetime_ident_token().filter(|lt| lt.text() == "'_").map(Change::Replace) + })) + .chain( + ast::RefType::cast(it) + .filter(|ty| ty.lifetime().is_none()) + .and_then(|ty| ty.amp_token()) + .map(Change::Insert), + ) + }) + .collect() +} + #[cfg(test)] mod tests { use crate::tests::{check_assist, check_assist_not_applicable}; @@ -164,6 +167,24 @@ mod tests { check_assist_not_applicable(add_lifetime_to_type, r#"struct Foo { a: &'a$0 i32 }"#); } + #[test] + fn add_lifetime_to_nested_types() { + check_assist( + add_lifetime_to_type, + r#"struct Foo { a: &$0i32, b: &(&i32, fn(&str) -> &str) }"#, + r#"struct Foo<'a> { a: &'a i32, b: &'a (&'a i32, fn(&str) -> &str) }"#, + ); + } + + #[test] + fn add_lifetime_to_explicit_infer_lifetime() { + check_assist( + add_lifetime_to_type, + r#"struct Foo { a: &'_ $0i32, b: &'_ (&'_ i32, fn(&str) -> &str) }"#, + r#"struct Foo<'a> { a: &'a i32, b: &'a (&'a i32, fn(&str) -> &str) }"#, + ); + } + #[test] fn add_lifetime_to_enum() { check_assist(