Skip to content

Commit 9331224

Browse files
A4-TacksA4-Tacks
committed
Trigger on non generic args path type
1 parent edf5375 commit 9331224

1 file changed

Lines changed: 70 additions & 45 deletions

File tree

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

Lines changed: 70 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
1+
use either::Either;
12
use syntax::{
2-
SyntaxKind, SyntaxNode, SyntaxToken,
3+
SyntaxKind, SyntaxNode,
34
ast::{self, AstNode, HasGenericParams, HasName},
5+
match_ast,
46
};
57

68
use crate::{AssistContext, AssistId, Assists};
@@ -23,10 +25,12 @@ use crate::{AssistContext, AssistId, Assists};
2325
// }
2426
// ```
2527
pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_, '_>) -> Option<()> {
26-
if !trigger_assist(ctx) {
27-
return None;
28-
}
2928
let node = ctx.find_node_at_offset::<ast::Adt>()?;
29+
// XXX: Maybe delete this and allow it to be triggered conveniently on ADT
30+
let _trigger = syntax::algo::ancestors_at_offset(node.syntax(), ctx.offset())
31+
.take_while(|it| node.syntax() != it)
32+
.filter_map(Either::<ast::Lifetime, Either<ast::RefType, ast::PathType>>::cast)
33+
.find_map(|it| Missing::from_node(dbg!(it.syntax().clone()), ctx))?;
3034
let has_lifetime = node
3135
.generic_param_list()
3236
.is_some_and(|gen_list| gen_list.lifetime_params().next().is_some());
@@ -35,7 +39,7 @@ pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_, '_
3539
return None;
3640
}
3741

38-
let changes = fetch_borrowed_types(&node)?;
42+
let changes = fetch_borrowed_types(&node, ctx)?;
3943
let target = node.syntax().text_range();
4044

