Skip to content

Commit 9591eaf

Browse files
committed
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>), } ```
1 parent 1f4e5e8 commit 9591eaf

File tree

1 file changed

+68
-47
lines changed

1 file changed

+68
-47
lines changed

crates/ide-assists/src/handlers/add_lifetime_to_type.rs

Lines changed: 68 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
1-
use syntax::ast::{self, AstNode, HasGenericParams, HasName};
1+
use syntax::{
2+
SyntaxKind, SyntaxNode, SyntaxToken,
3+
ast::{self, AstNode, HasGenericParams, HasName},
4+
};
25

36
use crate::{AssistContext, AssistId, Assists};
47

@@ -21,7 +24,7 @@ use crate::{AssistContext, AssistId, Assists};
2124
// ```
2225
pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
2326
let ref_type_focused = ctx.find_node_at_offset::<ast::RefType>()?;
24-
if ref_type_focused.lifetime().is_some() {
27+
if ref_type_focused.lifetime().is_some_and(|lifetime| lifetime.text() != "'_") {
2528
return None;
2629
}
2730

@@ -34,10 +37,10 @@ pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_>) -
3437
return None;
3538
}
3639

37-
let ref_types = fetch_borrowed_types(&node)?;
40+
let changes = fetch_borrowed_types(&node)?;
3841
let target = node.syntax().text_range();
3942

40-
acc.add(AssistId::generate("add_lifetime_to_type"), "Add lifetime", target, |builder| {
43+
acc.add(AssistId::quick_fix("add_lifetime_to_type"), "Add lifetime", target, |builder| {
4144
match node.generic_param_list() {
4245
Some(gen_param) => {
4346
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<'_>) -
5154
}
5255
}
5356

54-
for ref_type in ref_types {
55-
if let Some(amp_token) = ref_type.amp_token() {
56-
builder.insert(amp_token.text_range().end(), "'a ");
57+
for change in changes {
58+
match change {
59+
Change::Replace(it) => {
60+
builder.replace(it.text_range(), "'a");
61+
}
62+
Change::Insert(it) => {
63+
builder.insert(it.text_range().end(), "'a ");
64+
}
5765
}
5866
}
5967
})
6068
}
6169

62-
fn fetch_borrowed_types(node: &ast::Adt) -> Option<Vec<ast::RefType>> {
63-
let ref_types: Vec<ast::RefType> = match node {
70+
fn fetch_borrowed_types(node: &ast::Adt) -> Option<Vec<Change>> {
71+
let ref_types: Vec<_> = match node {
6472
ast::Adt::Enum(enum_) => {
6573
let variant_list = enum_.variant_list()?;
6674
variant_list
@@ -79,55 +87,50 @@ fn fetch_borrowed_types(node: &ast::Adt) -> Option<Vec<ast::RefType>> {
7987
}
8088
ast::Adt::Union(un) => {
8189
let record_field_list = un.record_field_list()?;
82-
record_field_list
83-
.fields()
84-
.filter_map(|r_field| {
85-
if let ast::Type::RefType(ref_type) = r_field.ty()?
86-
&& ref_type.lifetime().is_none()
87-
{
88-
return Some(ref_type);
89-
}
90-
91-
None
92-
})
93-
.collect()
90+
find_ref_types_from_field_list(&record_field_list.into())?
9491
}
9592
};
9693

9794
if ref_types.is_empty() { None } else { Some(ref_types) }
9895
}
9996

100-
fn find_ref_types_from_field_list(field_list: &ast::FieldList) -> Option<Vec<ast::RefType>> {
101-
let ref_types: Vec<ast::RefType> = match field_list {
102-
ast::FieldList::RecordFieldList(record_list) => record_list
103-
.fields()
104-
.filter_map(|f| {
105-
if let ast::Type::RefType(ref_type) = f.ty()?
106-
&& ref_type.lifetime().is_none()
107-
{
108-
return Some(ref_type);
109-
}
110-
111-
None
112-
})
113-
.collect(),
114-
ast::FieldList::TupleFieldList(tuple_field_list) => tuple_field_list
115-
.fields()
116-
.filter_map(|f| {
117-
if let ast::Type::RefType(ref_type) = f.ty()?
118-
&& ref_type.lifetime().is_none()
119-
{
120-
return Some(ref_type);
121-
}
122-
123-
None
124-
})
125-
.collect(),
97+
fn find_ref_types_from_field_list(field_list: &ast::FieldList) -> Option<Vec<Change>> {
98+
let ref_types: Vec<_> = match field_list {
99+
ast::FieldList::RecordFieldList(record_list) => {
100+
record_list.fields().flat_map(|f| infer_lifetimes(f.syntax())).collect()
101+
}
102+
ast::FieldList::TupleFieldList(tuple_field_list) => {
103+
tuple_field_list.fields().flat_map(|f| infer_lifetimes(f.syntax())).collect()
104+
}
126105
};
127106

128107
if ref_types.is_empty() { None } else { Some(ref_types) }
129108
}
130109

110+
enum Change {
111+
Replace(SyntaxToken),
112+
Insert(SyntaxToken),
113+
}
114+
115+
fn infer_lifetimes(node: &SyntaxNode) -> Vec<Change> {
116+
node.children()
117+
.filter(|it| !matches!(it.kind(), SyntaxKind::FN_PTR_TYPE | SyntaxKind::TYPE_BOUND_LIST))
118+
.flat_map(|it| {
119+
infer_lifetimes(&it)
120+
.into_iter()
121+
.chain(ast::Lifetime::cast(it.clone()).and_then(|lt| {
122+
lt.lifetime_ident_token().filter(|lt| lt.text() == "'_").map(Change::Replace)
123+
}))
124+
.chain(
125+
ast::RefType::cast(it)
126+
.filter(|ty| ty.lifetime().is_none())
127+
.and_then(|ty| ty.amp_token())
128+
.map(Change::Insert),
129+
)
130+
})
131+
.collect()
132+
}
133+
131134
#[cfg(test)]
132135
mod tests {
133136
use crate::tests::{check_assist, check_assist_not_applicable};
@@ -164,6 +167,24 @@ mod tests {
164167
check_assist_not_applicable(add_lifetime_to_type, r#"struct Foo { a: &'a$0 i32 }"#);
165168
}
166169

170+
#[test]
171+
fn add_lifetime_to_nested_types() {
172+
check_assist(
173+
add_lifetime_to_type,
174+
r#"struct Foo { a: &$0i32, b: &(&i32, fn(&str) -> &str) }"#,
175+
r#"struct Foo<'a> { a: &'a i32, b: &'a (&'a i32, fn(&str) -> &str) }"#,
176+
);
177+
}
178+
179+
#[test]
180+
fn add_lifetime_to_explicit_infer_lifetime() {
181+
check_assist(
182+
add_lifetime_to_type,
183+
r#"struct Foo { a: &'_ $0i32, b: &'_ (&'_ i32, fn(&str) -> &str) }"#,
184+
r#"struct Foo<'a> { a: &'a i32, b: &'a (&'a i32, fn(&str) -> &str) }"#,
185+
);
186+
}
187+
167188
#[test]
168189
fn add_lifetime_to_enum() {
169190
check_assist(

0 commit comments

Comments
 (0)