4145
acc.add(AssistId::quick_fix("add_lifetime_to_type"), "Add lifetime", target, |builder| {
@@ -54,28 +58,23 @@ pub(crate) fn add_lifetime_to_type(acc: &mut Assists, ctx: &AssistContext<'_, '_
5458

5559
for change in changes {
5660
match change {
57-
Change::Replace(it) => {
58-
builder.replace(it.text_range(), "'a");
61+
Missing::Lifetime(lt) => {
62+
builder.replace(lt.syntax().text_range(), "'a");
5963
}
60-
Change::Insert(it) => {
61-
builder.insert(it.text_range().end(), "'a ");
64+
Missing::RefType(ref_type) => {
65+
if let Some(amp_token) = ref_type.amp_token() {
66+
builder.insert(amp_token.text_range().end(), "'a ");
67+
}
68+
}
69+
Missing::PathType(path_type) => {
70+
builder.insert(path_type.syntax().text_range().end(), "<'a>");
6271
}
6372
}
6473
}
6574
})
6675
}
6776

68-
fn trigger_assist(ctx: &AssistContext<'_, '_>) -> bool {
69-
ctx.find_node_at_offset::<ast::RefType>()
70-
.is_some_and(|it| it.lifetime().is_none_or(|it| it.text() == "'_"))
71-
|| ctx
72-
.find_node_at_offset::<ast::PathType>()
73-
.map(ast::Type::from)
74-
.and_then(|it| it.generic_arg_list()?.lifetime_args().next()?.lifetime())
75-
.is_some_and(|it| it.text() == "'_")
76-
}
77-
78-
fn fetch_borrowed_types(node: &ast::Adt) -> Option<Vec<Change>> {
77+
fn fetch_borrowed_types(node: &ast::Adt, ctx: &AssistContext<'_, '_>) -> Option<Vec<Missing>> {
7978
let ref_types: Vec<_> = match node {
8079
ast::Adt::Enum(enum_) => {
8180
let variant_list = enum_.variant_list()?;
@@ -84,58 +83,71 @@ fn fetch_borrowed_types(node: &ast::Adt) -> Option<Vec<Change>> {
8483
.filter_map(|variant| {
8584
let field_list = variant.field_list()?;
8685

87-
find_ref_types_from_field_list(&field_list)
86+
find_ref_types_from_field_list(&field_list, ctx)
8887
})
8988
.flatten()
9089
.collect()
9190
}
9291
ast::Adt::Struct(strukt) => {
9392
let field_list = strukt.field_list()?;
94-
find_ref_types_from_field_list(&field_list)?
93+
find_ref_types_from_field_list(&field_list, ctx)?
9594
}
9695
ast::Adt::Union(un) => {
9796
let record_field_list = un.record_field_list()?;
98-
find_ref_types_from_field_list(&record_field_list.into())?
97+
find_ref_types_from_field_list(&record_field_list.into(), ctx)?
9998
}
10099
};
101100

102101
if ref_types.is_empty() { None } else { Some(ref_types) }
103102
}
104103

105-
fn find_ref_types_from_field_list(field_list: &ast::FieldList) -> Option<Vec<Change>> {
104+
fn find_ref_types_from_field_list(
105+
field_list: &ast::FieldList,
106+
ctx: &AssistContext<'_, '_>,
107+
) -> Option<Vec<Missing>> {
106108
let ref_types: Vec<_> = match field_list {
107109
ast::FieldList::RecordFieldList(record_list) => {
108-
record_list.fields().flat_map(|f| infer_lifetimes(f.syntax())).collect()
110+
record_list.fields().flat_map(|f| infer_lifetimes(f.syntax(), ctx)).collect()
109111
}
110112
ast::FieldList::TupleFieldList(tuple_field_list) => {
111-
tuple_field_list.fields().flat_map(|f| infer_lifetimes(f.syntax())).collect()
113+
tuple_field_list.fields().flat_map(|f| infer_lifetimes(f.syntax(), ctx)).collect()
112114
}
113115
};
114116

115117
if ref_types.is_empty() { None } else { Some(ref_types) }
116118
}
117119

118-
enum Change {
119-
Replace(SyntaxToken),
120-
Insert(SyntaxToken),
120+
enum Missing {
121+
RefType(ast::RefType),
122+
PathType(ast::PathType),
123+
Lifetime(ast::Lifetime),
124+
}
125+
126+
impl Missing {
127+
fn from_node(node: SyntaxNode, ctx: &AssistContext<'_, '_>) -> Option<Self> {
128+
match_ast! {
129+
match node {
130+
ast::Lifetime(it) => (it.syntax().text() == "'_").then_some(Missing::Lifetime(it)),
131+
ast::RefType(it) => (it.lifetime().is_none() && it.amp_token().is_some()).then_some(Missing::RefType(it)),
132+
ast::PathType(it) => {
133+
let needs_lifetime = match ctx.sema.resolve_path(&it.path()?)? {
134+
hir::PathResolution::Def(hir::ModuleDef::Adt(adt)) => adt.lifetime(ctx.db()).is_some(),
135+
// FIXME: check TypeAlias and Trait lifetime params
136+
_ => false,
137+
};
138+
(needs_lifetime && ast::Type::from(it.clone()).generic_arg_list().is_none())
139+
.then_some(Missing::PathType(it))
140+
},
141+
_ => None,
142+
}
143+
}
144+
}
121145
}
122146

123-
fn infer_lifetimes(node: &SyntaxNode) -> Vec<Change> {
147+
fn infer_lifetimes(node: &SyntaxNode, ctx: &AssistContext<'_, '_>) -> Vec<Missing> {
124148
node.children()
125149
.filter(|it| !matches!(it.kind(), SyntaxKind::FN_PTR_TYPE | SyntaxKind::TYPE_BOUND_LIST))
126-
.flat_map(|it| {
127-
infer_lifetimes(&it)
128-
.into_iter()
129-
.chain(ast::Lifetime::cast(it.clone()).and_then(|lt| {
130-
lt.lifetime_ident_token().filter(|lt| lt.text() == "'_").map(Change::Replace)
131-
}))
132-
.chain(
133-
ast::RefType::cast(it)
134-
.filter(|ty| ty.lifetime().is_none())
135-
.and_then(|ty| ty.amp_token())
136-
.map(Change::Insert),
137-
)
138-
})
150+
.flat_map(|it| infer_lifetimes(&it, ctx).into_iter().chain(Missing::from_node(it, ctx)))
139151
.collect()
140152
}
141153

@@ -188,13 +200,13 @@ mod tests {
188200
fn add_lifetime_to_explicit_infer_lifetime() {
189201
check_assist(
190202
add_lifetime_to_type,
191-
r#"struct Foo { a: &'_ $0i32, b: &'_ (&'_ i32, fn(&str) -> &str) }"#,
203+
r#"struct Foo { a: &'_$0 i32, b: &'_ (&'_ i32, fn(&str) -> &str) }"#,
192204
r#"struct Foo<'a> { a: &'a i32, b: &'a (&'a i32, fn(&str) -> &str) }"#,
193205
);
194206

195207
check_assist(
196208
add_lifetime_to_type,
197-
r#"struct Foo { a: &'_ $0i32, b: Foo<'_> }"#,
209+
r#"struct Foo { a: &'_$0 i32, b: Foo<'_> }"#,
198210
r#"struct Foo<'a> { a: &'a i32, b: Foo<'a> }"#,
199211
);
200212

@@ -205,6 +217,19 @@ mod tests {
205217
);
206218
}
207219

220+
#[test]
221+
fn add_lifetime_to_implicit_infer_lifetime() {
222+
check_assist(
223+
add_lifetime_to_type,
224+
r#"
225+
struct Ref<'a>(&'a ());
226+
struct Foo { a: &'_ i32, b: Ref$0 }"#,
227+
r#"
228+
struct Ref<'a>(&'a ());
229+
struct Foo<'a> { a: &'a i32, b: Ref<'a> }"#,
230+
);
231+
}
232+
208233
#[test]
209234
fn add_lifetime_to_enum() {
210235
check_assist(

0 commit comments

Comments
 (0)